fix(trtllm-mla): make spec-decode CUDA graph capture causal

python/tokenspeed/runtime/execution/cuda_graph_wrapper.py

@@ -443,6 +443,35 @@ def _capture_one(self, bs: int, variant: str = CUDA_GRAPH_VARIANT_DEFAULT):
             grammar_backend=self.grammar_backend,
         )
 
+        def prepare_attention_capture_buffers():
+            # Spec-decode capture runs a synthetic multi-token decode. Keep the
+            # dummy cache state internally consistent so attention capture does
+            # not depend on impossible q_len > seq_len inputs or private pages.
+            tokens_per_req = self.max_tokens_per_req
+            self.input_buffers.seq_lens_buf[:bs].fill_(tokens_per_req)
+
+            page_size = self.input_buffers.page_size
+            dummy_slot = int(self.input_buffers.dummy_kv_slot)
+            dummy_page = dummy_slot // page_size
+            for backend in (self.attn_backend, self.draft_attn_backend):
+                if backend is not None:
+                    backend.cuda_graph_capture_dummy_page = dummy_page
+            self.input_buffers.out_cache_loc_buf[: bs * tokens_per_req].fill_(
+                dummy_slot
+            )
+
+            dummy_page_start = dummy_page * page_size
+            dummy_page_end = dummy_page_start + page_size
+            for pool in (self.token_to_kv_pool, self.draft_token_to_kv_pool):
+                if pool is None or not hasattr(pool, "kv_buffer"):
+                    continue
+                for layer_buf in pool.kv_buffer:
+                    if isinstance(layer_buf, (tuple, list)):
+                        for sub_buf in layer_buf:
+                            sub_buf[dummy_page_start:dummy_page_end].zero_()
+                    else:
+                        layer_buf[dummy_page_start:dummy_page_end].zero_()
+
         def run_once():
             # Dummy add_batch keeps the grammar queue 1:1 with replays —
             # fetch_batch pops once per forward, so warmup + capture
@@ -460,7 +489,7 @@ def run_once():
             self._prepare_sampling_capture(bs=bs, variant=variant)
             # Keep warmup seq_lens >= q_len_per_req so no query row gets an
             # empty causal span; a stale seq_len of 1 overflows to non-finite KV.
-            self.input_buffers.seq_lens_buf[:bs].fill_(self.max_tokens_per_req)
+            prepare_attention_capture_buffers()
             self._init_capture_metadata(bs)
             run_once()
 
@@ -475,12 +504,14 @@ def run_once():
         # Warmups can switch a backend back to eager metadata objects. Restore
         # the graph-backed metadata immediately before capture so replay-time
         # metadata refreshes update the same tensors recorded by the graph.
+        prepare_attention_capture_buffers()
         self._init_capture_metadata(bs)
 
         # Fill sampler buffers OUTSIDE the capture so RNG ops aren't recorded.
         self._prepare_sampling_capture(bs=bs, variant=variant)
         # Warmup forwards can mutate aliased metadata buffers, so refresh
         # them again immediately before graph capture records the final views.
+        prepare_attention_capture_buffers()
         self._init_capture_metadata(bs)
 
         self.deepep_adapter.capture()

python/tokenspeed/runtime/layers/attention/backends/trtllm_mla.py

@@ -344,7 +344,11 @@ def init_forward_metadata_capture_cuda_graph(
         max_blocks = self._calc_padded_blocks(self.max_context_len)
         block_kv_indices = self.decode_cuda_graph_kv_indices[:bs, :max_blocks]
 
-        # For capture we don't have req_to_page yet; just zero-fill the block indices.
+        # For capture we don't have req_to_page yet. Use the reserved padding
+        # page for all synthetic rows so capture does not depend on extra KV
+        # capacity for private per-request pages.
+        block_kv_indices.fill_(int(getattr(self, "cuda_graph_capture_dummy_page", 0)))
+
         # The actual indices will be filled on replay. seq_lens_k aliases
         # seq_lens_buf (set in init_cuda_graph_state).
         metadata = TRTLLMMLADecodeMetadata(
@@ -420,21 +424,29 @@ def forward_decode(
         num_extends = metadata.num_extends
         q_len_per_req = q.shape[0] // bs if bs > 0 else 1
 
-        if q_len_per_req > 1 and self.is_draft:
-            # First draft step catching up its KV after verify: one query entry per token;
-            # per-token seq_lens advance by 1 so each successive token sees its own KV write.
+        if q_len_per_req > 1:
+            # Multi-token decode is used by target verification and by the
+            # draft first-step catch-up path. Flatten to one decode query per
+            # token so each token gets its own causal seq_len instead of using
+            # the grouped decode path with one full-context bound for all rows.
             query = q.view(-1, layer.tp_q_head_num, layer.head_dim).unsqueeze(1)
             block_tables = metadata.block_kv_indices[num_extends:].repeat_interleave(
                 q_len_per_req, dim=0
             )
             base_lens = metadata.seq_lens_k[num_extends:].repeat_interleave(
                 q_len_per_req
             )
+            if not self.is_draft:
+                # Target verification receives seq_lens at the end of the
+                # speculative window. Convert to per-token causal bounds.
+                base_lens = torch.clamp(base_lens - (q_len_per_req - 1), min=1)
             offsets = torch.arange(
                 q_len_per_req, device=base_lens.device, dtype=base_lens.dtype
             ).repeat(bs)
             seq_lens = base_lens + offsets
-            max_seq_len = metadata.max_seq_len_k + q_len_per_req
+            max_seq_len = metadata.max_seq_len_k + (
+                q_len_per_req if self.is_draft else 0
+            )
         else:
             # Plain decode (q_len=1) or bs-grouped multi-token decode.
             query = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim)