Skip to content

[dev](perf): DSA pre-attention: plan-driven scorer dispatch, host-built CP layout metadata - #6206

Open
jingqiny-99 wants to merge 18 commits into
NVIDIA:devfrom
jingqiny-99:jingqiny/dsa-pre-attention-opt
Open

[dev](perf): DSA pre-attention: plan-driven scorer dispatch, host-built CP layout metadata#6206
jingqiny-99 wants to merge 18 commits into
NVIDIA:devfrom
jingqiny-99:jingqiny/dsa-pre-attention-opt

Conversation

@jingqiny-99

@jingqiny-99 jingqiny-99 commented Aug 3, 2026

Copy link
Copy Markdown
Contributor

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):

before after
aten::nonzero per step 1840 80

1. Indexer scoring dispatch as a plan

Choosing between the fused cuDNN scorer and the per-head fallback loop was spread
across two files: DSAttention.forward assembled use_local_indexer_varlen,
single_packed_thd_sequence and the packed cu_seqlens metadata, and
_indexer_topk_bshd re-combined them into a branch. A layout no kernel claimed fell
through to the slow loop silently. Measured with a branch-tracing build, 3 of 5
configurations took the unfused loop
:

configuration before after
non-packed, CP1, sharing off fused fused
non-packed, CP1, sharing on unfused fused
packed THD, CP16 fused fused
packed THD, CP1 unfused fused
non-packed, CP > 1 unfused unfused, now warns once

2. Packed-CP layout metadata built on the host

build_packed_allgather_cp_local_positions rebuilt positions per layer per rank from
inputs 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:

  • Memoize per microbatch on PackedSeqParams (constructed per microbatch, so the
    memo 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_size times.
    aten::nonzero 1840 → 80 per step; 708.8 → 597.6 ms (−15.7%).
  • Build from host integers. Since [dev] moe(perf): Refactor CP layout organization for Qwen3.5-style hybrid attention model #6387, prebuild_thd_cp_partition_routes
    already syncs the compacted cu_seqlens to the host at batch-construction time.
    Store those lists on PackedSeqParams; the position table becomes a closed form
    over 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:

  • Bit-equality unit tests for every layout builder (batched vs per-rank loop,
    host vs device, zero-length sequences, padding) and plan resolution pinned to the
    branch-trace measurements.
  • Top-k tie analysis for the fused-scorer change: over 240 comparisons the fused
    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_break already declines to pin down on CUDA.

@copy-pr-bot

copy-pr-bot Bot commented Aug 3, 2026

Copy link
Copy Markdown

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
jingqiny-99 force-pushed the jingqiny/dsa-pre-attention-opt branch from 99f9d47 to a539bfc Compare August 4, 2026 09:43
@jingqiny-99 jingqiny-99 changed the title [Dev] Cut per-layer CPU overhead in the DSA pre-attention path [Dev] Reach the fused DSA indexer scorer at CP1, and instrument the pre-attention path Aug 4, 2026
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>
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:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we let it reach the fused scorer when packed_seq_params is not None?

@jingqiny-99 jingqiny-99 changed the title [Dev] Reach the fused DSA indexer scorer at CP1, and instrument the pre-attention path [Dev] Resolve the DSA indexer scoring path as a plan, and close two gaps it exposed Aug 13, 2026
@jingqiny-99

Copy link
Copy Markdown
Contributor Author

/ok to test bd85b78

@jingqiny-99

Copy link
Copy Markdown
Contributor Author

/ok to test 4888213

@jingqiny-99

Copy link
Copy Markdown
Contributor Author

/ok to test 9fd365a

@jingqiny-99

Copy link
Copy Markdown
Contributor Author

/ok to test 6d619af

@jingqiny-99

Copy link
Copy Markdown
Contributor Author

/ok to test 130e767

@jingqiny-99

Copy link
Copy Markdown
Contributor Author

/ok to test 6fff840

@jingqiny-99

Copy link
Copy Markdown
Contributor Author

/ok to test b6fdf1c

@jingqiny-99

Copy link
Copy Markdown
Contributor Author

/ok to test 020ff52

@jingqiny-99

Copy link
Copy Markdown
Contributor Author

/ok to test 9715079

@jingqiny-99

Copy link
Copy Markdown
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>
@copy-pr-bot

copy-pr-bot Bot commented Aug 25, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@jingqiny-99

Copy link
Copy Markdown
Contributor Author

/ok to test c97d83b

@jingqiny-99

Copy link
Copy Markdown
Contributor Author

/ok to test 1a1006a

jingqiny-99 and others added 4 commits August 25, 2026 13:50
…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>
@jingqiny-99

Copy link
Copy Markdown
Contributor Author

/ok to test c2a35e4

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants