[SGL ATOM] enable qwen3.8-flash mtp in sgl atom

atom/plugin/config.py

@@ -268,6 +268,48 @@ def _build_atom_speculative_config_from_vllm(vllm_spec_config: Any):
     )
 
 
+def _build_atom_speculative_config_from_sglang(server_args: Any, hf_config: Any):
+    """SGLang EAGLE/NEXTN → ATOM ``SpeculativeConfig``.
+
+    Native Qwen4Exp returns the unmixed HC bundle only when
+    ``atom_config.speculative_config is not None``. SGLang owns scheduling;
+    this object is that ABI flag plus the Native draft-token horizon
+    (``speculative_num_steps``, not SGLang's ``steps+1`` draft-token count).
+    """
+    algo = getattr(server_args, "speculative_algorithm", None)
+    name = str(getattr(algo, "name", algo)).strip().upper()
+    if name in {"", "NONE", "NULL", "FALSE"}:
+        return None
+
+    text = getattr(hf_config, "text_config", None) or hf_config
+    if getattr(text, "model_type", None) not in {"qwen4_exp", "qwen4_exp_text"}:
+        return None
+    if name not in {"EAGLE", "NEXTN"}:
+        raise ValueError("Flash MTP requires speculative_algorithm=EAGLE (or NEXTN)")
+    if getattr(server_args, "speculative_eagle_topk", 1) != 1:
+        raise ValueError("Flash MTP requires speculative_eagle_topk=1")
+    draft_path = getattr(server_args, "speculative_draft_model_path", None)
+    if draft_path and draft_path != server_args.model_path:
+        raise ValueError(
+            "Flash MTP requires the target and draft to share a checkpoint"
+        )
+
+    from atom.config import SpeculativeConfig
+
+    steps = getattr(server_args, "speculative_num_steps", None)
+    num_spec = int(steps) if steps is not None else None
+    logger.info(
+        "SGLang+ATOM speculative_config method=mtp tokens=%s model=%s",
+        num_spec,
+        server_args.model_path,
+    )
+    return SpeculativeConfig(
+        method="mtp",
+        model=server_args.model_path,
+        num_speculative_tokens=num_spec,
+    )
+
+
 def _generate_atom_config_from_vllm_config(config: Any) -> PluginConfig:
     from atom.config import CompilationConfig, Config
 
@@ -586,6 +628,10 @@ def _generate_atom_config_from_sglang_config(config: Any):
     if str(atom_kv_cache_dtype).startswith("fp8"):
         atom_kv_cache_dtype = "fp8"
 
+    atom_speculative_config = _build_atom_speculative_config_from_sglang(
+        server_args, sgl_model_config.hf_config
+    )
+
     return Config(
         model=server_args.model_path,
         trust_remote_code=server_args.trust_remote_code,
@@ -617,6 +663,7 @@ def _generate_atom_config_from_sglang_config(config: Any):
         plugin_config=plugin_config,
         online_quant_config=online_quant_config,
         hf_overrides=hf_overrides,
+        speculative_config=atom_speculative_config,
     )
 
 

atom/plugin/register.py

@@ -269,6 +269,13 @@ def _is_glm52_runner(runner) -> bool:
         except Exception:
             return False
 
+    def _is_qwen4_exp_nextn_runner(runner) -> bool:
+        model_config = getattr(runner, "model_config", None)
+        hf_config = getattr(model_config, "hf_config", None)
+        return "Qwen4ExpForCausalLMNextN" in (
+            getattr(hf_config, "architectures", None) or []
+        )
+
     def _is_dsv4_or_glm52_nextn_runner(runner) -> bool:
         return _is_dsv4_nextn_runner(runner) or _is_glm52_nextn_runner(runner)
 
@@ -788,15 +795,22 @@ def init_cuda_graphs(self):
                         backend.get_cuda_graph_seq_len_fill_value = (
                             lambda value=GLM52_GRAPH_SEQ_LEN_CAPACITY: value
                         )
+            # Flash draft stays eager: mRoPE positions are [3, N], same as
+            # Native EagleProposer skipping mid-step graphs.
             skip_all_draft_graphs = (
                 _env_flag("ATOM_SGLANG_V4_DISABLE_DRAFT_CG")
                 and not _draft_extend_graph_enabled(draft_runner)
                 and _is_dsv4_or_glm52_nextn_runner(draft_runner)
-            )
+            ) or _is_qwen4_exp_nextn_runner(draft_runner)
             original_capture_cuda_graphs = None
             if skip_all_draft_graphs:
                 original_capture_cuda_graphs = self._capture_cuda_graphs
-                self._capture_cuda_graphs = lambda: None
+
+                def _skip_draft_cuda_graphs(_self=self):
+                    _self.cuda_graph_runner = None
+                    _self.cuda_graph_runner_for_draft_extend = None
+
+                self._capture_cuda_graphs = _skip_draft_cuda_graphs
             original_draft_extend_backend = None
             hide_draft_extend_backend = (
                 not _draft_extend_graph_enabled(draft_runner)