diff --git a/python/freetoken/spec/mtp/chain_graph.py b/python/freetoken/spec/mtp/chain_graph.py index 2b73e61f9..007a8b364 100644 --- a/python/freetoken/spec/mtp/chain_graph.py +++ b/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() diff --git a/python/freetoken/spec/mtp/shared_proposer.py b/python/freetoken/spec/mtp/shared_proposer.py index b0e9b7205..5b000b3a1 100644 --- a/python/freetoken/spec/mtp/shared_proposer.py +++ b/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] diff --git a/tests/spec/test_verify_contract.py b/tests/spec/test_verify_contract.py index 1c7a892ef..077a6bf6f 100644 --- a/tests/spec/test_verify_contract.py +++ b/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: