feat(spec): capture the Qwen3.8 MTP draft steps in CUDA graphs (MRoPE)

atom/models/qwen4_exp_mtp.py

@@ -155,6 +155,13 @@ class Qwen4ExpMTP(nn.Module):
     disable_fused_shared_loading = True
     fused_shared_experts = 0
     get_expert_mapping = Qwen4ExpForConditionalGeneration.get_expert_mapping
+    draft_graph_supports_mrope = True
+    reuse_draft_graph_step_buffers = True
+
+    @staticmethod
+    def draft_graph_hidden_state_shape(draft_hf) -> tuple[int, ...]:
+        """The draft carries the whole hyper-connection bundle between steps."""
+        return (draft_hf.hc_count, draft_hf.hidden_size)
 
     def __init__(self, atom_config: Config):
         super().__init__()

atom/spec_decode/draft_graph.py

@@ -7,6 +7,7 @@
 passes still live with the flavor; only the machine is here.
 """
 
+import math
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
@@ -38,10 +39,16 @@ def _owned(value: Any, staged: tuple[torch.Tensor, ...] = ()) -> Any:
 
 @dataclass(frozen=True)
 class StagedInput:
-    """One staged input's type: everything but the leading batch axis."""
+    """One staged input's type: everything but the batch axis.
+
+    The batch axis leads unless ``batch_last`` (MRoPE positions, ``[3, N]``):
+    each batch then views a flat buffer's front as a contiguous ``[*shape, bs]``,
+    the layout the runner gives its own MRoPE positions.
+    """
 
     shape: tuple[int, ...] = ()
     dtype: torch.dtype = torch.int32
+    batch_last: bool = False
 
 
 @dataclass
@@ -113,14 +120,21 @@ def bind(self, config: "Config", device) -> "DraftGraph":
         )
         self._buffers = {
             role: torch.zeros(
-                (config.max_num_seqs, *staged.shape),
+                (
+                    (config.max_num_seqs * math.prod(staged.shape),)
+                    if staged.batch_last
+                    else (config.max_num_seqs, *staged.shape)
+                ),
                 dtype=staged.dtype,
                 device=device,
             )
             for role, staged in self.inputs.items()
         }
         return self
 
+    def _batch_of(self, role: str, src: torch.Tensor) -> int:
+        return src.shape[-1] if self.inputs[role].batch_last else src.shape[0]
+
     @property
     def will_capture(self) -> bool:
         """Capture only a supported pass when the deployment enables it."""
@@ -147,7 +161,12 @@ def stage(
         How many rows are real comes from the sources, never from the caller: a
         count that cannot disagree with the tensor it describes.
         """
-        counts = {t.shape[0] for t in srcs.values()}
+        undeclared = sorted(srcs.keys() - self._buffers.keys())
+        assert not undeclared, (
+            f"{self.name} has no staged inputs {undeclared}; "
+            f"declared roles are {sorted(self._buffers)}"
+        )
+        counts = {self._batch_of(role, t) for role, t in srcs.items()}
         assert len(counts) <= 1, (
             f"{self.name} was handed batches of {sorted(counts)}; they "
             f"describe one step, so one of them is not this step's"
@@ -167,6 +186,10 @@ def buffer(self, role: str, bs: int) -> torch.Tensor:
             f"{self.name} has no staged input {role!r}; "
             f"declared roles are {sorted(self._buffers)}"
         )
+        assert not self.inputs[role].batch_last, (
+            f"{self.name} stages {role!r} batch-last, whose view at one batch is "
+            f"not the front of a wider one; hand it to stage instead"
+        )
         buf = self._buffers[role]
         assert 0 <= bs <= buf.shape[0], (
             f"{self.name} requested {bs} rows from draft input {role!r}, "
@@ -175,7 +198,8 @@ def buffer(self, role: str, bs: int) -> torch.Tensor:
         return buf[:bs]
 
     def _stage_one(self, role: str, running_bs: int, src: torch.Tensor | None):
-        """Copy ``src`` into this input's fixed buffer, tail-repeating its last row.
+        """Copy ``src`` into this input's fixed buffer, tail-repeating its last row
+        (its last column if ``batch_last``).
 
         Every input repeats the same last index, so a pad row is a coherent copy
         of the last real one and only redoes work that row already does. Zeros
@@ -186,21 +210,25 @@ def _stage_one(self, role: str, running_bs: int, src: torch.Tensor | None):
         batch, not a particular one.
         """
         buf = self._buffers[role]
-        out = buf[:running_bs]
+        shape, batch_last = self.inputs[role].shape, self.inputs[role].batch_last
+        out = (
+            buf[: math.prod(shape) * running_bs].view(*shape, running_bs)
+            if batch_last
+            else buf[:running_bs]
+        )
         if src is None:
             return out
         assert src.dtype == buf.dtype, (
             f"draft input '{role}' of {self.name} arrived as {src.dtype}, but "
             f"its storage is {buf.dtype}; a capture has that address baked"
         )
-        assert src.ndim == buf.ndim, (
+        assert src.ndim == out.ndim, (
             f"draft input '{role}' of {self.name} arrived as {tuple(src.shape)}, "
-            f"whose leading axis is not the batch; this pass stages "
-            f"{tuple(buf.shape[1:])} per sequence. MRoPE positions ([3, N]) are "
-            f"the case that reaches here"
+            f"whose {'last' if batch_last else 'leading'} axis is not the batch; "
+            f"this pass stages {shape} per sequence"
         )
-        scheduled_bs = src.shape[0]
-        dst = out[:scheduled_bs]
+        scheduled_bs = self._batch_of(role, src)
+        dst = out[..., :scheduled_bs] if batch_last else out[:scheduled_bs]
         assert src.shape == dst.shape, (
             f"draft input '{role}' of {self.name} arrived as {tuple(src.shape)}, "
             f"but its fixed storage expects {tuple(dst.shape)}"
@@ -210,7 +238,10 @@ def _stage_one(self, role: str, running_bs: int, src: torch.Tensor | None):
         if src.data_ptr() != dst.data_ptr() or src.stride() != dst.stride():
             dst.copy_(src)
         if running_bs > scheduled_bs:
-            out[scheduled_bs:].copy_(src[-1])  # broadcasts; no expand needed
+            if batch_last:
+                out[..., scheduled_bs:].copy_(src[..., -1:])
+            else:
+                out[scheduled_bs:].copy_(src[-1])  # broadcasts; no expand needed
         return out
 
     def is_captured(self, running_bs: int) -> bool:

atom/spec_decode/eagle_proposer.py

@@ -126,9 +126,14 @@ def _declare_draft_graphs(self):
         pad and nothing a graph could pin. Declaring anyway used to warm and
         capture a ONE-dimensional-positions shape that serving never runs --
         graphs nobody replays, paid for at every startup.
+        A draft model that sets ``draft_graph_supports_mrope`` is the exception:
+        it stages them batch-last.
         """
         self.step = None
-        if self.runner.use_mrope or self.mtp_k < 2:
+        mrope = self.runner.use_mrope
+        if (
+            mrope and not getattr(self.model, "draft_graph_supports_mrope", False)
+        ) or self.mtp_k < 2:
             # mtp_k == 1 has no step 1+, so warming one would capture a graph
             # `propose` can never reach.
             self._reuse_step_buffers = False
@@ -148,7 +153,11 @@ def _declare_draft_graphs(self):
             # `input_ids` from that buffer to this one, and `stage` asserts the
             # two agree.
             "input_ids": StagedInput(dtype=torch.int32),
-            "positions": StagedInput(dtype=torch.int64),
+            "positions": (
+                StagedInput(shape=(3,), dtype=torch.int64, batch_last=True)
+                if mrope
+                else StagedInput(dtype=torch.int64)
+            ),
             "hidden_states": StagedInput(
                 shape=hidden_shape,
                 dtype=self.dtype,
@@ -809,10 +818,6 @@ def propose(
                         # from it, so it has to be `running_bs` long before THAT --
                         # padding it with the other rows, just before the
                         # forward, is a step too late.
-                        #
-                        # No MRoPE case to guard: a model whose positions keep
-                        # the token axis last declares no pass, so `self.step` is
-                        # None and these stay the target's own 2-D positions.
                         if self.step is not None:
                             positions = self.step.stage(
                                 running_bs, {"positions": positions}

tests/test_draft_graph.py

@@ -157,6 +157,14 @@ def test_fixed_storage_accessor_checks_role_and_capacity():
         g.buffer("row", 257)
 
 
+def test_staging_an_undeclared_role_names_it():
+    g = _graph()
+    with pytest.raises(
+        AssertionError, match=r"\['missing'\]; declared roles are \['row'\]"
+    ):
+        g.stage(1, {"missing": torch.zeros(1, dtype=torch.int32)})
+
+
 def test_staging_rejects_a_wrong_trailing_shape_before_identity_check():
     g = _graph(inputs={"row": StagedInput(shape=(4,), dtype=torch.int32)})
     wrong = g._buffers["row"][:2, :1]
@@ -171,10 +179,10 @@ def test_staging_a_wrong_dtype_fails_loudly():
 
 
 def test_staging_a_source_whose_leading_axis_is_not_the_batch_fails_loudly():
-    """Every staged source is indexed as `[rows, *StagedInput.shape]`.
+    """Every leading-batch source is indexed as `[rows, *StagedInput.shape]`.
 
     MRoPE positions are the live case: they keep the token axis LAST (`[3, N]`),
-    so a drafter must not hand them to a pass at all. Without this the mismatch
+    so a drafter must declare them `batch_last`. Without this the mismatch
     surfaced as a `copy_` size error two frames down, which reads like a
     batch-size bug rather than a layout one.
     """
@@ -183,6 +191,36 @@ def test_staging_a_source_whose_leading_axis_is_not_the_batch_fails_loudly():
         g._stage_one("row", 8, torch.zeros(3, 8, dtype=torch.int32))
 
 
+def _mrope_graph():
+    return _graph(
+        inputs={
+            "pos": StagedInput(shape=(3,), dtype=torch.int64, batch_last=True),
+            "row": StagedInput(dtype=torch.int32),
+        }
+    )
+
+
+def test_batch_last_input_is_contiguous_and_tail_repeats_the_last_column():
+    """MRoPE positions stage as a contiguous `[3, bs]`, the runner's own layout,
+    beside a leading-batch input a producer filled in place (`hidden_states`)."""
+    g = _mrope_graph()
+    src = torch.arange(6, dtype=torch.int64).view(3, 2) * 10
+    g.buffer("row", 2).copy_(torch.tensor([1, 2], dtype=torch.int32))
+    staged = g.stage(4, {"pos": src, "row": g.buffer("row", 2)})
+    pos = staged["pos"]
+    assert pos.shape == (3, 4) and pos.is_contiguous()
+    assert pos.tolist() == [[0, 10, 10, 10], [20, 30, 30, 30], [40, 50, 50, 50]]
+    assert staged["row"].tolist() == [1, 2, 2, 2]
+    assert pos.data_ptr() == g._buffers["pos"].data_ptr()
+
+
+def test_a_batch_last_input_has_no_producer_view():
+    """Its view at `scheduled_bs` overlaps the `running_bs` one at another
+    stride, so staging a fill of it would copy the storage onto itself."""
+    with pytest.raises(AssertionError, match="batch-last"):
+        _mrope_graph().buffer("pos", 2)
+
+
 def test_sources_that_disagree_on_the_batch_are_refused():
     """They describe one step, so one of them is not this step's."""
     g = _graph(
@@ -201,6 +239,17 @@ def test_sources_that_disagree_on_the_batch_are_refused():
         )
 
 
+def test_batch_last_sources_must_agree_with_leading_ones():
+    with pytest.raises(AssertionError, match="batches of"):
+        _mrope_graph().stage(
+            4,
+            {
+                "pos": torch.zeros(3, 2, dtype=torch.int64),
+                "row": torch.zeros(3, dtype=torch.int32),
+            },
+        )
+
+
 def test_warmup_runs_the_epilogue_too_not_just_the_capturable_forward(monkeypatch):
     """Warming only the capturable forward leaves the rest of the pass cold.
 

tests/test_qwen4_exp_mtp.py

@@ -19,6 +19,7 @@
     qsa_select_paged_tokens,
     qsa_sparse_paged_gqa,
 )
+from atom.spec_decode.draft_graph import StagedInput
 
 
 def test_qwen_mtp_does_not_import_deepseek_when_pcp_disabled(monkeypatch):
@@ -290,6 +291,37 @@ def prepare(self, batch, running_bs, running_tokens, max_seqlen_q):
         assert not graph_positions[:, tokens:].count_nonzero()
 
 
+@pytest.mark.parametrize("opted_in", [True, False])
+def test_mrope_draft_graph_declares_positions_batch_last(opted_in):
+    """Qwen3.8's MTP head declares its steps 1+ with `[3, bs]` positions; an
+    MRoPE draft model that has not opted in still declares no pass."""
+    from atom.models.qwen4_exp_mtp import Qwen4ExpMTP
+    from atom.spec_decode.eagle_proposer import EagleProposer
+
+    proposer = SimpleNamespace(
+        runner=SimpleNamespace(use_mrope=True),
+        mtp_k=3,
+        speculative_config=SimpleNamespace(
+            draft_model_hf_config=SimpleNamespace(hidden_size=8, hc_count=2)
+        ),
+        model=object.__new__(Qwen4ExpMTP) if opted_in else SimpleNamespace(),
+        dtype=torch.bfloat16,
+        _step_forward=lambda *a, **kw: None,
+        _step_head=lambda *a, **kw: None,
+        _step_warmup_inputs=lambda *a, **kw: None,
+    )
+    graphs = EagleProposer._declare_draft_graphs(proposer)
+    if not opted_in:
+        assert graphs == () and proposer.step is None
+        return
+    (graph,) = graphs
+    assert graph.inputs["positions"] == StagedInput(
+        shape=(3,), dtype=torch.int64, batch_last=True
+    )
+    assert graph.inputs["hidden_states"].shape == (2, 8)
+    assert proposer._reuse_step_buffers
+
+
 def test_draft_qsa_view_drops_only_graph_padding():
     from atom.model_ops.attentions.qwen4_exp_attn import Qwen4ExpQSAMetadata