From 3993b94c58d114e11deb18aa3a28084acf4271b2 Mon Sep 17 00:00:00 2001 From: Moustafa Saleh Date: Thu, 9 Jul 2026 02:19:33 +0100 Subject: [PATCH] fix(trtllm-mla): make spec-decode cuda graph capture causal Signed-off-by: Moustafa Saleh --- .../runtime/execution/cuda_graph_wrapper.py | 33 ++++++++++++++++++- .../layers/attention/backends/trtllm_mla.py | 22 ++++++++++--- 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/python/tokenspeed/runtime/execution/cuda_graph_wrapper.py b/python/tokenspeed/runtime/execution/cuda_graph_wrapper.py index 95f40dbbc..aaa33cbb9 100644 --- a/python/tokenspeed/runtime/execution/cuda_graph_wrapper.py +++ b/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() diff --git a/python/tokenspeed/runtime/layers/attention/backends/trtllm_mla.py b/python/tokenspeed/runtime/layers/attention/backends/trtllm_mla.py index 65e21af59..9cbef9a5e 100644 --- a/python/tokenspeed/runtime/layers/attention/backends/trtllm_mla.py +++ b/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,9 +424,11 @@ 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 @@ -430,11 +436,17 @@ def forward_decode( 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)