[None][perf] Overlap DSA heuristic prev_topk write-back on the aux stream - #16666
[None][perf] Overlap DSA heuristic prev_topk write-back on the aux stream#16666longcheng-nv wants to merge 1 commit into
Conversation
767b036 to
90c4b5c
Compare
|
/bot run |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
WalkthroughThe change adds asynchronous heuristic TopK write-back coordination to ChangesDSA TopK write-back
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 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
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
PR_Github #64192 [ run ] triggered by Bot. Commit: |
|
PR_Github #64192 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64243 [ run ] triggered by Bot. Commit: |
|
PR_Github #64243 [ run ] completed with state
|
90c4b5c to
4d90b7b
Compare
|
/bot run --disable-fail-fast |
|
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. |
|
PR_Github #64444 [ run ] triggered by Bot. Commit: |
|
PR_Github #64444 [ run ] completed with state
|
…ream 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 (NVIDIA#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 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
4d90b7b to
52d3f9b
Compare
|
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. |
|
Rebased/ported onto the sparse-attention framework refactor (#12733), new head 52d3f9b: the Indexer fork moved from |
|
/bot run --disable-fail-fast |
|
PR_Github #64772 [ run ] triggered by Bot. Commit: |
|
PR_Github #64772 [ run ] completed with state
|
2 participants
Description
The per-layer heuristic top-k feedback copy in
Indexer.sparse_attn_indexer(this step's decode top-k → next step's
pre_idxhint,metadata.heuristic_prev_topk[local_layer].copy_(...)) is a strided gatherthat sits on the main stream's critical path, once per indexer layer per
decode step. Nothing in the current step consumes it — the next reader is the
next decode step's same layer — so this PR forks it onto the Indexer's
existing aux stream right after the top-k kernel and joins it in the same
layer's MLA forward once core sparse attention has been enqueued. The copy
then overlaps with the layer's heaviest decode work instead of delaying it.
Design points:
forward, which CUDA graph capture requires (
cudaStreamEndCapturerejectsunjoined forks), and the join restores ordering before the next layer
overwrites the shared
topk_indices_bufferrows the copy reads. Nocross-layer or cross-step event state is needed.
graph-pool buffers; the captured copy stays valid on every replay and no
record_streambookkeeping is needed.do_multi_stream(), matchingmaybe_execute_in_parallelpolicy: the fork engages inside CUDA graph capture (where replay makes the
stream/event host overhead free); eager execution keeps the original inline
copy byte-for-byte.
forward_dsa_attnand theDeepSeek-V4 path in
forward_impl_with_deepseek_v4; DSA "shared" indexerlayers (
indexer is None) are unaffected.This also benefits the upcoming GVR top-k e2e wiring (#16420), which consumes
the same
heuristic_prev_topkfeedback loop.Test Coverage
join succeeds; replayed feedback values are step-correct across replays
(write@replay N is read@replay N+1); capture without the join fails
with
cudaErrorStreamCaptureUnjoined— confirming the join placement isload-bearing, not defensive.
test_fp8_blockscale[heuristic_topk_mtp1](TestDeepSeekV32,
cuda_graph=True) exercises the forked path end-to-endon Blackwell.
Performance Evidence (measured)
Paired interleaved A/B, 8xB200 (TEP8, BS=1), DeepSeek-V3.2-Exp FP4, one
SWE-bench prompt (ISL 68,656 tokens), OSL 1024, CUDA graphs +
enable_heuristic_topk: truein both arms. Arms differ only in this PR's twofiles (same build; per-repetition file swap, hashes logged); baseline arm runs
first in each pair. Each MTP config runs on a single node (pairs are
within-node; absolute times are not compared across nodes). Warm-up pairs
carrying first-load cost are excluded; TTFT (unaffected by this decode-only
change) is matched across arms in every reported pair.
The MTP=2 entry is a single TTFT-matched pair; its sweep was interrupted and
is being rerun. Remaining MTP 2-3 pairs on the same protocol will be added.
Per-pair data
MTP=0 (node A; pair 0 excluded as cold-start — its baseline leg carried
first-load cost, TTFT 20.9 s vs 3.15 s):
MTP=1 (node B; pair 1 showed no cold-start contamination — TTFT matched —
but is excluded for protocol consistency; including it: +2.30% over 10 pairs,
10/10 wins):
MTP=2 (node B; sweep interrupted after pair 1, rerun in progress; TTFT
matched, 3324.5 vs 3321.2 ms):
PR Checklist
Dev Engineer Review
do_multi_stream()is disabled.Indexer.maybe_join_prev_topk_copy()with pending-state cleanup.QA Engineer Review
No test changes.