perf(sglang-atom): CUDA graphs for the Qwen3.8-Flash-Next MTP draft (decode + extend)

atom/plugin/register.py

@@ -6,6 +6,7 @@
 
 from atom.config import Config
 from atom.plugin.prepare import is_rtpllm, is_sglang, is_vllm
+from atom.utils import envs
 
 logger = logging.getLogger("atom")
 
@@ -410,8 +411,18 @@ def _is_dsv4_pro_runner(runner) -> bool:
     def _draft_extend_graph_enabled(runner) -> bool:
         if _env_flag("ATOM_SGLANG_V4_DISABLE_DRAFT_EXTEND_CG"):
             return False
-        return _env_flag("ATOM_SGLANG_V4_ENABLE_DRAFT_EXTEND_CG") or (
-            _is_dsv4_nextn_runner(runner) and _is_dsv4_flash_runner(runner)
+        return (
+            _env_flag("ATOM_SGLANG_V4_ENABLE_DRAFT_EXTEND_CG")
+            or (_is_dsv4_nextn_runner(runner) and _is_dsv4_flash_runner(runner))
+            or _qwen4_draft_extend_graph_enabled(runner)
+        )
+
+    def _qwen4_draft_extend_graph_enabled(runner) -> bool:
+        """PROTOTYPE: Flash draft-extend graph, on top of the draft decode graph."""
+        return (
+            _is_qwen4_exp_nextn_runner(runner)
+            and envs.ATOM_SGLANG_QWEN4_DRAFT_CG
+            and envs.ATOM_SGLANG_QWEN4_DRAFT_EXTEND_CG
         )
 
     def _target_verify_graph_enabled() -> bool:
@@ -523,6 +534,13 @@ def _replay(self, forward_batch):
             setattr(EAGLEDraftCudaGraphRunner, draft_internal_method, _replay)
 
         def replay(self, forward_batch):
+            if _is_qwen4_exp_nextn_runner(getattr(self, "model_runner", None)):
+                from atom.plugin.sglang.qwen4_exp_bridge import (
+                    stage_draft_graph_mrope,
+                )
+
+                stage_draft_graph_mrope(self.buffers.mrope_positions, forward_batch)
+                return original_draft_replay(self, forward_batch)
             if not _is_dsv4_or_glm52_nextn_runner(getattr(self, "model_runner", None)):
                 return original_draft_replay(self, forward_batch)
             if _env_flag("ATOM_SGLANG_V4_DISABLE_DRAFT_CG"):
@@ -606,16 +624,25 @@ def can_run(self, forward_batch):
                 return False
             return _dsv4_draft_extend_graph_layout_ok(self, forward_batch)
 
-        def replay(self, forward_batch):
+        # SGLang 0.5.20's ``execute(forward_batch, select_index)`` passes an
+        # extra argument; forward whatever the runner API passes.
+        def replay(self, forward_batch, *args, **kwargs):
             model_runner = getattr(self, "model_runner", None)
             if _is_glm52_nextn_runner(model_runner):
                 if not _draft_extend_graph_enabled(model_runner):
                     raise RuntimeError(
                         "GLM-5.2 draft-extend cuda graph replay was disabled"
                     )
-                return original_extend_replay(self, forward_batch)
+                return original_extend_replay(self, forward_batch, *args, **kwargs)
+            if _is_qwen4_exp_nextn_runner(model_runner):
+                from atom.plugin.sglang.qwen4_exp_bridge import (
+                    stage_draft_graph_mrope,
+                )
+
+                stage_draft_graph_mrope(self.buffers.mrope_positions, forward_batch)
+                return original_extend_replay(self, forward_batch, *args, **kwargs)
             if not _is_dsv4_nextn_runner(model_runner):
-                return original_extend_replay(self, forward_batch)
+                return original_extend_replay(self, forward_batch, *args, **kwargs)
             if not _draft_extend_graph_enabled(getattr(self, "model_runner", None)):
                 raise RuntimeError(
                     "DSV4 draft-extend cuda graph replay was disabled after capture; "
@@ -651,7 +678,7 @@ def replay(self, forward_batch):
                                 :num_tokens
                             ],
                         )
-                out = original_extend_replay(self, forward_batch)
+                out = original_extend_replay(self, forward_batch, *args, **kwargs)
                 # EAGLE V2 consumes draft-extend logits with a fixed
                 # `seq * speculative_num_draft_tokens + offset` layout.
                 # SGLang's runner trims to the actual compact token count,
@@ -695,6 +722,46 @@ def replay(self, forward_batch):
         setattr(EAGLEDraftExtendCudaGraphRunner, extend_replay_method, replay)
         EAGLEDraftExtendCudaGraphRunner._atom_dsv4_replay_patched = True
 
+    if not getattr(EagleDraftWorker, "_atom_qwen4_draft_mrope_patched", False):
+        original_draft_forward = EagleDraftWorker.draft_forward
+
+        def draft_forward(self, forward_batch):
+            # EAGLE advances only the 1-D positions after each draft decode
+            # forward, and the eager / graph runners hand the model a copy of
+            # the batch. Advance the caller's mRoPE rows (the Flash drafter's
+            # RoPE input) alongside; inside a draft graph capture the add is
+            # recorded on the graph's mRoPE buffer.
+            runner = getattr(self, "draft_runner", None)
+            if (
+                not _is_qwen4_exp_nextn_runner(runner)
+                or forward_batch.forward_mode.is_idle()
+                or getattr(forward_batch, "mrope_positions", None) is None
+            ):
+                return original_draft_forward(self, forward_batch)
+            from atom.plugin.sglang.qwen4_exp_bridge import (
+                advance_draft_decode_mrope,
+            )
+
+            shadowed = vars(runner).get("forward")
+            runner_forward = runner.forward
+
+            def forward(step_batch, *args, **kwargs):
+                out = runner_forward(step_batch, *args, **kwargs)
+                advance_draft_decode_mrope(forward_batch)
+                return out
+
+            runner.forward = forward
+            try:
+                return original_draft_forward(self, forward_batch)
+            finally:
+                if shadowed is None:
+                    del runner.forward
+                else:
+                    runner.forward = shadowed
+
+        EagleDraftWorker.draft_forward = draft_forward
+        EagleDraftWorker._atom_qwen4_draft_mrope_patched = True
+
     if not getattr(EagleDraftWorker, "_atom_dsv4_draft_extend_accept_patched", False):
         original_draft_extend_for_decode = EagleDraftWorker._draft_extend_for_decode
 
@@ -860,13 +927,18 @@ 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.
+            # Flash draft graphs are opt-in: ATOM_SGLANG_QWEN4_DRAFT_CG captures
+            # the draft decode graph (ATOM_SGLANG_QWEN4_DRAFT_EXTEND_CG adds the
+            # draft-extend graph below). Their [3, N] mRoPE input is staged on
+            # replay by stage_draft_graph_mrope.
             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)
+            ) or (
+                _is_qwen4_exp_nextn_runner(draft_runner)
+                and not envs.ATOM_SGLANG_QWEN4_DRAFT_CG
+            )
             original_capture_cuda_graphs = None
             if skip_all_draft_graphs:
                 original_capture_cuda_graphs = self._capture_cuda_graphs
@@ -898,8 +970,13 @@ def _skip_draft_cuda_graphs(_self=self):
                     getattr(self, "draft_runner", None)
                 ):
                     self.cuda_graph_runner = None
+                # PROTOTYPE: the Flash drafter also takes this plugin-built
+                # draft-extend runner; SGLang's own HIP gate only admits aiter
+                # multi-step / DSV4 / DSA backends.
                 supports_plugin_graph = _is_dsv4_or_glm52_nextn_runner(
                     getattr(self, "draft_runner", None)
+                ) or _qwen4_draft_extend_graph_enabled(
+                    getattr(self, "draft_runner", None)
                 )
                 if (
                     getattr(self, "cuda_graph_runner_for_draft_extend", None)
@@ -924,24 +1001,31 @@ def _skip_draft_cuda_graphs(_self=self):
                             or 1
                         ),
                     )
