[None][perf] Overlap DSA heuristic prev_topk write-back on the aux stream

[longcheng-nv]

Description

Indexer.sparse_attn_indexer writes each decode step's top-k back to metadata.heuristic_prev_topk (the next step's pre_idx hint). This copy is a strided gather on the main stream, once per indexer layer per decode step, and nothing in the current step reads it. This PR forks the copy onto the indexer's existing aux stream right after the top-k kernel and joins it in the same layer's MLA forward after core sparse attention is enqueued, so the copy overlaps with the layer's heaviest decode work instead of delaying it.

Notes:

• Fork and join stay within one layer's forward: CUDA graph capture rejects unjoined forks, and the join restores ordering before the next layer overwrites the topk_indices_buffer rows the copy reads.
• Source and destination are persistent graph-pool buffers, so the captured copy is valid on every replay; no record_stream needed.
• Gated on do_multi_stream(), same policy as maybe_execute_in_parallel. Eager execution keeps the original inline copy.
• Covers both consumers: forward_dsa_attn (V3.2) and forward_impl_with_deepseek_v4. Shared DSA indexer layers (indexer is None) are unchanged.
• Also benefits the GVR top-k e2e wiring ([None][feat] top-k: route decode to CuTe DSL GVR top-k in e2e #16420), which consumes the same feedback loop.

Test Coverage

• Standalone CUDA graph smoke test: capture with the same-layer join succeeds and replayed feedback values are step-correct across replays; capture without the join fails with cudaErrorStreamCaptureUnjoined.
• Existing e2e: test_fp8_blockscale[heuristic_topk_mtp1] (TestDeepSeekV32, cuda_graph=True) exercises the forked path on Blackwell.

Performance

nsys per-iteration A/B: DeepSeek-V4 Pro FP4, 8xB200, TEP8, BS=1, ISL 64K, OSL 2048, MTP=0, enable_heuristic_topk: true, CUDA graphs ON. Arms differ only in this PR's files on the same C++ build; ABBA run order; nsys window = decode iters 500-550 (49 per-iteration deltas per run).

per-iter decode	base	this PR	delta
10% trimmed mean	10.6501 ms	10.6351 ms	-15.0 µs/iter
median	10.6488 ms	10.6335 ms	-15.3 µs
min	10.600 ms	10.575 ms	-25 µs

Mann-Whitney U z = -4.60 (p ≈ 4e-6); all four runs direction-consistent. In the base traces 0/300 feedback copies overlap other-stream work; with this PR 300/300 overlap core sparse attention and the join adds no main-stream bubble.

-15 µs/iter over 30 indexer layers ≈ 0.5 µs/layer, i.e. the copy-kernel duration leaving the critical path. That is 0.14-0.23% of TPOT for this config and ≤ ~0.4% for any realistic config, below what e2e timing can resolve. A matching e2e TPOT A/B (DeepSeek-V3.2-Exp FP4, 8xB200, TEP8, BS=1, MTP=3, ISL ~68.7K, 10 warm pairs) reads -0.42% mean / -0.16% median, SE ≈ 0.7% — statistically zero, consistent with the per-iteration bound.

Note (2026-08-12): an earlier revision reported +2-4% e2e TPOT gains (MTP=0/1/2). Those numbers are retracted as measurement artifacts — their baselines are inconsistent with later same-node reruns, and the magnitude is 10-20x above the per-iteration bound above.

PR Checklist

• Please check this after reviewing the above items as appropriate for this PR.

Dev Engineer Review

• Added auxiliary-stream fork/join handling for heuristic TopK copies in DSA and DeepSeek-V4 paths.
• Preserved synchronous copying when do_multi_stream() is disabled.
• Added Indexer.maybe_join_prev_topk_copy() with pending-state cleanup.
• Placed joins after core attention enqueue to preserve stream ordering.
• Preserved CUDA graph capture and persistent-buffer usage.
• Left shared DSA indexer layers unchanged.
• No configuration or test-list changes were identified.
• Reported a 15 µs per-iteration nsys improvement. End-to-end TPOT results were statistically null and consistent with the estimated 0.14–0.23% effect. Earlier 2–4% results were retracted as measurement artifacts.
• Fix #17416 initializes in_mtp_draft_loop in the synthetic metadata stub.
• Later CI run #65523 succeeded after the rebase and inclusion of fix #17416.

QA Engineer Review

No test changes.

[longcheng-nv]

/bot run

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c58d4c84-b5fe-41fc-9839-4269058fcaab

📥 Commits

Reviewing files that changed from the base of the PR and between 3253b64 and 842357f.

📒 Files selected for processing (3)

• tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py
• tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
• tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py

🚧 Files skipped from review as they are similar to previous changes (3)

• tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py
• tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py
• tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py

Included review availability: Your plan includes up to 12 reviews per rolling hour; 10 remain after this review.

---

Walkthrough

The change adds asynchronous heuristic TopK write-back coordination to Indexer. DSA and DeepSeek-V4 sparse attention paths join pending copies after scheduling core attention work.

Changes

DSA TopK write-back

Layer / File(s)	Summary
Asynchronous TopK copy coordination tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py	Indexer adds CUDA events and pending state. Decode write-back uses aux_stream when multi-stream execution is enabled and retains synchronous copying otherwise. A join method clears pending state after synchronization.
Sparse attention forward-path synchronization tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py, tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py	Both sparse attention paths call maybe_join_prev_topk_copy() after core attention scheduling when an indexer is present.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 84235

This change moves a per-layer feedback copy onto an auxiliary stream only when multi-stream execution is enabled, while preserving the eager path and matching each fork with a same-layer join; no actionable merge-blocking risk remains beyond normal checks and review.

Suggested reviewers: lfr-0531, kefeng-duan, yunruis

Sequence Diagram(s)

sequenceDiagram
  participant SparseAttention
  participant Indexer
  participant aux_stream
  participant CUDA_events
  SparseAttention->>Indexer: Schedule heuristic TopK write-back
  Indexer->>aux_stream: Fork asynchronous copy
  aux_stream->>CUDA_events: Record copy completion
  SparseAttention->>Indexer: maybe_join_prev_topk_copy()
  Indexer->>CUDA_events: Wait for pending copy
  CUDA_events-->>Indexer: Signal completion
  Indexer-->>SparseAttention: Clear pending state

Loading

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Linked Issues check	⚠️ Warning	The linked issue requires initializing in_mtp_draft_loop in a synthetic test stub, but the provided changes do not show that fix.	Add the synthetic metadata stub initialization for in_mtp_draft_loop = False and include the affected test change.
Out of Scope Changes check	⚠️ Warning	The asynchronous heuristic TopK copy changes are unrelated to the sole linked issue, which only requires a test-stub initialization fix.	Link the performance issue that scopes the auxiliary-stream TopK changes, or separate those changes from the linked-issue fix.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Description check	✅ Passed	The description explains the problem, solution, test coverage, performance results, and checklist items clearly.
Title check	✅ Passed	The title clearly and concisely describes moving DSA heuristic TopK write-back to the auxiliary stream.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[tensorrt-cicd]

PR_Github #64192 [ run ] triggered by Bot. Commit: 90c4b5c Link to invocation

[tensorrt-cicd]

PR_Github #64192 [ run ] completed with state FAILURE. Commit: 90c4b5c
/LLM/main/L0_MergeRequest_PR pipeline #52105 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[longcheng-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #64243 [ run ] triggered by Bot. Commit: 90c4b5c Link to invocation

[tensorrt-cicd]

PR_Github #64243 [ run ] completed with state FAILURE. Commit: 90c4b5c
/LLM/main/L0_MergeRequest_PR pipeline #52149 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[longcheng-nv]

/bot run --disable-fail-fast

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[tensorrt-cicd]

PR_Github #64444 [ run ] triggered by Bot. Commit: 4d90b7b Link to invocation

[tensorrt-cicd]

PR_Github #64444 [ run ] completed with state FAILURE. Commit: 4d90b7b
/LLM/main/L0_MergeRequest_PR pipeline #52320 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[longcheng-nv]

Rebased/ported onto the sparse-attention framework refactor (#12733), new head 52d3f9b: the Indexer fork moved from sparse/dsa.py to sparse/dsa/indexer.py (single site now covers both V3.2 and the DeepseekV4Indexer subclass), and the two same-layer join sites moved from modules/mla.py to sparse/dsa/module.py::_forward_dsa_attn and sparse/deepseek_v4/module.py::forward_sparse_attn. Logic is unchanged (diff-of-diffs is the import split only). This also picks up #17416, which fixes the test_on_update_kv_lens_rebuilds_stale_map stub failure that hit the previous CI run on main-side code.

[longcheng-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #64772 [ run ] triggered by Bot. Commit: 52d3f9b Link to invocation

[tensorrt-cicd]

PR_Github #64772 [ run ] completed with state FAILURE. Commit: 52d3f9b
/LLM/main/L0_MergeRequest_PR pipeline #52616 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[longcheng-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65523 [ run ] triggered by Bot. Commit: 52d3f9b Link to invocation

[tensorrt-cicd]

PR_Github #65523 [ run ] completed with state SUCCESS. Commit: 52d3f9b
/LLM/main/L0_MergeRequest_PR pipeline #53259 completed with status: 'SUCCESS'

CI Report

Link to invocation

[longcheng-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #66682 [ run ] triggered by Bot. Commit: 366e531 Link to invocation

[tensorrt-cicd]

PR_Github #66682 [ run ] completed with state SUCCESS. Commit: 366e531
/LLM/main/L0_MergeRequest_PR pipeline #54298 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[longcheng-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #67019 [ run ] triggered by Bot. Commit: 842357f Link to invocation

[tensorrt-cicd]

PR_Github #67019 [ run ] completed with state SUCCESS. Commit: 842357f
/LLM/main/L0_MergeRequest_PR pipeline #54565 completed with status: 'SUCCESS'

CI Report

Link to invocation

[longcheng-nv]

/bot reuse-pipeline

[tensorrt-cicd]

PR_Github #67623 [ reuse-pipeline ] triggered by Bot. Commit: 83f6c83 Link to invocation

[tensorrt-cicd]

PR_Github #67623 [ reuse-pipeline ] completed with state SUCCESS. Commit: 83f6c83
Reusing PR_Github #67019 for commit 83f6c83

Link to invocation