[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 |
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
|
|
/bot run --disable-fail-fast |
|
PR_Github #65523 [ run ] triggered by Bot. Commit: |
|
PR_Github #65523 [ run ] completed with state |
| 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 |
There was a problem hiding this comment.
Please simplify comments from agent. :)
|
|
||
| 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 |
There was a problem hiding this comment.
Same here, first line is enough.
3 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)
Two instruments, mutually consistent: a per-iteration nsys A/B that resolves
the effect (statistically significant), and an end-to-end TPOT A/B that
correctly reads zero — the effect size (~0.2% of TPOT) is below what e2e
timing can resolve.
nsys per-iteration A/B (mechanism + magnitude)
DeepSeek-V4 Pro FP4, 8xB200, TEP8, BS=1, ISL 64K (real-data prompt), OSL 2048,
MTP=0,
enable_heuristic_topk: true, CUDA graphs ON. Arms differ only in thisPR's files on the same C++ build; ABBA run order (PR r1 → base r1 → base r2 →
PR r2); nsys window = decode iters 500–550; per-iteration times taken from a
once-per-iteration anchor kernel (49 deltas per run).
Mann-Whitney U z = −4.60 (p ≈ 4e-6); all four runs direction-consistent.
Mechanism, from the same traces: in the base arm 0/300 feedback copies overlap
any other-stream work — the copy is a serial node between the top-k kernel and
the index transform, with ~1.5 µs node-dispatch gaps on each edge. With this
PR, 300/300 copies overlap core sparse attention on the aux stream, and the
join introduces no main-stream bubble.
−15 µs/iter over 30 indexer layers ≈ 0.5 µs/layer removed, i.e. roughly the
copy-kernel duration leaving the critical path (the per-edge dispatch gaps are
not recovered — they re-materialize on the surviving top-k→transform edge).
This bounds the effect at 0.14–0.23% of TPOT for this config, and ≤ ~0.4%
for any realistic config (more indexer layers / MTP raise the per-iter µs, but
iteration time grows alongside).
e2e TPOT A/B (null result, consistent with the bound)
DeepSeek-V3.2-Exp FP4, 8xB200, TEP8, BS=1, MTP=3, same paired interleaved
protocol (one SWE-bench prompt, ISL ~68.7K, within-node pairs, TTFT-matched),
10 warm pairs: paired mean TPOT reduction −0.42% (mean-TPOT) / −0.16%
(median-TPOT), SE ≈ 0.7% — statistically zero. This is the expected reading:
the predicted effect for this config is +0.18–0.36% (≈64 copies/iter ×
~0.5–0.8 µs on a ~5.6 ms TPOT), and resolving it against the observed per-pair
σ ≈ 2.2–2.4% would take on the order of 250–1400 pairs. The e2e instrument
cannot see this change; the per-iteration instrument does.
What this PR is (and isn't)
A structural win with a small, rigorously measured latency gain: it removes a
serial copy node from every indexer layer's decode critical path (proven
off-path in traces), at zero cost to the eager path, reusing the per-layer
fork/join pattern already shipped for
maybe_execute_in_parallel. It is notan e2e-visible speedup on its own; its value compounds with anything that
lengthens the per-layer overlap window or shortens iteration time.
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.