[dev](perf): DSA pre-attention: plan-driven scorer dispatch, host-built CP layout metadata - #6206
Open
jingqiny-99 wants to merge 18 commits into
Open
[dev](perf): DSA pre-attention: plan-driven scorer dispatch, host-built CP layout metadata#6206jingqiny-99 wants to merge 18 commits into
jingqiny-99 wants to merge 18 commits into
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
jingqiny-99
force-pushed
the
jingqiny/dsa-pre-attention-opt
branch
from
August 4, 2026 09:43
99f9d47 to
a539bfc
Compare
buptzyb
pushed a commit
to buptzyb/Megatron-LM
that referenced
this pull request
Aug 4, 2026
…tention (NVIDIA#6206) Signed-off-by: Robin Zhang <robinz@nvidia.com>
buptzyb
reviewed
Aug 6, 2026
| force on this common path. Genuine packed/CP/custom-position varlen is left | ||
| untouched (``varlen_is_plain_causal`` is False there). | ||
| """ | ||
| if cp_size == 1 and packed_seq_params is None and varlen_is_plain_causal: |
Contributor
There was a problem hiding this comment.
Can we let it reach the fused scorer when packed_seq_params is not None?
Open
33 tasks
Contributor
Author
|
/ok to test bd85b78 |
Contributor
Author
|
/ok to test 4888213 |
Contributor
Author
|
/ok to test 9fd365a |
Contributor
Author
|
/ok to test 6d619af |
Contributor
Author
|
/ok to test 130e767 |
Contributor
Author
|
/ok to test 6fff840 |
Contributor
Author
|
/ok to test b6fdf1c |
Contributor
Author
|
/ok to test 020ff52 |
Contributor
Author
|
/ok to test 9715079 |
Contributor
Author
|
/claude strict-review |
Enabling cross-layer indexer top-k sharing (dsa_indexer_topk_freq > 1) sets index_share for the whole model, so the `not self.index_share` guard on the fused whole-DSA path routes every layer through the split path -- including layers that own an indexer and compute their own top-k. The split path then loses the fused single-kernel indexer scorer. Plain causal attention still gets explicit varlen bounds from build_dsattention_forward_mask, and run_fused_qk_topk / run_fused_qk_topk_with_loss had no way to know those bounds were plain causal, so score_seq_lens stayed non-None and scoring fell into _compute_indexer_scores_chunk_with_global_rows -- a Python loop over dsa_indexer_n_heads issuing roughly five kernels per head. run_fused_dsa_attention already normalized these bounds away for exactly this reason. Extract that into _normalize_plain_causal_bounds and apply it in the two split-path entry points as well, threading varlen_is_plain_causal through the dispatcher the same way run_fused_dsa_attention already receives it. This preserves semantics rather than approximating them: with _INDEXER_RATIO == 1, _causal_seq_lens rebuilds ends = arange(1, sq+1).clamp(max=sk) and starts = 0, which is bit-for-bit what the mask builder emits for plain causal. The top-k validity mask degrades the same way -- _topk_in_bounds with starts None reduces to topk_indices < seq_lens, and with identity key positions the bounds branch computes that same predicate. Using the flag rather than a torch.equal comparison keeps the host/device sync off this path. The normalization runs after the `starts is None` decline check, so bounds the caller could not build at all are still a genuine decline. varlen_starts and varlen_ends in dsa.py are deliberately left untouched, because the sparse attention call downstream still needs them; only the indexer scoring scope is normalized. Measured on GB200 (9-layer GLM-5.2 proxy, TP1/PP1/EP4/CP1, MXFP8, topk_freq 4): the split path issued 256 extra bmm calls per profiled step (8 top-k computations x 32 indexer heads) and zero indexer_forward kernels, against 20 for the fused path. Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
…filer The pre-attention path carries no NVTX ranges and no record_function markers, and AbsorbedMLASelfAttention.forward overrides Attention.forward so MCore's own ...qkv / ...core_attention ranges never fire either. In a trace the entire region -- MLA down/up projection, the absorb, the packed-CP position and reorder build, the mask and bound construction -- collapses into one undifferentiated stretch of aten ops, and has to be identified by guessing at landmark operators. Add dsa_mark_begin/dsa_mark_end next to the existing NVTX helpers. They emit both torch.cuda.nvtx (nsys) and torch.profiler.record_function (chrome trace), gated on the same --nvtx-ranges switch, so normal training pays nothing. Begin/end rather than a context manager so the marked regions need no re-indenting; the stack is module-level and assumes a single-threaded forward. Regions: mla.down_proj, mla.up_proj_absorb, dsa.pre.cp_metadata, dsa.pre.cp_gather, dsa.pre.mask_bounds, dsa.indexer.qk_proj, dsa.sparse_attn. EOF Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
…g it Choosing between the fused indexer scorer and the per-head fallback was spread across two files: DSAttention.forward assembled use_local_indexer_varlen, single_packed_thd_sequence and the packed cu_seqlens metadata, threaded them through the dispatcher, and _indexer_topk_bshd recombined them into a branch. The decision only existed as the emergent result of those conditions, so a layout that no fused kernel claimed fell through to the slow loop silently. Add resolve_indexer_scoring_plan, which is total: every configuration maps to a plan and carries the reason, including UNFUSED_BOUNDS. report_unfused_scoring_once surfaces such a configuration once per distinct reason, so an unfused layout is visible rather than merely slow. No call site is switched over yet; this commit only introduces the decision and its tests. Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
…onditions _indexer_topk_bshd re-derived its branch from use_local_indexer_varlen, single_packed_thd_sequence and the packed cu_seqlens metadata, duplicating the conditions DSAttention.forward had already assembled. Resolve the plan once at the dispatcher entry and switch on it. This also removes the reason the split path had to delete information to go fast: _normalize_plain_causal_bounds returned (None, None) for plain causal bounds purely so the dispatcher would take its 'starts is None' branch. The plan now says PLAIN_CAUSAL directly, so both split entry points keep the bounds they were given. run_fused_dsa_attention still normalizes, because there the call additionally suppresses has_varlen to select the dense loss path; that one is a separate concern and is left alone. Behaviour is unchanged: the plan mapping reproduces each existing branch. Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
Threading varlen_is_plain_causal into FusedQKTopKWithSparseLossFunc and FusedIndexerSparseAttnFunc added a forward input to each without extending the backward return tuple, so training died at the first backward pass with 'returned an incorrect number of gradients (expected 25, got 24)'. Both Functions were short by one; only one of them surfaced in the failure. Add the missing gradient slots and a source-level test that pairs each Function's forward arity with its backward return. The mismatch is invisible to any test that stops at the forward, and the check needs no GPU. Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
…allelism The zigzag-segment kernel was gated behind cp_size > 1 in two places, so packed sequences at cp_size=1 fell to the per-head loop -- measured on GB200, where a packed CP1 configuration takes CHUNKS.SLOW_global_rows. The gate was never about what the kernel can do: _indexer_topk_single_packed_cp_segments already handles a rank holding the whole sequence, where its sk == local_query_len branch reduces to plain causal split into two halves. Restate both conditions as facts about the layout -- whether the pack holds one sequence, whether the mask is causal -- and let the scoring plan decide which kernel accepts it. Multi-sequence packs at cp_size=1 stay unfused and now say why: the cu_seqlens kernel genuinely requires cp_size > 1. Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
Widening the single-pack gate exposed that _indexer_topk_single_packed_cp_segments raises, rather than declining, when its structural preconditions fail. A packed cp_size=1 run therefore died with 'requires b=1, even local query length, ratio=1' instead of falling back to the unfused loop. Make the preconditions part of the decision: _segment_kernel_applicable lives next to the kernel so the two cannot drift, and the plan degrades to UNFUSED_BOUNDS with that reason when the shape does not fit. This also covers the pre-existing cp_size>1 path, which had the same latent crash. Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
_indexer_topk_single_packed_cp_segments split the local query in half because the zigzag context-parallel layout gives every rank a front and a back chunk of equal length. It then demanded an even local query length, which is a consequence of the split rather than a property of the scoring. A rank that holds the whole sequence has no zigzag to undo. Give that case one causal segment over the local range: per-row key lengths are still built from arange inside the loop, so the masking is unchanged, and the parity requirement disappears. Measured on GB200, a packed cp_size=1 configuration reaches this path with local_query_len=16383 and was rejected purely for being odd. Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
It duplicated a check autograd already performs, with a worse message: a
forward/backward mismatch raises 'returned an incorrect number of gradients
(expected N, got M)' at the first backward pass. Counting positional arguments
in the AST also approximates the real rule and would drift on *args or
keyword-only parameters.
Its one run cost a CI round: resolving the module path via
inspect.getfile(__import__('megatron')) raises TypeError because megatron is a
namespace package, and the collection error blocked the other 405 tests. The
gradient slots it was added alongside stay.
Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
_indexer_topk_bshd gained varlen_is_plain_causal and scoring_plan, so the three monkeypatch stubs that mirror its signature exactly stopped accepting the call: 'fake_indexer_topk() got an unexpected keyword argument varlen_is_plain_causal'. The other stubs in the file take **kwargs and were unaffected. Extend the three explicit signatures rather than loosening them to **kwargs, so they keep rejecting arguments the dispatcher is not supposed to pass. Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
CP query positions and the gathered-KV reorder index depend only on cu_seqlens_q/cu_seqlens_kv, cp_size, cp_rank and the requested output sizes. Every one of those is identical for all layers in a microbatch, yet each DSA layer rebuilt them from scratch. Each rebuild loops cp_size times over build_packed_allgather_cp_local_positions, and each of those runs five boolean-mask index operations whose output shape is data dependent, so eager PyTorch must read the resulting size back to the host. The explicit syncs on this path were already removed -- repeat_interleave is passed output_size= and the four .item() sites in dsa_layout are unreachable here -- but the mask-driven readbacks remain and cannot be avoided without either changing the algorithm or not running it per layer. Hang a memo on PackedSeqParams and reuse it across layers. _build_kv_reorder_idx uses the same key shape, so when its output sizes coincide with the earlier branch it reuses that result rather than rerunning the loop a second time. PackedSeqParams is deliberately the only carrier. It is constructed per microbatch, so an entry cannot outlive the layout it describes, which matters under dynamic CP where the layout changes between microbatches. The index-share top-k holders fall back to attention_mask/config, but that is only safe because every computing layer overwrites its slot before any sharing layer reads it; a layout memo has no such write-before-read ordering and config outlives the microbatch. Measured on GB200 (9-layer GLM-5.2 proxy, TP1/PP1/EP4/CP4, THD + dynamic CP, seq 16384): the unmemoized path issued 1040 aten::nonzero calls per step costing 265 ms of CPU, matching 10 DSA layers x 4 microbatches x (cp_size+1) builds x 5 masks, and the GPU idled 41.6% of the profiled window. Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
build_packed_allgather_cp_query_positions_and_key_reorder looped cp_size times over build_packed_allgather_cp_local_positions to assemble the gathered-KV order. Each iteration issued roughly twenty small kernels and five boolean-mask index operations whose output shape is data dependent, so eager PyTorch had to read the resulting size back to the host. At CP64 that is 65 builds and 325 readbacks per microbatch, on a path where larger CP also means fewer tokens per rank and so less GPU work to hide them behind. Every expensive step in that loop is rank-invariant. cp_rank enters only through front_starts and back_starts, which are cheap elementwise expressions. Critically segment_lens is stack(half_seq_lens, half_seq_lens), so it -- and therefore the nonzero and nonempty_segments masks, segment_ids, segment_offsets, and every data-dependent shape -- is identical for all ranks. Add build_packed_allgather_cp_all_rank_positions, which computes the invariant part once and broadcasts the rank offset over a leading cp_size dimension, returning [cp_size, output_size]. Row-major flattening reproduces exactly the rank0-local, rank1-local, ... concatenation the caller built by hand, so the argsort input is bit-identical. The readback count drops from 5 * cp_size to 5. The single-rank builder is retained: dsa.py still needs one rank's positions on the sequence-parallel path, and computing all ranks there would be wasteful. Tests assert row-for-row equality against the per-rank loop across cp_size 2-64, multi-sequence and zero-length packings, covered and uncovered cu_seqlens, and padded output sizes, plus that the public key reorder index is unchanged. Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
The packed-CP position and key-reorder builders run ~20 kernels over a few dozen integers, pay two device-to-host size readbacks for their boolean-mask filters, and finish with an argsort over the gathered key positions. All of it derives from cu_seqlens alone. Since NVIDIA#6387, prebuild_thd_cp_partition_routes already synchronizes the compacted cu_seqlens to the host once per microbatch at batch-construction time, where the CUDA queue is shallow. Store those host lists on PackedSeqParams (q and kv, deduplicated when they alias), and give dsa_layout host-side builders that consume them: the position table becomes a closed form over Python ints, and the key reorder becomes a direct inverse permutation -- the argsort was only ever sorting values that already tile [0, total) exactly once, so the rank of a value is the value, and the padding pseudo-positions are unique and already ordered. Zero kernels, zero synchronization, one asynchronous copy of each finished table. The device builders remain as the fallback for callers that never ran finalize_packed_seq_params. Correctness contract is bit-equality against the device builders, pinned by parametrized tests across cp_size, sequence mixes including zero-length sequences, and padded output sizes; the closed-form inverse was additionally fuzzed against an argsort reference over 300 randomized layouts. Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
Contributor
Author
|
/ok to test c97d83b |
Contributor
Author
|
/ok to test 1a1006a |
…and hooks From an adversarial review of the branch: - report_unfused_scoring_once had no production caller, so the promised once-per-reason warning for configurations landing on the unfused loop never fired. Wire it at the dispatch choke point in _indexer_topk_bshd. - The host layout builders diverged from the device builders they mirror at cp_size <= 1 with odd lengths: the span form floors seq_len // 2 and dropped the middle token, where the device builder returns the identity. Give both host builders the same dedicated cp_size <= 1 identity path. On host integers the 2*cp_size divisibility check is free (the device builders can only afford it without a sync on CPU inputs), so enforce it and assert the spans tile the key stream before writing into uninitialized reorder slots. New tests cover cp1-odd bit-equality and the non-divisible ValueError; the existing suite padded every length to even and could not see either. - Backend hooks are resolved dynamically and may live out of tree; passing the new varlen_is_plain_causal keyword unconditionally would TypeError an older backend at the first fused call. Filter kwargs by the hook's signature. - dsa_mark_end now honours what dsa_mark_begin actually did (a None sentinel is pushed when profiling is disabled), so toggling profiling inside a region can no longer desynchronize MCore's shared NVTX stack; the safety previously depended on toggle placement in the training loop, a different file. - Export the host builders in __all__, and document that non-packed callers of the kernel-side plan resolver are treated as causal (the dispatch outcome for a non-causal layout is unchanged; only the diagnostic reason is generic). Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
…used warning - Document in the three dsa_kernels hook docstrings that single_packed_thd_sequence / use_local_indexer_varlen now describe layout only and no longer imply cp_size > 1, so out-of-tree backends consult cp_size instead of assuming a zigzag split with even local length. - Extend the diagnostic-honesty comment in _resolve_scoring_plan_for_call to cover the packed_thd inference as well as the mask half. - Emit the unfused-scoring warning on the affected rank (once per process) instead of rank 0 only, so a non-zero rank sitting on the unfused cliff under heterogeneous packs is no longer silent. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
The pre-existing unit tests for prebuild_thd_cp_partition_routes pass SimpleNamespace stand-ins carrying only the q-side fields, which is the function's de-facto duck-typed contract. The host-list stash read cu_seqlens_kv_padded/cu_seqlens_kv directly and widened that contract, failing test_prebuild_thd_cp_partition_routes_populates_direct_fields with AttributeError. Read all four fields via getattr with a None default; a missing kv side degrades to aliasing the q side exactly as a None kv does. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
Review feedback: a dedicated begin/end layer (record_function handle stack + toggle sentinel) is more machinery than the job needs. Use the same nvtx_range_push/pop convention as the rest of MCore (e.g. TransformerLayer), keeping the range names unchanged so existing nsys comparisons stay valid. Also drop is_nvtx_profiling_enabled, which only the removed helpers used; utils.py is now untouched by this branch. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
Contributor
Author
|
/ok to test c2a35e4 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The DSA pre-attention path — indexer scoring dispatch, packed-CP layout metadata — was losing most of its wall time to dispatch overhead and host synchronization
rather than compute. Measured end to end on GB200 (GLM-5.2 proxy, THD 32K, CP16,
16 GPU, MXFP8):
aten::nonzeroper step1. Indexer scoring dispatch as a plan
Choosing between the fused cuDNN scorer and the per-head fallback loop was spread
across two files:
DSAttention.forwardassembleduse_local_indexer_varlen,single_packed_thd_sequenceand the packedcu_seqlensmetadata, and_indexer_topk_bshdre-combined them into a branch. A layout no kernel claimed fellthrough to the slow loop silently. Measured with a branch-tracing build, 3 of 5
configurations took the unfused loop:
2. Packed-CP layout metadata built on the host
build_packed_allgather_cp_local_positionsrebuilt positions per layer per rank frominputs identical across layers; each rebuild ran boolean-mask filters whose
data-dependent shapes force a device-to-host readback, and the key-reorder used a
full argsort. Fixed in three steps, each verified against a branch-trace or kernel
diff:
PackedSeqParams(constructed per microbatch, so thememo cannot outlive the layout it describes) and batch the per-rank loop — the
masks are rank-invariant, so the readbacks are paid once instead of
cp_sizetimes.aten::nonzero1840 → 80 per step; 708.8 → 597.6 ms (−15.7%).prebuild_thd_cp_partition_routesalready syncs the compacted
cu_seqlensto the host at batch-construction time.Store those lists on
PackedSeqParams; the position table becomes a closed formover Python ints and the reorder a direct inverse permutation — the argsort only
ever sorted values that tile
[0, total)exactly once. Zero kernels, zero syncs,one async H2D copy each.
Numerical verification
This workload is not run-to-run reproducible: the same commit run twice differs on
25/25 iterations (rel. 4e-6…4e-4) — MXFP8, HybridEP routing, and CP collectives all
contribute — so end-to-end loss cannot discriminate. The contracts used instead:
host vs device, zero-length sequences, padding) and plan resolution pinned to the
branch-trace measurements.
and split paths diverge only on rows whose selected score multisets are identical
or within 2 ULP of fp32, and the fused selection never scores lower — the same tie
freedom
_use_dense_indexer_topk_tie_breakalready declines to pin down on CUDA.