-                    for backend in (
-                        getattr(
-                            getattr(self, "draft_runner", None), "attn_backend", None
-                        ),
-                        getattr(self, "draft_extend_attn_backend", None),
+                    # DSV4 / GLM-5.2 pad rows need a long fill; Flash QSA keeps
+                    # the backend default.
+                    if _is_dsv4_or_glm52_nextn_runner(
+                        getattr(self, "draft_runner", None)
                     ):
-                        if backend is not None and hasattr(
-                            backend, "_cuda_graph_seq_len_fill_value"
+                        for backend in (
+                            getattr(
+                                getattr(self, "draft_runner", None),
+                                "attn_backend",
+                                None,
+                            ),
+                            getattr(self, "draft_extend_attn_backend", None),
                         ):
-                            backend._cuda_graph_seq_len_fill_value = seq_len_fill
+                            if backend is not None and hasattr(
+                                backend, "_cuda_graph_seq_len_fill_value"
+                            ):
+                                backend._cuda_graph_seq_len_fill_value = seq_len_fill
                     self.cuda_graph_runner_for_draft_extend = (
                         EAGLEDraftExtendCudaGraphRunner(self)
                     )
                 elif supports_plugin_graph:
                     self.cuda_graph_runner_for_draft_extend = None
             except Exception as exc:
                 logger.warning(
-                    "Failed to enable DSV4 draft-extend cuda graph in ATOM plugin: %s",
+                    "Failed to enable draft-extend cuda graph in ATOM plugin: %s",
                     exc,
                 )
             return ret

atom/plugin/sglang/qwen4_exp_bridge.py

@@ -37,7 +37,7 @@
     resolve_attn_backend,
     resolve_mamba_req_pool,
 )
-from atom.utils import CpuGpuBuffer
+from atom.utils import CpuGpuBuffer, envs
 from atom.utils.forward_context import get_forward_context
 
 logger = logging.getLogger(__name__)
@@ -234,10 +234,53 @@ def bind_graph_ple(
 
 
 _DECODE_GRAPH = _Qwen4ExpDecodeGraphBuffers()
-# Draft stays eager (Flash mRoPE) but must not write the buffers baked into
+# The draft must not write the buffers baked into
 # the target verify graph. A second set still keeps page tables and seq_lens
 # on device, so a draft step does not allocate or sync seq_lens back to host.
 _DRAFT_QSA = _Qwen4ExpDecodeGraphBuffers()
+# PROTOTYPE: non-None marker returned to the draft runner's out-of-graph hook
+# before the first draft forward has populated _DRAFT_QSA.last_qsa.
+_DRAFT_GRAPH_METADATA_IN_GRAPH = object()
+
+
+def _is_draft_decode_batch(forward_batch: Any) -> bool:
+    """PROTOTYPE: draft-runner decode batch (EAGLE draft input), graphs enabled."""
+    if not envs.ATOM_SGLANG_QWEN4_DRAFT_CG:
+        return False
+    spec = getattr(forward_batch, "spec_info", None)
+    return type(spec).__name__ == "EagleDraftInput"
+
+
+def stage_draft_graph_mrope(mrope_buffer: torch.Tensor, forward_batch: Any) -> None:
+    """Stage the live mRoPE rows into a draft graph runner's ``[3, N]`` buffer.
+
+    SGLang's EAGLE draft decode / draft-extend runners capture
+    ``forward_batch.mrope_positions`` from that buffer but refresh only the
+    1-D ``positions`` on replay, and the buffer is shared by name and shape
+    with other runners. The Flash drafter reads the 3-row positions
+    (``_lm_positions_for_runtime``), so copy SGLang's live rows in (they carry
+    multimodal offsets), or broadcast the 1-D positions when SGLang computed
+    none. Rows past the batch are zeroed, as SGLang zeroes padded ``positions``.
+    """
+    live = getattr(forward_batch, "mrope_positions", None)
+    if live is None:
+        live = forward_batch.positions.unsqueeze(0).expand(3, -1)
+    num_tokens = int(live.shape[-1])
+    mrope_buffer[:, :num_tokens].copy_(live)
+    mrope_buffer[:, num_tokens:].zero_()
+
+
+def advance_draft_decode_mrope(forward_batch: Any) -> None:
+    """Advance mRoPE in step with EAGLE ``draft_forward``'s ``positions.add_(1)``.
+
+    SGLang moves only the 1-D positions after each draft decode step, so the
+    3-row positions the Flash drafter reads would stay at step 0. The plugin's
+    ``draft_forward`` wrapper calls this on the caller's batch after every draft
+    decode forward; inside a draft decode graph capture the add is recorded.
+    """
+    mrope = getattr(forward_batch, "mrope_positions", None)
+    if mrope is not None:
+        mrope.add_(1)
 
 
 def _is_draft_forward() -> bool:
@@ -1592,6 +1635,12 @@ def prepare_qwen4_exp_decode_graph_metadata(
     is_verify = _is_target_verify(forward_batch)
     if not (mode.is_decode_or_idle() or is_verify):
         return None
+    # PROTOTYPE: a Flash draft decode graph rebuilds its QSA metadata into
+    # _DRAFT_QSA inside the graph (positions / out_cache_loc change per draft
+    # step). Do not fill the target's _DECODE_GRAPH buffers from the draft
+    # runner's out-of-graph hook: TARGET_VERIFY replays from them.
+    if mode.is_decode_or_idle() and _is_draft_decode_batch(forward_batch):
+        return _DRAFT_QSA.last_qsa or _DRAFT_GRAPH_METADATA_IN_GRAPH
     if atom_config is None:
         try:
             from atom.config import get_current_atom_config
@@ -1653,7 +1702,12 @@ def qwen4_exp_metadata_from_forward_batch(
     that storage in place; there is no extra PLE copy.
     """
     bind_qsa_caches(model, forward_batch, atom_config)
-    if _is_capturing() and _DECODE_GRAPH.last_qsa is not None:
+    # Draft graphs (PROTOTYPE) build into _DRAFT_QSA inside the graph instead.
+    if (
+        _is_capturing()
+        and not _is_draft_forward()
+        and _DECODE_GRAPH.last_qsa is not None
+    ):
         ple = _DECODE_GRAPH.last_ple
         if ple is None:
             ple = _build_decode_graph_ple_metadata(

atom/utils/envs.py

@@ -234,6 +234,15 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float:
     "ATOM_SGLANG_EAGLE3_TP_VERIFY_BROADCAST": lambda: (
         os.getenv("ATOM_SGLANG_EAGLE3_TP_VERIFY_BROADCAST", "0") == "1"
     ),
+    # Qwen3.8-Flash-Next MTP on SGLang: capture the drafter's decode CUDA graph
+    # (default: the draft runs eagerly). Read by both the graph gate in
+    # atom/plugin/register.py and the QSA metadata hook in qwen4_exp_bridge.py.
+    "ATOM_SGLANG_QWEN4_DRAFT_CG": lambda: _flag_env("ATOM_SGLANG_QWEN4_DRAFT_CG"),
+    # With ATOM_SGLANG_QWEN4_DRAFT_CG, also capture the draft-extend graph that
+    # runs after every verify.
+    "ATOM_SGLANG_QWEN4_DRAFT_EXTEND_CG": lambda: _flag_env(
+        "ATOM_SGLANG_QWEN4_DRAFT_EXTEND_CG"
+    ),
     # Dynamic CPP follows SGLang's quadratic chunk-size predictor. This factor
     # interpolates between the configured initial chunk size (0) and the fitted
     # equal-latency size (1). Values around 0.6-0.85 are typically most stable.

docs/environment_variables.md

@@ -111,6 +111,8 @@ cannot affect a server that does not use the feature.
 |----------|------|---------|-------------|
 | **ATOM_DISABLE_VLLM_PLUGIN** | bool | 0 (false) | If set to `1`, disable the vLLM plugin registration entirely. |
 | **ATOM_V4_CG_INSTR_DIR** | str | — (off) | vLLM plugin, DeepSeek-V4 diagnostic: directory for per-process cudagraph dispatch counters (dispatched modes, steps whose uniform decode was overridden for a still-prefilling row, guard trips), one `cg_<pid>.json` per process, rewritten on every step. |
+| **ATOM_SGLANG_QWEN4_DRAFT_CG** | bool | 0 (false) | SGLang plugin, Qwen3.8-Flash-Next EAGLE MTP: capture the draft decode CUDA graph instead of running each draft step eagerly. Any value other than empty, `0`, `false`, `no` or `off` enables it. |
+| **ATOM_SGLANG_QWEN4_DRAFT_EXTEND_CG** | bool | 0 (false) | With `ATOM_SGLANG_QWEN4_DRAFT_CG`, also capture the draft-extend CUDA graph that runs after every verify. Ignored without `ATOM_SGLANG_QWEN4_DRAFT_CG`. |
 
 ## Kernel / backend selection