perf(spec): capture the gemma4 assistant draft in cuda graphs

python/freetoken/spec/mtp/chain_graph.py

@@ -15,6 +15,7 @@
     from freetoken.core import Req
 
     from .proposer import MTPProposer
+    from .shared_proposer import SharedKVMTPProposer
 
 
 class _TritonSteps:
@@ -163,3 +164,92 @@ def replay(self, reqs: List[Req], cs: List[int]) -> torch.Tensor | None:
         self.inputs[:, :padded].copy_(host, non_blocking=True)
         self.graphs[padded].replay()
         return self.drafts[:bs].clone()
+
+
+class SharedDraftGraphs:
+    """Replays ``SharedKVMTPProposer.propose``: the draft stores no KV, so every one of the k steps
+    reads the same [0, cached_len) window and one metadata set per batch size holds all k launches."""
+
+    def __init__(self, proposer: "SharedKVMTPProposer", bs_list: List[int], pool) -> None:
+        self.p = proposer
+        ctx = proposer.ctx
+        device, max_bs = ctx.device, max(bs_list)
+        self.page_table = ctx.page_table
+        self.dummy_row = ctx.dummy_req.table_idx
+        # inputs: the table row and its entry position; the pending token rides its own buffer
+        self.inputs = torch.zeros(2, max_bs, dtype=torch.int64, device=device)
+        self.token = torch.zeros(max_bs, dtype=torch.int64, device=device)
+        self.positions = torch.zeros(max_bs, dtype=torch.int32, device=device)
+        self.indptr = torch.zeros(1, max_bs + 1, dtype=torch.int32, device=device)
+        self.indices = torch.zeros(1, max_bs * ctx.page_table.shape[1], dtype=torch.int32, device=device)
+        self.arange = torch.arange(max_bs + 1, dtype=torch.int32, device=device)
+        self.drafts = torch.zeros(max_bs, proposer.k, dtype=torch.int64, device=device)
+        self.graphs: Dict[int, torch.cuda.CUDAGraph] = {}
+        # the metadata of these batches holds the decode scratch the graphs read
+        self._batches: Dict[int, Batch] = {}
+        for bs in sorted(bs_list, reverse=True):
+            # warmup rows gather one dummy key: a zero-length window is a decode path nothing else runs
+            self.inputs[0, :bs] = self.dummy_row
+            self.inputs[1, :bs] = 1
+            batch = self._batch(bs)
+            # outside the graph: triton compiles and the decode scratch is allocated here
+            self._run(bs, batch)
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, pool=pool, stream=ctx.stream):
+                self._run(bs, batch)
+            if pool is None:
+                pool = graph.pool()
+            self.graphs[bs] = graph
+            self._batches[bs] = batch
+
+    def _batch(self, bs: int) -> Batch:
+        batch = Batch(reqs=[None] * bs, phase="decode")
+        batch.padded_reqs = batch.reqs
+        batch.positions = self.positions[:bs]
+        batch.attn_metadata = TritonMetadata(
+            cu_seqlens_q_gpu=self.arange[: bs + 1],
+            indptr=self.indptr[0, : bs + 1],
+            indices=self.indices[0],
+            q_to_req=self.arange[:bs],
+            q_positions=batch.positions,
+            is_decode=True,
+            prefix_lens=batch.positions,
+            max_q_len=1,
+        )
+        return batch
+
+    def _run(self, bs: int, batch: Batch) -> None:
+        p = self.p
+        rows, entry = self.inputs[:, :bs]
+        # the window [0, entry) serves every step: the draft writes no KV between them, so only the
+        # carried token and hidden move
+        self.positions[:bs] = entry
+        self.indptr[0, 1 : bs + 1] = torch.cumsum(entry, 0)
+        gather_kv_indices(self.page_table, rows, self.indptr[:, : bs + 1], self.indices)
+        if getattr(p.kv, "swa_paged", False):
+            batch.attn_metadata.swa_indices = p.kv.translate_loc_from_full_to_swa(self.indices[0])
+        token = self.token[:bs]
+        hidden = p.seed_hidden[rows]
+        for step in range(p.k):
+            out = p._forward(batch, p.chain_backend, p.draft.embed_tokens.forward(token), hidden)
+            token = p._draft_token(out)
+            hidden = p._relayed(out)
+            self.drafts[:bs, step] = token
+
+    def replay(self, reqs: List[Req], cs: List[int], pending: torch.Tensor) -> torch.Tensor | None:
+        """The drafts [bs, k], or None when no captured batch size fits."""
+        bs = len(reqs)
+        padded = next((b for b in sorted(self.graphs) if b >= bs), None)
+        if padded is None:
+            return None
+        pad = padded - bs
+        # padded rows draft over the dummy page's first key: their drafts are discarded
+        host = torch.tensor(
+            [[r.table_idx for r in reqs] + [self.dummy_row] * pad, cs + [1] * pad],
+            dtype=torch.int64, pin_memory=True,
+        )
+        self.inputs[:, :padded].copy_(host, non_blocking=True)
+        self.token[bs:padded].zero_()
+        self.token[:bs].copy_(pending)
+        self.graphs[padded].replay()
+        return self.drafts[:bs].clone()

