From b9eeaa577af1d7a88a953d073a82538320afc536 Mon Sep 17 00:00:00 2001 From: ZhiweiYan-96 Date: Wed, 20 May 2026 08:35:46 +0000 Subject: [PATCH 1/2] [ATOM-SGL][Runtime] Introduce model adapter specs Co-authored-by: Cursor --- atom/plugin/prepare.py | 12 ++--- .../sglang/models/base_model_wrapper.py | 13 ++--- atom/plugin/sglang/runtime/__init__.py | 9 +++- atom/plugin/sglang/runtime/model_arch.py | 52 +++++++++++++++---- 4 files changed, 59 insertions(+), 27 deletions(-) diff --git a/atom/plugin/prepare.py b/atom/plugin/prepare.py index 3c7b722d29..bbcb369bcd 100644 --- a/atom/plugin/prepare.py +++ b/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() diff --git a/atom/plugin/sglang/models/base_model_wrapper.py b/atom/plugin/sglang/models/base_model_wrapper.py index 85bc6608e1..02ac6edc77 100644 --- a/atom/plugin/sglang/models/base_model_wrapper.py +++ b/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): diff --git a/atom/plugin/sglang/runtime/__init__.py b/atom/plugin/sglang/runtime/__init__.py index f63890418f..fa8468c177 100644 --- a/atom/plugin/sglang/runtime/__init__.py +++ b/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", diff --git a/atom/plugin/sglang/runtime/model_arch.py b/atom/plugin/sglang/runtime/model_arch.py index 1ddd3772c9..4765baa483 100644 --- a/atom/plugin/sglang/runtime/model_arch.py +++ b/atom/plugin/sglang/runtime/model_arch.py @@ -1,22 +1,56 @@ -"""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) -MODEL_ARCH_SPECS = { - "DeepseekV3ForCausalLM": ModelArchSpec(apply_deepseek_patch=True), - "Qwen3MoeForCausalLM": ModelArchSpec(), - "Qwen3NextForCausalLM": ModelArchSpec(wrapper_binds_gdn_context=True), +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, + ), } +# Backwards-compatible alias for callers that only need generated EntryClass names. +MODEL_ARCH_SPECS = MODEL_ADAPTER_SPECS + -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()) From cfda3b018b5348fd5d04c52594bb1f5b723554b0 Mon Sep 17 00:00:00 2001 From: ZhiweiYan-96 Date: Thu, 21 May 2026 02:32:48 +0000 Subject: [PATCH 2/2] [ATOM-SGL][Runtime] Keep custom wrappers out of generated entries Co-authored-by: Cursor --- atom/plugin/sglang/runtime/model_arch.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/atom/plugin/sglang/runtime/model_arch.py b/atom/plugin/sglang/runtime/model_arch.py index 4765baa483..f103cb9427 100644 --- a/atom/plugin/sglang/runtime/model_arch.py +++ b/atom/plugin/sglang/runtime/model_arch.py @@ -48,8 +48,17 @@ def _install_deepseek_mla_adapters(model: Any) -> None: ), } -# Backwards-compatible alias for callers that only need generated EntryClass names. -MODEL_ARCH_SPECS = MODEL_ADAPTER_SPECS +# 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 = { + key: MODEL_ADAPTER_SPECS[key] + for key in ( + "DeepseekV3ForCausalLM", + "Qwen3MoeForCausalLM", + "Qwen3NextForCausalLM", + ) +} def get_model_arch_spec(model_arch: str) -> SGLangModelAdapterSpec: