Zhiwei/sglang model adapter

atom/plugin/prepare.py

@@ -74,15 +74,11 @@ def prepare_model(config: Any, engine: str):
     model_cls = _ATOM_SUPPORTED_MODELS[model_arch]
     logger.info(f"ATOM model class for {model_arch} is {model_cls}")
 
-    if model_arch in {
-        "Qwen3_5ForConditionalGeneration",
-        "Qwen3_5MoeForConditionalGeneration",
-    }:
-        from atom.plugin.sglang.models.qwen3_5 import (
-            apply_prepare_model_adaptations,
-        )
+    from atom.plugin.sglang.runtime import get_model_arch_spec
 
-        apply_prepare_model_adaptations(atom_config, model_arch)
+    model_adapter = get_model_arch_spec(model_arch)
+    if model_adapter.prepare_config is not None:
+        model_adapter.prepare_config(atom_config, model_arch)
 
     register_ops_to_sglang(atom_config=atom_config)
     set_attn_cls()

atom/plugin/sglang/models/base_model_wrapper.py

@@ -92,15 +92,10 @@ def __init__(
             config, skip_all_gather=plugin_skip_all_gather
         )
 
-        # Apply ds model-specific sglang patches (attn dispatch, weight hooks, etc.)
-        # TODO: will remove this after sglang supports atom attention backend
-        if self.model_arch_spec.apply_deepseek_patch:
-            from atom.plugin.sglang.models.deepseek_mla import (
-                setup_deepseek_for_sglang,
-            )
-
+        # Apply model-specific install-time adapters (attn dispatch, weight hooks, etc.).
+        if self.model_arch_spec.install_adapters is not None:
             with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config):
-                setup_deepseek_for_sglang(self.model)
+                self.model_arch_spec.install_adapters(self.model)
 
     def get_embed_and_head(self):
         if hasattr(self.model, "get_embed_and_head"):
@@ -180,7 +175,7 @@ def forward(
                     inputs_embeds=runtime.input_embeds,
                 )
                 uses_context_only_forward = (
-                    self.model_arch_spec.apply_deepseek_patch
+                    self.model_arch_spec.install_adapters is not None
                     or self.model_arch_spec.wrapper_binds_gdn_context
                 )
                 with SGLangForwardBatchMetadata.bind(metadata):

atom/plugin/sglang/runtime/__init__.py

@@ -7,11 +7,18 @@
     plugin_runtime_scope,
 )
 from atom.plugin.sglang.runtime.forward_context import SGLangPluginRuntime
-from atom.plugin.sglang.runtime.model_arch import MODEL_ARCH_SPECS, get_model_arch_spec
+from atom.plugin.sglang.runtime.model_arch import (
+    MODEL_ADAPTER_SPECS,
+    MODEL_ARCH_SPECS,
+    SGLangModelAdapterSpec,
+    get_model_arch_spec,
+)
 
 __all__ = [
+    "MODEL_ADAPTER_SPECS",
     "MODEL_ARCH_SPECS",
     "SGLangForwardBatchMetadata",
+    "SGLangModelAdapterSpec",
     "SGLangPluginRuntime",
     "bind_current_forward_batch",
     "get_current_forward_batch",

atom/plugin/sglang/runtime/model_arch.py

@@ -1,22 +1,65 @@
-"""Runtime behavior flags for SGLang plugin model wrappers."""
+"""SGLang plugin model adapter registry."""
 
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import Any, Callable, Optional
 
 
 @dataclass(frozen=True)
-class ModelArchSpec:
+class SGLangModelAdapterSpec:
+    """Adapter hooks for one SGLang plugin model architecture.
+
+    The first version keeps the existing runtime flags while adding function
+    hooks for config preparation and install-time model adaptation. This avoids
+    growing a long list of booleans in the generic wrapper as new models arrive.
+    """
+
     wrapper_binds_gdn_context: bool = False
-    apply_deepseek_patch: bool = False
+    prepare_config: Optional[Callable[[Any, str], None]] = None
+    install_adapters: Optional[Callable[[Any], None]] = None
+
+
+def _prepare_qwen35_config(atom_config: Any, model_arch: str) -> None:
+    from atom.plugin.sglang.models.qwen3_5 import apply_prepare_model_adaptations
+
+    apply_prepare_model_adaptations(atom_config, model_arch)
 
 
+def _install_deepseek_mla_adapters(model: Any) -> None:
+    from atom.plugin.sglang.models.deepseek_mla import setup_deepseek_for_sglang
+
+    setup_deepseek_for_sglang(model)
+
+
+MODEL_ADAPTER_SPECS = {
+    "DeepseekV3ForCausalLM": SGLangModelAdapterSpec(
+        install_adapters=_install_deepseek_mla_adapters,
+    ),
+    "Qwen3MoeForCausalLM": SGLangModelAdapterSpec(),
+    "Qwen3NextForCausalLM": SGLangModelAdapterSpec(
+        wrapper_binds_gdn_context=True,
+    ),
+    "Qwen3_5ForConditionalGeneration": SGLangModelAdapterSpec(
+        prepare_config=_prepare_qwen35_config,
+    ),
+    "Qwen3_5MoeForConditionalGeneration": SGLangModelAdapterSpec(
+        prepare_config=_prepare_qwen35_config,
+    ),
+}
+
+# Architectures whose SGLang EntryClass is generated by base_model_wrapper.
+# Custom outer-wrapper modules, such as Qwen3.5 multimodal wrappers, keep their
+# own EntryClass and should not appear here or SGLang will see duplicate classes.
 MODEL_ARCH_SPECS = {
-    "DeepseekV3ForCausalLM": ModelArchSpec(apply_deepseek_patch=True),
-    "Qwen3MoeForCausalLM": ModelArchSpec(),
-    "Qwen3NextForCausalLM": ModelArchSpec(wrapper_binds_gdn_context=True),
+    key: MODEL_ADAPTER_SPECS[key]
+    for key in (
+        "DeepseekV3ForCausalLM",
+        "Qwen3MoeForCausalLM",
+        "Qwen3NextForCausalLM",
+    )
 }
 
 
-def get_model_arch_spec(model_arch: str) -> ModelArchSpec:
-    return MODEL_ARCH_SPECS.get(model_arch, ModelArchSpec())
+def get_model_arch_spec(model_arch: str) -> SGLangModelAdapterSpec:
+    return MODEL_ADAPTER_SPECS.get(model_arch, SGLangModelAdapterSpec())