python/freetoken/spec/mtp/shared_proposer.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import os
 from itertools import accumulate
 from typing import TYPE_CHECKING, List, Protocol
 
@@ -12,12 +13,28 @@
 from freetoken.spec.mtp.proposer import _DraftReq, _PrefillSink, _VerifySink, with_resident_moe
 from freetoken.spec.mtp.shared_kv import ReadOnlyTargetKVView
 from freetoken.spec.proposer import DraftContext, Proposer, SpecProposal, check_draft_path
+from freetoken.utils import init_logger
 
 if TYPE_CHECKING:
     from freetoken.core import Req
     from freetoken.layers.embedding import VocabParallelEmbedding
     from freetoken.models.config import ModelConfig
 
+    from .chain_graph import SharedDraftGraphs
+
+logger = init_logger(__name__)
+
+
+# the off switch keeps the draft eager for the contract test's per-step logit instrumentation
+def _shared_draft_graphs() -> bool:
+    raw = os.environ.get("FREETOKEN_SHARED_DRAFT_GRAPHS", "1")
+    if raw not in ("0", "1"):
+        raise ValueError(f"FREETOKEN_SHARED_DRAFT_GRAPHS must be 0 or 1, got {raw!r}")
+    return raw == "1"
+
+
+_SHARED_DRAFT_GRAPHS = _shared_draft_graphs()
+
 
 class SharedKVMTPDraft(Protocol):
     """The class a family registers as ``ModelSpec.spec_drafts["shared_kv_mtp"]``. It owns its vocab table and stores no KV."""
@@ -117,6 +134,8 @@ def __init__(self, config, target_model) -> None:
             raise RuntimeError(f"{draft_cls.__name__} must build its own MoE layers, resident on the GPU")
         self._prefill_sink: _PrefillSink | None = None
         self.kv = None
+        self.draft_graphs: SharedDraftGraphs | None = None
+        self.draft_graph_fallbacks = 0
 
     # ------------------------------------------------------------------ lifecycle
 
@@ -126,7 +145,9 @@ def attach(self, ctx: DraftContext) -> None:
         super().attach(ctx)
         cfg = self.draft.config
         rows = self.config.max_running_req + 1
-        # a rebuild swaps the target pool, so the read-only view is rebuilt with every attach
+        # a rebuild swaps the target pool, so the read-only view is rebuilt with every attach; the
+        # graphs bind this view, this page table and this backend, so they are re-captured after it
+        self.draft_graphs = None
         self.kv = ReadOnlyTargetKVView(ctx.kv_cache, self.draft.config, self.config.model_config)
         width = self.draft.relay_width(self.config.model_config)
         # the target hidden that seeds the next proposal, one row per table row
@@ -138,6 +159,7 @@ def attach(self, ctx: DraftContext) -> None:
             self.chain_backend = create_attention_backend("triton", cfg)
 
     def detach(self, *, keep_kv: bool) -> None:
+        self.draft_graphs = None  # the graphs hold the old pool's slabs alive; drop them before the rebuild frees it
         self.kv = None  # a view over the old pool is dead once a rebuild swaps it, keep_kv included
         self.seed_hidden = self.verify_hidden = None
         self.chain_backend = None
@@ -183,11 +205,26 @@ def _relayed(self, out: torch.Tensor) -> torch.Tensor:
     # ------------------------------------------------------------------ one step
 
     def capture_graphs(self, bs_list: List[int], pool) -> None:
-        # read-only drafts stay eager until token-identity is proven; capture is mechanical after that
-        return
+        from .chain_graph import SharedDraftGraphs
+
+        self.draft_graphs = None
+        if not bs_list or not _SHARED_DRAFT_GRAPHS:
+            return
+        self.draft_graphs = SharedDraftGraphs(self, bs_list, pool)
+        logger.info_rank0(f"speculative decoding: shared draft cuda graphs captured for bs {sorted(bs_list)}")
 
     def propose(self, reqs: List[Req], rows: torch.Tensor, pending: torch.Tensor) -> SpecProposal:
         cs = [r.cached_len for r in reqs]
