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

python/tokenspeed/runtime/execution/cuda_graph_wrapper.py

@@ -349,6 +349,41 @@ def _capture_one(self, bs: int):
             grammar_backend=self.grammar_backend,
         )
 
+        # Spec-decode capture runs a synthetic multi-token decode. Keep the
+        # dummy cache lengths internally consistent with that token count so
+        # attention warmup does not read an impossible q_len > seq_len state.
+        tokens_per_req = self.max_tokens_per_req
+        self.input_buffers.seq_lens_buf[:bs].fill_(tokens_per_req)
+
+        # Capture block tables point at synthetic per-request pages. Write the
+        # dummy KV tokens into those same slots so attention warmup/capture reads
+        # initialized keys instead of the reserved padding slot.
+        page_size = self.input_buffers.page_size
+        pages_per_req = (tokens_per_req + page_size - 1) // page_size
+        token_offsets = torch.arange(tokens_per_req, dtype=torch.int32, device=self.device)
+        request_offsets = (
+            torch.arange(bs, dtype=torch.int32, device=self.device).unsqueeze(1)
+            * pages_per_req
+            * page_size
+        )
+        self.input_buffers.out_cache_loc_buf[: bs * tokens_per_req].copy_(
+            (request_offsets + token_offsets).reshape(-1)
+        )
+
+        # Some fused decode kernels may read full page vectors during capture
+        # even when seq_lens bounds the logical context. Clear the synthetic
+        # pages so any padding read is deterministic zero, not allocator noise.
+        capture_slots = bs * pages_per_req * 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[:capture_slots].zero_()
+                else:
+                    layer_buf[:capture_slots].zero_()
+
         self._init_capture_metadata(bs)
 
         def run_once():

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

@@ -341,7 +341,18 @@ 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. Point each synthetic request
+        # at private dummy pages that match CudaGraphWrapper's out_cache_loc_buf
+        # setup, so warmup/capture attention reads initialized KV entries.
+        block_kv_indices.zero_()
+        tokens_per_req = max(int(self.spec_num_tokens or 1), 1)
+        pages_per_req = triton.cdiv(tokens_per_req, self.page_size)
+        fill_pages = min(pages_per_req, block_kv_indices.shape[1])
+        if fill_pages > 0:
+            page_offsets = torch.arange(fill_pages, dtype=torch.int32, device=self.device)
+            request_offsets = torch.arange(bs, dtype=torch.int32, device=self.device).unsqueeze(1) * pages_per_req
+            block_kv_indices[:, :fill_pages] = request_offsets + page_offsets
+
         # The actual indices will be filled on replay. seq_lens_k aliases
         # seq_lens_buf (set in init_cuda_graph_state).
         metadata = TRTLLMMLADecodeMetadata(
@@ -414,21 +425,28 @@ 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 bounds:
+                # [valid+1, valid+2, ..., valid+q_len].
+                base_lens = base_lens - (q_len_per_req - 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)