From b231986050c811acfd12687afaeb87665bd767b7 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Sat, 8 Aug 2026 07:36:47 +0000 Subject: [PATCH 1/3] [None][perf] Overlap DSA heuristic prev_topk write-back on the aux stream The per-layer heuristic top-k feedback copy (this step's decode top-k -> next step's pre_idx hint) is a strided gather sitting on the main stream's critical path, once per indexer layer per decode step. Nothing in the current step consumes it, so fork it onto the Indexer's existing aux stream right after the top-k kernel and join it in the same layer's MLA forward once core sparse attention is enqueued -- the copy overlaps with the layer's heaviest decode work. Same-layer fork/join keeps CUDA graph capture free of unjoined forks (cudaStreamEndCapture rejects them) and restores ordering before the next layer overwrites the shared topk_indices_buffer rows the copy reads. Source and destination are persistent stable-address buffers, so replays stay valid with no record_stream bookkeeping. The fork engages only under do_multi_stream() (i.e. inside CUDA graph capture, where replay makes the stream/event host overhead free); eager execution keeps the original inline copy unchanged. Validated with a pattern-level CUDA graph smoke test: capture with the join succeeds, replayed feedback values are step-correct, and capture without the join fails with cudaErrorStreamCaptureUnjoined. Ported onto the sparse-attention framework refactor (#12733): the Indexer changes moved from sparse/dsa.py to sparse/dsa/indexer.py, and the two MLA join sites moved from modules/mla.py to sparse/dsa/module.py (_forward_dsa_attn) and sparse/deepseek_v4/module.py (forward_sparse_attn). Made-with: Claude Code (Fable 5) Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../sparse/deepseek_v4/module.py | 7 +++ .../attention_backend/sparse/dsa/indexer.py | 49 ++++++++++++++++++- .../attention_backend/sparse/dsa/module.py | 7 +++ 3 files changed, 61 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py index 68bb1ee8d1fb..7306618dbd39 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py @@ -1219,6 +1219,13 @@ def _indexer_branch(): self._fused_kv_norm_active = False self._fused_kv_norm_hoisted = False + # Join the aux-stream heuristic prev_topk write-back forked in + # sparse_attn_indexer, now that this layer's core attention is + # enqueued (the copy overlaps with it). Must stay within this + # layer's forward: CUDA graph capture rejects unjoined forks. + if self.indexer is not None: + self.indexer.maybe_join_prev_topk_copy() + class DeepSeekV4Hooks(MLASparseHooks): """Typed DeepSeek-V4 adapter for the shared MLA module.""" diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index eae55ac98ae7..ef3236911075 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -21,7 +21,10 @@ from tensorrt_llm._torch.distributed.ops import allgather from tensorrt_llm._torch.modules.layer_norm import LayerNorm from tensorrt_llm._torch.modules.linear import Linear -from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.modules.multi_stream_utils import ( + do_multi_stream, + maybe_execute_in_parallel, +) from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding from tensorrt_llm._torch.utils import Fp4QuantizedTensor, maybe_compile from tensorrt_llm._utils import get_sm_version, maybe_pin_memory, prefer_pinned @@ -682,6 +685,11 @@ def __init__( self.use_fp4 = sparse_params.indexer_k_dtype == "fp4" self.aux_stream = aux_stream self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] + # Fork/join pair for the aux-stream heuristic prev_topk write-back: + # [0] orders the copy after the top-k kernel, [1] is waited on by + # maybe_join_prev_topk_copy() in the owning MLA layer. + self.prev_topk_copy_events = [torch.cuda.Event(), torch.cuda.Event()] + self._prev_topk_copy_pending = False self.use_cute_dsl_topk = sparse_params.use_cute_dsl_topk and IS_CUTLASS_DSL_AVAILABLE self.use_cute_dsl_paged_mqa_logits = ( sparse_params.use_cute_dsl_paged_mqa_logits and IS_CUTLASS_DSL_AVAILABLE @@ -1879,7 +1887,29 @@ def sparse_attn_indexer( local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx] decode_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens] last_mtp_topk = decode_topk[next_n - 1 :: next_n] - metadata.heuristic_prev_topk[local_layer, :num_generations].copy_(last_mtp_topk) + prev_topk_dst = metadata.heuristic_prev_topk[local_layer, :num_generations] + if do_multi_stream() and self.aux_stream is not None: + # Fork the write-back onto the aux stream so the strided + # gather copy overlaps with this layer's core sparse + # attention instead of sitting on the critical path. + # Nothing in this step reads it back — the next consumer + # is the next decode step's pre_idx for this same layer. + # Source and destination are persistent stable-address + # buffers, so no record_stream bookkeeping is needed. The + # fork MUST be joined within this layer's forward via + # maybe_join_prev_topk_copy(): that both keeps CUDA graph + # capture free of unjoined forks (cudaStreamEndCapture + # rejects them) and restores ordering before the next + # layer overwrites the shared topk_indices_buffer rows + # this copy reads. + self.prev_topk_copy_events[0].record() + with torch.cuda.stream(self.aux_stream): + self.prev_topk_copy_events[0].wait() + prev_topk_dst.copy_(last_mtp_topk) + self.prev_topk_copy_events[1].record() + self._prev_topk_copy_pending = True + else: + prev_topk_dst.copy_(last_mtp_topk) elif has_decode and metadata.skip_indexer_for_gen_reqs: # Fill topk_indices_buffer with pre-defined dense topk indices @@ -1929,6 +1959,21 @@ def _mtp_last_accepted_rows( offset = (gen_num_accepted - 1).clamp(0, next_n - 1) return gen_topk[base + offset] + def maybe_join_prev_topk_copy(self) -> None: + """Join the aux-stream heuristic prev_topk write-back, if forked. + + Called by the owning MLA layer after this layer's core sparse + attention has been enqueued, so the copy forked in + sparse_attn_indexer overlaps with it. Joining within the same + layer's forward keeps every fork matched with a join inside a + single captured region (CUDA graph capture rejects unjoined + forks) and orders the copy's read of topk_indices_buffer before + the next layer's indexer overwrites those rows. + """ + if self._prev_topk_copy_pending: + self.prev_topk_copy_events[1].wait() + self._prev_topk_copy_pending = False + def _weight_scale(self, weights: torch.Tensor, q_scale: torch.Tensor) -> torch.Tensor: """Apply quantization scale to indexer attention weights.""" weights = _scale(weights, q_scale, self.weight_scale_factor) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py index 5202067f14a6..bacbb9e28f69 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py @@ -387,6 +387,13 @@ def _forward_dsa_attn( indexer_intermediates=indexer_intermediates, ) + # Join the aux-stream heuristic prev_topk write-back forked in + # sparse_attn_indexer, now that this layer's core attention is + # enqueued (the copy overlaps with it). Must stay within this + # layer's forward: CUDA graph capture rejects unjoined forks. + if self.mqa.indexer is not None: + self.mqa.indexer.maybe_join_prev_topk_copy() + def should_use_short_mha( self, attn_metadata: AttentionMetadata, position_ids: Optional[torch.Tensor] From 842357f70f23c8fe297367a8e70f656fdd8b942b Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Mon, 17 Aug 2026 06:48:41 +0000 Subject: [PATCH 2/3] [None][chore] Simplify aux-stream fork/join comments per review Trim the fork-site comment and maybe_join_prev_topk_copy docstring to their essentials, as requested in review. Made-with: Claude Code (Fable 5) Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 26 +++---------------- 1 file changed, 3 insertions(+), 23 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index ef3236911075..d858b8a6ca70 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -1889,19 +1889,8 @@ def sparse_attn_indexer( last_mtp_topk = decode_topk[next_n - 1 :: next_n] prev_topk_dst = metadata.heuristic_prev_topk[local_layer, :num_generations] if do_multi_stream() and self.aux_stream is not None: - # Fork the write-back onto the aux stream so the strided - # gather copy overlaps with this layer's core sparse - # attention instead of sitting on the critical path. - # Nothing in this step reads it back — the next consumer - # is the next decode step's pre_idx for this same layer. - # Source and destination are persistent stable-address - # buffers, so no record_stream bookkeeping is needed. The - # fork MUST be joined within this layer's forward via - # maybe_join_prev_topk_copy(): that both keeps CUDA graph - # capture free of unjoined forks (cudaStreamEndCapture - # rejects them) and restores ordering before the next - # layer overwrites the shared topk_indices_buffer rows - # this copy reads. + # Fork the write-back onto the aux stream to overlap with this + # layer's core sparse attention; joined by maybe_join_prev_topk_copy(). self.prev_topk_copy_events[0].record() with torch.cuda.stream(self.aux_stream): self.prev_topk_copy_events[0].wait() @@ -1960,16 +1949,7 @@ def _mtp_last_accepted_rows( return gen_topk[base + offset] def maybe_join_prev_topk_copy(self) -> None: - """Join the aux-stream heuristic prev_topk write-back, if forked. - - Called by the owning MLA layer after this layer's core sparse - attention has been enqueued, so the copy forked in - sparse_attn_indexer overlaps with it. Joining within the same - layer's forward keeps every fork matched with a join inside a - single captured region (CUDA graph capture rejects unjoined - forks) and orders the copy's read of topk_indices_buffer before - the next layer's indexer overwrites those rows. - """ + """Join the aux-stream heuristic prev_topk write-back, if forked.""" if self._prev_topk_copy_pending: self.prev_topk_copy_events[1].wait() self._prev_topk_copy_pending = False From 83f6c83209781509e9ffd46614c4371700d26eb9 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Thu, 20 Aug 2026 02:34:59 +0000 Subject: [PATCH 3/3] [None][chore] Trim aux-stream fork/join comments to essentials Made-with: Claude Code (Fable 5) Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../_torch/attention_backend/sparse/deepseek_v4/module.py | 6 ++---- tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py | 4 +--- tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py | 6 ++---- 3 files changed, 5 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py index 7306618dbd39..2c09926c9835 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py @@ -1219,10 +1219,8 @@ def _indexer_branch(): self._fused_kv_norm_active = False self._fused_kv_norm_hoisted = False - # Join the aux-stream heuristic prev_topk write-back forked in - # sparse_attn_indexer, now that this layer's core attention is - # enqueued (the copy overlaps with it). Must stay within this - # layer's forward: CUDA graph capture rejects unjoined forks. + # Join the prev_topk copy forked in sparse_attn_indexer; CUDA graph + # capture requires the join within the same layer's forward. if self.indexer is not None: self.indexer.maybe_join_prev_topk_copy() diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index d858b8a6ca70..3e67594b1690 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -685,9 +685,7 @@ def __init__( self.use_fp4 = sparse_params.indexer_k_dtype == "fp4" self.aux_stream = aux_stream self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] - # Fork/join pair for the aux-stream heuristic prev_topk write-back: - # [0] orders the copy after the top-k kernel, [1] is waited on by - # maybe_join_prev_topk_copy() in the owning MLA layer. + # Fork/join events for the aux-stream prev_topk write-back. self.prev_topk_copy_events = [torch.cuda.Event(), torch.cuda.Event()] self._prev_topk_copy_pending = False self.use_cute_dsl_topk = sparse_params.use_cute_dsl_topk and IS_CUTLASS_DSL_AVAILABLE diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py index bacbb9e28f69..bf36a14655ed 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py @@ -387,10 +387,8 @@ def _forward_dsa_attn( indexer_intermediates=indexer_intermediates, ) - # Join the aux-stream heuristic prev_topk write-back forked in - # sparse_attn_indexer, now that this layer's core attention is - # enqueued (the copy overlaps with it). Must stay within this - # layer's forward: CUDA graph capture rejects unjoined forks. + # Join the prev_topk copy forked in sparse_attn_indexer; CUDA graph + # capture requires the join within the same layer's forward. if self.mqa.indexer is not None: self.mqa.indexer.maybe_join_prev_topk_copy()