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

@@ -38,10 +38,12 @@ 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, which sits at
+    ``batch_dim`` (``1`` for MRoPE positions, ``[3, N]``)."""
 
     shape: tuple[int, ...] = ()
     dtype: torch.dtype = torch.int32
+    batch_dim: int = 0
 
 
 @dataclass
@@ -113,7 +115,11 @@ def bind(self, config: "Config", device) -> "DraftGraph":
         )
         self._buffers = {
             role: torch.zeros(
-                (config.max_num_seqs, *staged.shape),
+                (
+                    *staged.shape[: staged.batch_dim],
+                    config.max_num_seqs,
+                    *staged.shape[staged.batch_dim :],
+                ),
                 dtype=staged.dtype,
                 device=device,
             )
@@ -147,7 +153,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 = {t.shape[self.inputs[role].batch_dim] 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,15 +178,16 @@ 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)}"
         )
-        buf = self._buffers[role]
-        assert 0 <= bs <= buf.shape[0], (
+        buf, d = self._buffers[role], self.inputs[role].batch_dim
+        assert 0 <= bs <= buf.shape[d], (
             f"{self.name} requested {bs} rows from draft input {role!r}, "
-            f"whose capacity is {buf.shape[0]}"
+            f"whose capacity is {buf.shape[d]}"
         )
-        return buf[:bs]
+        return buf.narrow(d, 0, 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 index along ``batch_dim``).
 
         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
@@ -185,8 +197,8 @@ def _stage_one(self, role: str, running_bs: int, src: torch.Tensor | None):
         ``src=None`` hands back the view unwritten: warmup wants a well-formed
         batch, not a particular one.
         """
-        buf = self._buffers[role]
-        out = buf[:running_bs]
+        buf, d = self._buffers[role], self.inputs[role].batch_dim
+        out = buf.narrow(d, 0, running_bs)
         if src is None:
             return out
         assert src.dtype == buf.dtype, (
@@ -195,12 +207,11 @@ def _stage_one(self, role: str, running_bs: int, src: torch.Tensor | None):
         )
         assert src.ndim == buf.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"but this pass stages {self.inputs[role].shape} per sequence with "
+            f"the batch at axis {d}"
         )
-        scheduled_bs = src.shape[0]
-        dst = out[:scheduled_bs]
+        scheduled_bs = src.shape[d]
+        dst = out.narrow(d, 0, 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 +221,9 @@ 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
+            out.narrow(d, scheduled_bs, running_bs - scheduled_bs).copy_(
+                src.narrow(d, scheduled_bs - 1, 1)
+            )  # broadcasts; no expand needed
         return out
 
     def is_captured(self, running_bs: int) -> bool:

atom/spec_decode/eagle_proposer.py

@@ -126,9 +126,15 @@ 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:
+        ``StagedInput.batch_dim`` stages them with the batch at axis 1. Other
+        MRoPE drafts have not been validated with draft graphs and stay eager.
         """
         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 +154,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_dim=1)
+                if mrope
+                else StagedInput(dtype=torch.int64)
+            ),
             "hidden_states": StagedInput(
                 shape=hidden_shape,
                 dtype=self.dtype,
@@ -768,7 +778,7 @@ def propose(
                 ):
                     # (`_step_forward`) Step 0 of a decode runs every row but
                     # hands back only each sequence's last one; a step 1+ that
-                    # no declared pass covers (MRoPE) runs here, eagerly.
+                    # no declared pass covers runs here, eagerly.
                     ret_hidden_states, head_input = self.model.forward_with_head(
                         d_input_ids,
                         d_positions,
@@ -848,10 +858,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,18 +179,55 @@ 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_dim=1`. 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.
     """
     g = _graph()
-    with pytest.raises(AssertionError, match="leading axis is not the batch"):
+    with pytest.raises(AssertionError, match="batch at axis 0"):
         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_dim=1),
+            "row": StagedInput(dtype=torch.int32),
+        }
+    )
+
+
+def test_a_batch_dim_1_input_stages_columns_and_tail_repeats_the_last_one():
+    """MRoPE positions stage as the first `bs` columns of one `[3, max_num_seqs]`
+    buffer, 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)
+    assert pos.stride() == (g._buffers["pos"].shape[1], 1)
+    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_dim_1_producer_view_is_the_front_of_the_staged_one():
+    """`buffer(role, scheduled_bs)` is a prefix of the `running_bs` view at the
+    same stride, so staging a fill of it pads after it and copies nothing over
+    it."""
+    g = _mrope_graph()
+    fill = g.buffer("pos", 2)
+    fill.copy_(torch.tensor([[1, 2], [3, 4], [5, 6]]))
+    pos = g.stage(4, {"pos": fill})["pos"]
+    assert pos.data_ptr() == fill.data_ptr() and pos.stride() == fill.stride()
+    assert pos.tolist() == [[1, 2, 2, 2], [3, 4, 4, 4], [5, 6, 6, 6]]
+
+
 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 +246,17 @@ def test_sources_that_disagree_on_the_batch_are_refused():
         )
 
 
+def test_batch_dim_1_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_with_the_batch_at_axis_1(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_dim=1
+    )
+    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