+        if self.draft_graphs is not None:
+            drafts = self.draft_graphs.replay(reqs, cs, pending)
+            if drafts is not None:
+                return SpecProposal(drafts)
+            self.draft_graph_fallbacks += 1
+            if self.draft_graph_fallbacks == 1:
+                logger.warning_rank0(
+                    f"speculative decoding: no shared draft graph fits bs {len(reqs)} "
+                    f"(captured {sorted(self.draft_graphs.graphs)}), drafting eagerly"
+                )
         table_idxs = [r.table_idx for r in reqs]
         token = pending
         hidden = self.seed_hidden[rows]

tests/spec/test_verify_contract.py

@@ -126,17 +126,19 @@ def _wrapped(self, *args, _original=_original, _key=_key):
 
 # the draft side: whether the drafts came from the draft's own head reading the target's slabs
 from freetoken.models.gemma4.mtp import Gemma4MTPDraft
+from freetoken.spec.mtp.chain_graph import SharedDraftGraphs
 from freetoken.spec.mtp.shared_kv import ReadOnlyTargetKVView
 
-draft = {"head_calls": 0, "reads": {}, "top3": []}
+draft = {"head_calls": 0, "reads": {}, "top3": [], "stashed": []}
 draft_logits = Gemma4MTPDraft.draft_logits
 
 
 def counting_draft_logits(self, hidden):
     draft["head_calls"] += 1
     logits = draft_logits(self, hidden)
+    # tensors are stashed, never listed here: a host sync inside this head would poison graph capture
     top = logits.float().topk(3, dim=-1)
-    draft["top3"].append((top.indices.tolist(), top.values.tolist()))
+    draft["stashed"].append((top.indices, top.values))
     return logits
 
 
@@ -152,17 +154,36 @@ def counting_k_cache(self, index):
 
 ReadOnlyTargetKVView.k_cache = counting_k_cache
 
+# a replayed draft graph runs no python, so its proposals are invisible to the wrappers above: count
+# the replays themselves to pin that the graphs served the run
+draft_graph_replays = 0
+replay_orig = SharedDraftGraphs.replay
+
+
+def counting_replay(self, reqs, cs, pending):
+    global draft_graph_replays
+    out = replay_orig(self, reqs, cs, pending)
+    draft_graph_replays += out is not None
+    return out
+
+
+SharedDraftGraphs.replay = counting_replay
+
 # the identity sidecar's drafts are known exactly, so record each proposal to pin the seed row
 propose_orig = SharedKVMTPProposer.propose
 proposals = []
 
 
 def recording_propose(self, reqs, rows, pending):
+    before = len(draft["stashed"])
     proposal = propose_orig(self, reqs, rows, pending)
-    # this call's k head calls: per request row, one [ids, vals] top-3 per step
-    tops = draft["top3"][len(draft["top3"]) - self.k:]
+    # this call's k head calls: per request row, one [ids, vals] top-3 per step; a replayed graph
+    # runs no head, so its proposals carry only the drafts, compared against the eager run's steps
+    fresh = draft["stashed"][before:]
+    tops = [(i.tolist(), v.tolist()) for i, v in fresh[-self.k:]] if len(fresh) >= self.k else []
+    drafted = proposal.draft_token_ids.tolist()
     proposals.append([
-        (r.uid, p, [[ids[i], vals[i]] for ids, vals in tops])
+        (r.uid, p, [[ids[i], vals[i]] for ids, vals in tops], drafted[i])
         for r, p, i in zip(reqs, pending.tolist(), range(len(reqs)))
     ])
     return proposal
@@ -337,6 +358,9 @@ def __init__(self, num_rows, head_dim, dtype=torch.bfloat16):
 params = [SamplingParams(temperature=0.0, max_tokens=n, ignore_eos=True, presence_penalty=p, frequency_penalty=p)
           for n, p in zip((48, 37, 41), penalties)]
 res = llm.generate(prompts, params)
+draft["top3"] = [(i.tolist(), v.tolist()) for i, v in draft["stashed"]]
+proposer = getattr(getattr(llm.engine, "spec", None), "proposer", None)
+graphs_of = getattr(proposer, "draft_graphs", None)
 pool = llm.engine.linear_state_pool
 window = list(pool.slot_state("dflash_window").shape) if pool is not None and pool.has_slot_state("dflash_window") else None
 cache = llm.cache_manager
@@ -345,6 +369,8 @@ def __init__(self, num_rows, head_dim, dtype=torch.bfloat16):
            "swa_window": cache.sliding_window_size, "swa": cache.swa_paged, "rows": llm.config.decode_rows_per_req,
            "prompt_lens": [len(p) for p in prompts], "penalties": list(penalties),
            "evict": evict, "proposals": proposals,
+           "draft_graphs": sorted(graphs_of.graphs) if graphs_of is not None else [],
+           "draft_graph_replays": draft_graph_replays,
            "draft": {"head_calls": draft["head_calls"], "reads": sorted(draft["reads"].items())}}, open(out, "w"))
 """
 
@@ -593,9 +619,37 @@ def test_verify_steps_reproduce_greedy_decode(tmp_path, plain_runs, monkeypatch,
         draft = {"speculative_draft_model_path": _tiny_mtp_draft(str(tmp_path / "draft"), target, family)}
         extra = {"use_dummy_weight": False}
         model = target
+    if algorithm == "mtp" and family == "gemma4":
+        # the strict pins below read the eager draft's per-step logits and head calls, which a replay
+        # hides; the graphs run is pinned against this one by its replayed drafts, not just its output
+        monkeypatch.setenv("FREETOKEN_SHARED_DRAFT_GRAPHS", "0")
     spec = _run(model, str(tmp_path / "spec.json"), decode, speculative_algorithm=algorithm, speculative_num_steps=K,
                 **draft, **extra, **kwargs)
+    monkeypatch.delenv("FREETOKEN_SHARED_DRAFT_GRAPHS", raising=False)
     if algorithm == "mtp" and family == "gemma4":
+        graphs = _run(model, str(tmp_path / "spec-graphs.json"), decode, speculative_algorithm=algorithm,
+                      speculative_num_steps=K, **draft, **extra, **kwargs)
+        assert spec["draft_graphs"] == [] and graphs["draft_graphs"], (spec["draft_graphs"], graphs["draft_graphs"])
+        # every proposal of the graphs run replayed a captured draft, or the graphs prove nothing
+        assert graphs["draft_graph_replays"] == graphs["verify_steps"], graphs
+        assert graphs["outputs"] == spec["outputs"] and graphs["verify_steps"] == spec["verify_steps"], (
+            "the draft graphs changed the run")
+        # rejection sampling can hide a draft difference behind an identical output, so compare the
+        # drafts per step, tolerating only a reorder inside the eager head's top-3 within two fp16 ulps
+        assert len(graphs["proposals"]) == len(spec["proposals"]), (len(graphs["proposals"]), len(spec["proposals"]))
+        for call, replayed in zip(spec["proposals"], graphs["proposals"]):
+            eager = {uid: (pending, steps, drafts) for uid, pending, steps, drafts in call}
+            graph = {uid: (pending, drafts) for uid, pending, _, drafts in replayed}
+            assert eager.keys() == graph.keys(), (eager.keys(), graph.keys())
+            for uid, (pending, steps, drafts) in eager.items():
+                assert graph[uid][0] == pending, (uid, pending, graph[uid][0])
+                for step, (a, b) in enumerate(zip(drafts, graph[uid][1])):
+                    ids, vals = steps[step]
+                    ulp = 2.0 ** (math.floor(math.log2(abs(vals[0]) or 1.0)) - 10)
+                    by_id = dict(zip(ids, vals))
+                    pair = {a, b}
+                    assert pair <= set(by_id) and max(by_id[t] for t in pair) - min(by_id[t] for t in pair) <= 2 * ulp, (
+                        uid, step, a, b, ids, vals)
         # the draft must earn its keep: its own head produced every draft (k per proposal), and each
         # forward read the target's slabs through the view (3 SWA layers -> 2, the full layer -> 3)
         assert spec["proposals"] and spec["draft"]["head_calls"] == len(spec["proposals"]) * K, spec["draft"]
@@ -605,7 +659,7 @@ def test_verify_steps_reproduce_greedy_decode(tmp_path, plain_runs, monkeypatch,
         # not be that row's argmax, and a row whose candidates sit within two fp16 ulps may still
         # flip under the head's own tiling
         for call in spec["proposals"]:
-            for uid, pending, steps in call:
+            for uid, pending, steps, _ in call:
                 if spec["penalties"][uid]:
                     continue
                 for ids, vals in steps: