Repository navigation
Add the min-memory DSA attention kernels - #7152
Merged
zhongbozhu merged 20 commits intoOct 14, 2026
Merged
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. |
alokpathy
marked this pull request as ready for review
September 8, 2026 22:43
5 of 6 tasks
alokpathy
force-pushed
the
atripathy/dsa-min-memory-kernels
branch
from
September 9, 2026 05:36
addf186 to
5759290
Compare
This was referenced Sep 9, 2026
Contributor
|
Could you add a README for the exposed kernels (w/o auxiliary kernels) covering inputs/outputs, supported configs, and limitations? |
Member
|
/claude strict-review |
zhongbozhu
approved these changes
Sep 9, 2026
DSA computes attention over a top-k subset of keys chosen per query by a learned indexer. The straightforward implementation materializes the full index-score matrix to pick that subset, which costs O(sq*sk) memory and defeats the point of attending sparsely. These kernels stream instead: scores are computed and reduced tile by tile, so peak memory scales with the tile rather than the sequence, and only the selected keys are ever gathered. Two backends, both here. dsa_min_memory.py is the PyTorch reference: readable, differentiable by autograd where it can be, and the definition the other backend is checked against. dsa_min_memory_triton.py is the fused implementation, with the top-k selection, the selected-score forward and backward, the scatter of selected gradients back to sequence positions, and the weight gradients each written as a Triton kernel. Every one has a reference counterpart, and the tests assert they agree rather than asserting either is self-consistent. The backward pass is where the memory argument is won or lost, so it is explicit rather than autograd-generated. Selected-score gradients scatter into a sequence-length accumulator with repeated indices, which is why the scatter accumulates in fp32: bf16 accumulation of a key selected by many queries loses the tail of the sum. The wgrad path recomputes the input normalization rather than storing it, and falls back to exact recomputation for LayerNorm where the RMS shortcut does not apply. Eighteen tests, all of the form "fused kernel equals reference". Kept in their own file rather than with the attention layer's tests: these take plain tensors, so a numerical failure is attributable to one kernel instead of to the layer that called it. First of four PRs splitting the DSA-over-GQA work. Nothing here imports the attention layer; the layer that consumes these kernels follows next. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
CI runs tools/autoformat.sh with SKIP_DOCS=false, so pylint's missing-function-docstring is enabled. Eighteen functions lacked one: the timing profiler's two methods, the forward and backward of both autograd Functions, and the ten Triton wrapper entry points. The two prints in _DSATimingProfiler.log are deliberate -- it is a per-rank diagnostic printer and its output is meant for stdout, not the logger -- so they carry a bad-builtin suppression rather than being converted, matching the treatment in megatron/core/pipeline_parallel/fine_grained_activation_offload.py. pylint exited 20 on this branch, which is its bitmask for convention (16) plus warning (4), not a failure of tools/check_golden_values.py in the same job. Both files now rate 10.00/10. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
dsa_min_memory_gqa and dsa_min_memory_gqa_forward_only accepted a use_cudnn argument and never read it: the simplified path they dispatch to has no cuDNN branch. The module also imported the cuDNN DSA wrapper under a comment claiming it was "gated at runtime by use_cudnn", but nothing here ever referenced the import either -- every other _DSA hit in the file is _DSATimingProfiler. An argument a caller can pass and the callee ignores is worse than no argument: it reads as backend selection and silently does nothing. cuDNN selection belongs to the indexer dispatch in dsa.py, which is where the wrapper is actually called. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
The query and key chunk sizes control how the index scores are tiled: they set the temporary memory and the loop counts, not what is computed. Two fixed defaults (512 and 1024) decided that for every shape, and callers could override them per call. _plan_execution now derives both from the shapes and the backend. The score tile is [batch, query_chunk, key_chunk] in fp32, so the query chunk is sized to keep the widest tile inside a byte budget, capped at 8192. That matters at long context: with a full-length key chunk, a fixed 8192 query chunk would be 268 MiB of tile at sequence length 8192 and 8.6 GiB at 262144, which is the opposite of what a min-memory path is for. Routing under the torch backend still reads the whole key length. That is a correctness rule rather than a tuning choice: streaming torch.topk over key chunks is not tie-equivalent to one full torch.topk, exact ties are common after ReLU, and min-memory-torch has to reproduce reference routing to be usable as the oracle. _DSAExecutionPlan keeps it as a separate field so the distinction is visible rather than implied. The entry points keep their chunk arguments as overrides so tests can force multi-chunk paths on small tensors; nothing in the layer or the config supplies them. _SCORE_TILE_BUDGET_BYTES is 256 MiB, chosen to sit well inside a layer's activation headroom at the sequence lengths this path targets. It has not been tuned against measured peak memory. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
The main-attention-K variant was never used: every launcher enables the learned K, and sharing the attention keys pinned the indexer head dim to kv_channels and left no separate tensor for the indexer-K cache to hold. Thread the learned K unconditionally through the min-memory kernels. The entry points now always pass indexer.linear_k, so the autograd functions drop ctx.use_learned_k along with the main-K arms of the selected-score and indexer-loss backward paths. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
The kernels carried their own stage timer, median-window history and rank-gated printing, duplicating the profiler and timer infrastructure the training loop already exposes through --profile and --profile-ranks. Delete _DSATimingProfiler and the profile arguments threaded through the min-memory forward and backward paths. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
The execution planner capped the routing key chunk at 1024 for Triton, so at sequence 131072 each query tile issued 128 routing launches and re-ran the top-k merge 128 times. Measured on 12 nodes at sequence 131072: 82.5 s/step and 104 TFLOP/s/GPU, against 26.6 s/step and 321 TFLOP/s/GPU once the whole key length is routed in one pass -- 3.1x, with no other change. The cap bought nothing. triton_topk_index_block already streams the key dimension itself in 256-wide sub-blocks and merges the per-sub-block top-k the same way the outer loop does, so it never materialises the [batch, query_chunk, key_chunk] score tile the budget was charging for. An outer key chunk therefore saves no memory; it only repeats work. Budgeting that phantom tile also cost occupancy: with the routing width in the denominator, the query chunk shrank to pay for an allocation that never happened. Triton now budgets against the router's internal sub-block width, promoted to ROUTER_KEY_SUB_BLOCK so the two stay in sync. The Torch path is unchanged -- it does materialise the tile, and its full-length routing remains the tie-equivalence correctness rule it always was. Give query_chunk_size and key_chunk_size None defaults on all three public entrypoints while here. The tile sizes are not a model option, no caller in the layer passes them, and leaving them required means a caller that omits them raises TypeError inside a forward -- which megatron.core.utils catches and logs rather than failing the run, so it goes unnoticed. A signature test now pins the contract. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
The min-memory kernels take indexer.linear_q and indexer.linear_k as raw tensors rather than calling the owning Linear modules. That skips the forward pre-hook, which is what a sharding backend relies on to publish a parameter: with DDP and overlap_param_gather the all-gather is issued ahead of time and overlapped with other layers' compute, so weight.data is still the unpublished shard until the hook waits on it. Reading it early gives wrong values and raises nothing. Put the wait on _module_weight rather than on its callers. Every read of an indexer weight already goes through it, so a new call site cannot forget the wait -- which is the same thing that went wrong here in the first place. The cost is negligible where it does not apply. A parameter no backend marked carries no callback and the call returns immediately, so nothing changes without overlap_param_gather. A bucket that is already published returns on its own hot path, so the second weight of a pair does not wait again. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
_simplified_topk_index_tile already returns the scores of the keys it selected, and _sort_topk_support_by_position reorders scores and indices together, so they arrive aligned with topk_indices. Both the forward and the backward discarded them and then rebuilt the same values with _simplified_selected_index_scores. Reuse them. That also removes the learned-K branch here: routing scores against the same K the loss would have projected, and _project_simplified_k_index_block is pointwise with absolute rotary positions, so chunking it differently cannot change a value. The support-chunked rescoring loop and the second [0:q_end] K projection go with it. The recompute path stays for the case that still needs it: with dsa_simple_routing under selective recomputation the backward reads indices from the routing cache and never calls the routing kernel, so it has no scores to reuse. The two paths are different Triton kernels -- triton_topk_index_block's merge versus triton_simplified_selected_index_scores -- each verified against an FP32 oracle to 5e-3 but not against each other, and no unit test runs a full forward with use_triton=True. Checked with a 12-node A/B on the production config: the indexer loss agrees to 1.7e-7 at the first step, before any optimizer step makes the runs diverge, and lm loss is identical for two steps. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
Routing reprojects the causal K prefix inside every query tile, and the values do not depend on the tile. At sequence 131072 with a query chunk of 8192 that is sixteen prefixes of 8192, 16384, ... 131072 rows -- 8.5 sequences of projection where one would do. The same happens again in the backward. cache_indexer_k avoided it, but only by also retaining the projection in save_for_backward, so the choice was 8.5x the work or seq * batch * index_head_dim held from this layer's forward to its backward. Those are separate decisions. Project once in either case; let cache_indexer_k decide only whether the result survives into the backward, which projects it once itself when it did not. The eval path projects once too and keeps nothing, having no backward to feed. Forward and backward both drop from 8.5S to S with the cache off, at the same retained memory as before -- strictly better than what that setting did. With the cache on nothing changes. The cost is a transient full-length index K alive during the tile loop where previously there was none; it is freed when the pass returns. A per-tile [0:q_end] projection in the backward's selected-score path goes too: full_k_index is now always populated there, so it was unreachable. test_simplified_learned_k_only_persists_full_k_when_cached still pins the retention behaviour. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
The result is small -- [sequence, batch, 1, index_head_dim], tens of megabytes -- but building it in one call is not. _apply_indexer_input_norm_tile casts its input to FP32 and forms the normalized product there, so a full-sequence call materialises several temporaries the size of hidden_states in FP32. At sequence 131072 with hidden size 4096 that is about 2 GB for the cast alone, against 33 MB of result. Fill it in 8192-row chunks. The output is unchanged: the projection and its input norm are pointwise, and the rotary positions are absolute, so a position's value does not depend on the chunk it was computed in. All three sites take the same path -- the training forward, the backward when nothing was retained, and the eval forward. This peak is not new. The uncached path reached it too, because its last query tile projects the whole prefix; making the projection unconditional only moved it somewhere easier to see. The chunks are copied into a preallocated result rather than collected and concatenated, so only one block is live at a time; a cat would hold every block and the joined tensor together, costing more than the unchunked version it replaces. _INDEXER_K_BUILD_CHUNK is set to 8192 to match _MAX_QUERY_CHUNK. It has not been tuned against measured peak memory, like _SCORE_TILE_BUDGET_BYTES above it. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
upstream/main now requires every kernel-bearing file under megatron/ to carry a manifest entry, enforced by test_every_kernel_bearing_file_is_registered and by tools/check_kernel_determinism_coverage.py in the linting job. Both min-memory backends match the content patterns and landed before that invariant existed, so the determinism/kernels unit-test bucket fails on this branch. dsa_min_memory_triton.py matches on @triton.jit, dsa_min_memory.py on index_add_ -- the key, value and indexer-K gradients accumulate there whenever top-k selects one key more than once, so the sum order follows kernel scheduling. Both point at test_dsa_min_memory_kernels.py, which asserts torch.equal against the torch reference and pins the FP32 accumulation. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
_plan_execution budgeted the query chunk against ROUTER_KEY_SUB_BLOCK on the Triton path, on the grounds that triton_topk_index_block streams the key dimension internally and so never materialises a [batch, query_chunk, key_chunk] tile. The routing key chunk was then set to the full key length, since an outer chunk would buy no memory. That reasoning only holds when the streaming router actually runs, and it does not: triton_topk_index_block rejects any key block over 2048 before reaching its sub-block loop, so routing falls to the dense scorer with a tile sized for the full causal prefix. At sequence 131072 the planner sized the query chunk for a 256-wide tile and the code allocated a 131072-wide one -- 4 GiB against a 256 MiB budget, 16x over, growing with the causal limit across query tiles. Invert the solve on the Triton path. Hold the query chunk at the occupancy cap, where shrinking it would cost launches without bounding anything, and derive the widest key block whose tile fits the budget. The outer routing loop and _merge_topk already iterate the remainder. Torch is unchanged: it may not chunk keys without breaking tie-equivalence with a single full topk, so it keeps the full key length and absorbs the budget by shrinking the query chunk instead. The budget now bounds the tile for every shape rather than for an assumed width: batch 2 and longer sequences narrow the key block instead of overrunning. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
The previous 256 MiB was a starting point, documented as untuned against any measurement. With the budget now driving the key block width rather than an assumed one, it is the dial that decides how many passes routing makes over the keys, so the value matters in a way it did not before. 4 GiB covers the whole 131072-key prefix in a single block at batch 1, which is the shape the 8B long-context runs use. One wide block costs one kernel launch and one merge where sixteen narrow ones cost sixteen of each, and the tile is transient -- freed before the next query chunk -- so this is a few percent of an 186 GB device held briefly rather than anything that accumulates. Smaller values stay correct and simply run more passes, which is the knob to reach for on a smaller device or at larger batch. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
_simplified_topk_index_tile tried triton_topk_index_block first and fell back to a dense score block when it returned None. The fused kernel selects while scoring, so it never holds a score tile, and the fallback exists for the shapes it cannot handle. Now that _plan_execution sizes the key block from the score budget, that division of labour has nothing left to divide. The block is bounded by construction, so the tile the fused path avoids is already affordable, and the key dimension is streamed by the loop around this call rather than inside the kernel. Selecting the dense path by "the fused one declined" also made the backend a function of a kernel precondition rather than of the plan, which is how a misscoped key-width check silently became the routing policy. Score the block and take its top-k unconditionally. unit_weights went with the fused call: it existed to drive the weighted standard-indexer kernel from the simplified path, which has no per-head weights. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
triton_topk_index_block selected the top-k while scoring, so it never held a score tile, and it streamed the key dimension itself in 256-wide sub-blocks. That made it the memory-safe routing path and the dense scorer its fallback. Nothing has run it since the execution planner began routing over the full key length: _can_use_index_scores rejects any key block over 2048, which is a BLOCK_N register bound belonging to _triton_topk_index_block_query_once, applied to a wrapper whose whole purpose is to make the key width not matter. Both backends have been scoring densely, min-memory-torch because Triton dispatch is off and min-memory-triton because the check rejects it. A five-step run at sequence 131072 takes the dense branch on all eighty routing calls and the fused branch on none. With the key block now sized from the score budget, the path has no role to reclaim. The budget bounds the tile the fused kernel existed to avoid, and the key dimension is streamed by the routing loop rather than inside a kernel. Its remaining advantage would be skipping the score write to global memory, but it pays for that with a cat/topk/gather per 256-wide sub-block -- more traffic than the tile it saves, and the 3.1x the dense path measured against it. Removes the two kernels, their autotune configs, the three wrappers, _merge_topk_tensors and ROUTER_KEY_SUB_BLOCK. _merge_topk, which merges the outer key blocks, is a different function and stays. _MAX_TRITON_SUPPORT_TOPK stays: six other call sites bound topk with it. The two deleted tests covered the sub-block merge. The equivalent property now lives on the outer merge, which production reaches whenever a prefix exceeds one tile, so the numerical-optimality test is ported there rather than dropped. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
dsa_dense_indexer_loss sized its chunks from _plan_execution, which models the routing tile: [batch, query_chunk, key_chunk] for the simplified indexer's single head. The dense teacher scores every main-attention head against the same keys, so its logit tile is [batch, num_query_heads, query_chunk, key_chunk] -- 32x the budgeted size at 32 heads, and the plan had no head count to notice with. _MAX_KEY_CHUNK kept it bounded in practice, since the teacher streams keys on the outside at 1024 wide. But that is a constant unrelated to the budget, which is the same way routing came to allocate 16x what it was sized for. Add _plan_dense_warmup, taking the head count and solving the query chunk against it. The two plans return disjoint fields -- routing reads routing_key_chunk, warmup reads key_chunk, neither reads the other -- so a caller can no longer pick up a bound computed for a different tile shape. _DSAExecutionPlan.key_chunk had no remaining reader and is dropped. Two full-size copies also had to go before the bound held. _dense_teacher_logits_ block collected per-group blocks and concatenated them, holding every group and the joined tensor at once; it now writes each group into a preallocated result. _mask_dense_causal_scores then returned masked_fill rather than masked_fill_, holding the masked copy beside the original, which by itself put the teacher at twice the budget. Both callers pass a tile they just built and never read again, and both run without autograd tracking -- routing under no_grad, the teacher inside an autograd.Function forward -- so masking in place is safe. Tests parameterize over head counts and assert the realized tile stays inside the budget, and one measures peak CUDA allocation against the planned tile so a reintroduced copy fails here rather than at a user's OOM. That test is what caught the masked_fill copy. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
_sparse_attention_backward_torch_fp32 was the Torch fallback for _output_consistency_backward, which computed the gradient of the sparse/dense output-consistency loss -- one of the two DSA main-attention auxiliary losses. Removing those losses deleted the caller and left the fallback behind, so it has had no production caller since. The main sparse-attention backward has always replayed the forward under torch.autograd.grad and never routed here. Its docstring still advertised it as "the Torch fallback", which sends a reader looking for a call site that no longer exists. Its test asserted that the Torch backward accumulates repeated keys in FP32 -- true of the removed aux-loss path, not of anything that now ships -- so a green test stood behind a claim about production that had stopped being true. _selected_causal_invalid_mask and _gather_selected_kv have other callers and stay. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
The fused sparse-attention backward runs only when query, key and value all need gradients. Below that the backward replays the forward under torch.autograd.grad instead -- which cannot work while Triton dispatch is on, because _sparse_attention_tile returns a raw kernel result with no grad_fn and autograd.grad has no graph to walk. The two gates disagree: the fused backward tests needs_input_grad, while the Triton forward tests only shapes and dtypes, so a mixed configuration declines the backward and still takes the Triton forward. Only min-memory-torch escapes it, by disabling Triton dispatch and reaching the differentiable einsum path. Nothing hits this today -- Q, K and V all come from linear_qkv and all need gradients -- but a frozen or detached key or value would, and the failure would surface several frames inside autograd with no connection to the cause. Raise at the gate instead, naming which of the three disagreed. All-false stays legal: nothing needing a gradient is an ordinary inference or frozen-backbone call, and the existing attention_inputs guard already handles it. This also narrows min-memory-torch, which would otherwise differ from min-memory-triton here. The backends accepting different configurations silently is the shape of the other bugs in this file, so the guard applies to both. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
Three expressions this review round introduced fit on one line under the project's 100-column black profile and were left wrapped. Collapses the _DSAExecutionPlan construction, the dense teacher's per-group einsum, and the all-or-nothing gradient gate. Signed-off-by: Alok Tripathy <atripathy@nvidia.com>
Contributor
Author
|
/ok to test de0dd04 |
kvareddy
approved these changes
Oct 10, 2026
cspades
approved these changes
Oct 13, 2026
kwangjunahn
approved these changes
Oct 13, 2026
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38496925256 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38502930273 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38509268666 |
This branch was previously deployed
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.
What does this PR do?
This is the first of four stacked PRs that introduces training support for a simplified version of DSA. These stacked PRs split PR #6869.
This PR includes Triton kernels and autograd support for simplified DSA with GQA attention. Included are triton kernels for both the indexer and sparse GQA attention, along with wrapper classes and torch autograd functions. A summary of the changes included in summary of the changes in each file is included below.
Miscellaneous notes:
seq_len x seq_lenQK matrix.megatron/core/transformer/experimental_attention_variant/dsa_min_memory_triton.pymegatron/core/transformer/experimental_attention_variant/dsa_min_memory.pyDSASimplifiedMinMemoryGQAFn,DSASimplifiedDenseIndexerLossFn), tile-wise projection/routing helpers, per-stage timing profiler.Part of a series
This PR is 1 of 4 splitting a single 13k-line PR #6869
Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.