From 2f99faf20087a38730ccc8afcff211d03b48c29f Mon Sep 17 00:00:00 2001 From: Jason Cook Date: Thu, 5 Nov 2026 13:12:47 -0500 Subject: [PATCH] Share the first GLM draft pass's sparse selection with later passes The GLM-5.3-Flash checkpoints declare `index_share_for_mtp_iteration: true`: every draft pass of a speculative round attends through the sparse selection of the round's first pass. Transformers has no GLM MTP layer and ignores the field. vLLM's GLM MTP proposer runs the draft indexer at the first draft step and reuses its top-k selection for later steps when prefill context parallelism is off. Flare's draft passes ran the MTP layer's indexer, and with context parallelism its candidate exchange, on every pass. When the checkpoint configuration declares the field true, rounds with two or more drafts run the MTP indexer only on their first draft pass and copy each request's physical selection and selected length into round-local buffers. Before each later pass, every proposing request's selection is placed in its packed draft row, and idle rows select nothing. Later passes skip the index query and gate projections, query quantization, C4 pool update, selection, expansion, and context-parallel candidate all-gather. They still write their latent records, and they attend to records at or before the first pass's position rather than to the records of the round's earlier drafts. Only the first pass advances the draft C4 tail ring, which the round restores before absorption. Target verification and absorption are unchanged, so the served distribution is unchanged; draft proposals and acceptance follow the declared selection semantics. Checkpoints that do not declare the field run the indexer on every pass. docs/mtp-design.md describes the shared selection. Compatibility: no package metadata or prepared-contract change. Exported packages keep running the indexer on every draft pass; re-export prepares the shared-selection draft graphs. Validation on main (972e28e4) with this change, on one RTX PRO 6000 (SM120): - tests/recipes/test_glm53_mla_cuda.py: test_reused_selection_attends_only_to_the_first_pass_records passed in 4 cases (TP2 and TP4 head partitions, query-chain overlap on and off) and skipped 2 (TP1), with layer 3 MLA weights of GLM-5.3-Flash-NVFP4-MXFP8-CSF-QAD. With 256 cached records, a first pass at position 256 selects 257 records. A reusing pass at position 257 leaves the selection, C4 cache and C4 tail unchanged, writes its latent record, and produces zero output for the idle row. Its output matches a fresh selection at position 256 within 1e-3 relative norm (bitwise in all four cases) and differs by 11 to 17 percent relative norm from a selection at its own position. The file's six TP1 native-replay cases fail during b12x BMM tuning ("BMM plan requires a supported BF16 x MXFP8 geometry"), identically without this change. - tests/recipes/test_glm53_speculative.py with the other GLM recipe, speculative-capacity and prepared-contract CPU tests: 223 passed, 47 skipped. When a request stops proposing and the packed rows shift, each remaining request receives its own first-pass selection. Six tests/recipes/test_declarations.py cases fail because the Kimi and Qwen checkpoints they read are absent on the test host; they fail identically without this change. - DGX Spark (GB10, SM121): GLM-5.3-Flash packages exported from a tree that contains this change and other GLM recipe changes, without context parallelism, served at TP2 on two hosts and TP4 on four hosts with five drafts and adaptive verification. Each package passed a 32-prompt record-retrieval and counting check serially and at eight concurrent requests, with identical answers in both modes. Limitations: under decode stream overlap the indexer runs beside the auxiliary-stream query projection and absorption, so skipping it on a later pass saves less than its duration. With context parallelism, each rank reuses its own first-pass local selection; that path is untested. Not run: context parallelism, SM103 hardware, and acceptance or decode-time measurements that isolate this change. Design-Signed-Off-By: flare-02; design-review=4d3a --- docs/mtp-design.md | 9 ++++ ignition/recipes/glm53/mla.py | 9 +++- ignition/recipes/glm53/mtp.py | 5 +- ignition/recipes/glm53/speculative.py | 34 +++++++++++- tests/recipes/test_glm53_mla_cuda.py | 72 +++++++++++++++++++++++++ tests/recipes/test_glm53_speculative.py | 31 +++++++++++ 6 files changed, 155 insertions(+), 5 deletions(-) diff --git a/docs/mtp-design.md b/docs/mtp-design.md index 90c7c9aa..5636bc82 100644 --- a/docs/mtp-design.md +++ b/docs/mtp-design.md @@ -129,6 +129,15 @@ snapshots therefore preserve the draft partial page at that endpoint as well as target pages and recurrent state. Absorption, proposals, verification, sampling, and commit execute inside prepared CUDA graphs. +The GLM checkpoints declare `index_share_for_mtp_iteration`. A round's first +draft pass runs the MTP layer's C4 indexer. Later passes reuse each request's +first-pass sparse selection and skip the indexer's projections, C4 writes, +selection, and, with context parallelism, candidate exchange. They attend to +records at or before the first pass's position, not to the records of the +round's earlier drafts. Transformers has no GLM MTP layer and ignores the +field; vLLM's GLM MTP proposer applies the same reuse without prefill +context parallelism. + ## DeepSeek DSpark proposals and continuation state The [DSpark preparation code](../ignition/recipes/deepseek41/dspark.py) composes diff --git a/ignition/recipes/glm53/mla.py b/ignition/recipes/glm53/mla.py index 6f9b77fb..64418393 100644 --- a/ignition/recipes/glm53/mla.py +++ b/ignition/recipes/glm53/mla.py @@ -396,7 +396,11 @@ def rebind_pages(self): for slice in self.slices: slice[4].rebind_cache(self.state.cache) - def __call__(self, source): + def __call__(self, source, *, select=True): + """Attend with the indexer's selection, or without `select` with the one already in index scratch. + + A pass that reuses the selection writes its latent records but leaves the C4 cache and tail unchanged. + """ s, p = self.scratch, self.projections p["qkv_a_proj"](source) rms_norm(s.query_input, self.module.q_a_layernorm.weight, s.query_residual, @@ -411,7 +415,8 @@ def __call__(self, source): rms_norm(s.latent_input, self.module.kv_a_layernorm.weight, s.latent, self.module.kv_a_layernorm.variance_epsilon, round_before_weight=True) self.writer() - self.indexer() + if select: + self.indexer() if self.context is None: s.visible.copy_(self.metadata.positions + 1) if self.end is not None: diff --git a/ignition/recipes/glm53/mtp.py b/ignition/recipes/glm53/mtp.py index edcc7549..df5febf7 100644 --- a/ignition/recipes/glm53/mtp.py +++ b/ignition/recipes/glm53/mtp.py @@ -238,7 +238,8 @@ def buffer(*shape): def rebind_pages(self): self.attention.rebind_pages() - def __call__(self): + def __call__(self, *, select=True): + """Without `select`, attention reuses the sparse selection already in its index scratch.""" w, m = self.weights.tensors, self.metadata self.embedding() self.embedded.masked_fill_((m.indexed.positions == 0)[:, None], 0) @@ -257,7 +258,7 @@ def __call__(self): index.gate_projection(self.normalized) glm_pool.update_pools(*index.pool_arguments, plan=index.pool_plan) return - self.attention_reduce(self.attention(self.normalized)) + self.attention_reduce(self.attention(self.normalized, select=select)) rms_norm_add_residual(self.attention_output, self.residual, w["post_attention_layernorm"], self.eps, x_out=self.normalized) self.mlp_reduce(self.moe(self.normalized)) diff --git a/ignition/recipes/glm53/speculative.py b/ignition/recipes/glm53/speculative.py index 1540a0c1..7a3d18a1 100644 --- a/ignition/recipes/glm53/speculative.py +++ b/ignition/recipes/glm53/speculative.py @@ -112,6 +112,10 @@ def __init__(self, weights, pool, requests, inputs, scratch, prefill_scratch, *, self.draft_tables = torch.full_like(inputs.page_tables, -1) self.draft_generation = torch.zeros_like(requests.generation) self.draft_offset = torch.zeros_like(requests.generation) + # A checkpoint declaring index_share_for_mtp_iteration attends every draft pass of a round through + # its first pass's sparse selection. + self.share_draft_selection = bool(getattr(weights.tree.config.text_config, + "index_share_for_mtp_iteration", False)) # Verification records every KDA layer's tokens, so each request keeps one recurrent row. layers = delta_layers(pool) recurrent = pool.decoder.layers[layers[0]].recurrent @@ -195,6 +199,11 @@ def __init__(self, owner, batch, width, name): self.draft = o.draft(self.draft_metadata, f"{name}.draft") if w > 1 else None self.draft_logits = (buffers.zeros(self.draft.logits.shape, self.draft.logits.dtype) if self.draft is not None else None) + self.selection = None + if o.share_draft_selection and w > 2: + index = self.draft.attention.scratch.index + self.selection = (buffers.zeros(index.physical.shape, index.physical.dtype), + buffers.zeros(index.selected_lengths.shape, index.selected_lengths.dtype)) self.absorb_metadata = DraftMetadata(b * w, b, blocks, device, shards=o.shards) self.absorb = o.draft(self.absorb_metadata, f"{name}.absorb", cache_only=True) self.before = torch.zeros(b, dtype=torch.int64, device=device) @@ -244,6 +253,24 @@ def rebind_pages(self, pages): draft.rebind_pages() self.validator.pages = self.owner.pool.memory_capacity[1:] + def retain_selection(self): + """Keep each request's first-pass selection from its output row. + + Requests that propose a later draft proposed the first one; other requests' copies are never read. + """ + index, rows = self.draft.attention.scratch.index, self.draft_metadata.output_indices + physical, lengths = self.selection + torch.index_select(index.physical, 0, rows, out=physical) + torch.index_select(index.selected_lengths, 0, rows, out=lengths) + + def reuse_selection(self): + """Place each proposing request's first-pass selection in its packed draft row; idle rows select nothing.""" + index, m = self.draft.attention.scratch.index, self.draft_metadata.indexed + physical, lengths = self.selection + requests = m.requests.long() + torch.index_select(physical, 0, requests, out=index.physical) + index.selected_lengths.copy_(torch.where(m.positions >= 0, lengths[requests], 0)) + def target_metadata(self, slots, limits): o, b, w, m = self.owner, self.batch, self.width, self.metadata g, index = m.recurrent, m.indexed @@ -324,7 +351,12 @@ def __call__(self): proposing = self.active & (i < limits - 1) self.draft_metadata.set(slots, proposing.long(), self.before - 1 + i, token, hidden, o.draft_tables, state_rows=o.inputs.request_state_rows, capacity=1) - self.draft() + reuse = self.selection is not None and i > 0 + if reuse: + self.reuse_selection() + self.draft(select=not reuse) + if self.selection is not None and i == 0: + self.retain_selection() torch.index_select(self.draft.logits, 0, self.draft_metadata.output_indices, out=self.draft_logits) self.offset.copy_(self.base_offset + 4 * i) proposal() diff --git a/tests/recipes/test_glm53_mla_cuda.py b/tests/recipes/test_glm53_mla_cuda.py index ecad87b8..0052f913 100644 --- a/tests/recipes/test_glm53_mla_cuda.py +++ b/tests/recipes/test_glm53_mla_cuda.py @@ -14,6 +14,7 @@ from cuda.bindings import driver as cuda from ignition.components.parallel import TensorParallelShard +from ignition.components.sparse_mla import PreparedLatentCacheWriter from ignition.preparation.checkpoint import inspect_indexed_checkpoint from ignition.preparation.compilation import compile_forward from ignition.preparation.execution import GraphSignature, PreparedExecution @@ -235,3 +236,74 @@ def forward(): if worker is not None: worker.close() graph.reset() + + +@pytest.mark.parametrize("overlap", [True, False]) +@torch.no_grad() +def test_reused_selection_attends_only_to_the_first_pass_records(model_weights, overlap): + """A draft pass that shares the first pass's selection excludes its own and later records.""" + module, weights, world = model_weights + if world == 1: + pytest.skip("GLM prepares draft graphs for TP2 and TP4") + device = torch.device("cuda:0") + source = torch.zeros((2, 4096), dtype=torch.bfloat16, device=device) + scratch = MLAScratch.allocate(2, 64 // world, 64, "decode", device) + state = MLAState.allocate(128, 4, device) + metadata = IndexMetadata.allocate(2, 2, 64, device) + history = torch.zeros((256, 512), dtype=torch.bfloat16, device=device) + history_slots = torch.zeros(256, dtype=torch.int64, device=device) + arena = ScratchArena(device, overlap=overlap) + with RecordedPreparationSession(device=device, checkpoint="GLM-5.3-Flash-NVFP4-MXFP8-CSF-QAD", + workload_schema="glm-mla-shared-selection-v1", numerical_policy="default") as session: + session.register_scratch(arena) + attention = PreparedGlmMLA(module, weights, source, scratch, state, metadata, + mode="decode", session=session, name="mla", workspace=arena) + writer = PreparedLatentCacheWriter(history, state.cache, history_slots, session=session, name="history") + session.tune() + state.cache.zero_() + state.index.cache.zero_() + state.index.tail.zero_() + generator = torch.Generator(device=device).manual_seed(173) + metadata.block_table.copy_(torch.randperm(128, device=device, generator=generator, dtype=torch.int32) + .reshape(2, 64)) + # Request zero is live in state slot 2; request one is idle. + metadata.state_slots.copy_(torch.tensor([2, -1], device=device)) + metadata.query_start_loc.copy_(torch.tensor([0, 1, 1], dtype=torch.int32, device=device)) + metadata.requests.zero_() + page = metadata.block_table[0, 0].long() + # Latent records of positions 0 through 255 precede the first draft pass at position 256. With fewer + # than 512 complete C4 pools every pool is selected, so the selection does not depend on the query. + history.normal_(generator=generator) + history_slots.copy_(page * 256 + torch.arange(256, device=device)) + writer() + + def run(position, *, select, write=True): + slot = metadata.block_table[0, position // 256].long() * 256 + position % 256 + metadata.positions.copy_(torch.tensor([position, -1], device=device)) + metadata.token_slots.copy_(torch.where(torch.tensor([write, False], device=device), slot, -1)) + return attention(source, select=select).clone() + + source[0].normal_(std=.1, generator=generator) + run(256, select=True) + physical, lengths = scratch.index.physical.clone(), scratch.index.selected_lengths.clone() + # Sixty-four complete pools and the tail token at position 256. + assert lengths.tolist() == [257, 0] + index_cache, index_tail = state.index.cache.clone(), state.index.tail.clone() + source[0].normal_(std=.1, generator=generator) + reused = run(257, select=False) + # The pass writes its latent record and leaves the selection and C4 state unchanged. + assert torch.equal(scratch.index.physical, physical) and torch.equal(scratch.index.selected_lengths, lengths) + assert torch.equal(state.index.cache, index_cache) and torch.equal(state.index.tail, index_tail) + own = (metadata.block_table[0, 1].long() * 256 + 1).item() + assert state.cache.view(-1, 528)[own].count_nonzero() + assert not reused[1].count_nonzero() + # Selecting at the preceding position with the same query attends to the same records. + preceding = run(256, select=True, write=False) + assert torch.equal(scratch.index.physical[0, :257].sort().values, physical[0, :257].sort().values) + assert ((preceding.float() - reused.float()).norm() / reused.float().norm()).item() < 1e-3 + # Selecting at the pass's own position also attends to its record. + state.index.cache.copy_(index_cache) + state.index.tail.copy_(index_tail) + current = run(257, select=True, write=False) + assert scratch.index.selected_lengths.tolist() == [258, 0] + assert ((current.float() - reused.float()).norm() / reused.float().norm()).item() > 2e-2 diff --git a/tests/recipes/test_glm53_speculative.py b/tests/recipes/test_glm53_speculative.py index 58ca0c38..5d45e0fc 100644 --- a/tests/recipes/test_glm53_speculative.py +++ b/tests/recipes/test_glm53_speculative.py @@ -96,6 +96,37 @@ def test_draft_rows_exclude_exhausted_requests_and_restore_output_order(): assert metadata.indexed.token_slots.tolist() == [-1, -1, -1] +def test_later_draft_passes_reuse_each_requests_first_selection(): + op = PreparedRound.__new__(PreparedRound) + op.draft_metadata = DraftMetadata(3, 3, 2, "cpu") + index = SimpleNamespace(physical=torch.zeros((3, 2051), dtype=torch.int32), + selected_lengths=torch.zeros(3, dtype=torch.int32)) + op.draft = SimpleNamespace(attention=SimpleNamespace(scratch=SimpleNamespace(index=index))) + op.selection = (torch.zeros((3, 2051), dtype=torch.int32), torch.zeros(3, dtype=torch.int32)) + slots, starts = torch.tensor([2, 0, 1]), torch.tensor([255, 510, 300]) + state_rows = torch.arange(1, 4)[:, None].expand(3, 8) + tables = torch.tensor([[10, 11], [20, 21], [30, 31]]) + hidden = torch.zeros((3, 4096), dtype=torch.bfloat16) + tokens = torch.tensor([7, 8, 9]) + # The first pass packs requests one and two; request zero has no draft to propose. + op.draft_metadata.set(slots, torch.tensor([0, 1, 1]), starts, tokens, hidden, tables, + capacity=1, state_rows=state_rows) + first = {1: torch.tensor([3, 5, 9]), 2: torch.tensor([4, 6, 8, 12])} + index.physical.fill_(-1) + for row, request in enumerate((1, 2)): + index.physical[row, :first[request].numel()] = first[request] + index.selected_lengths[row] = first[request].numel() + op.retain_selection() + # The second pass keeps only request two, which now occupies packed row zero. + op.draft_metadata.set(slots, torch.tensor([0, 0, 1]), starts + 1, tokens, hidden, tables, + capacity=1, state_rows=state_rows) + index.physical.fill_(-2) + index.selected_lengths.fill_(-2) + op.reuse_selection() + assert index.selected_lengths.tolist() == [4, 0, 0] + assert index.physical[0, :4].tolist() == first[2].tolist() + + @pytest.fixture(autouse=True) def records_only_cpu_commit(monkeypatch): # These state-machine tests run without CUDA; reduction and graph ordering