Share the first GLM draft pass's sparse selection with later passes

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

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:

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))

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()

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

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