Repository navigation
Conversation
Squash of NVIDIA#6472 at head 1aff38f (same tree): the PR's three commits, a merge of dev c4220af carrying the adaptations below, and a commit with the dev test doubles. Upstream commits: - 58de1df refactor(transformer): add generic forward sharing state Give the test doubles the state that dev's TransformerBlock.forward and checkpointed absorbed-MLA attention now read: TE CUDA-graph route metadata staging (dev NVIDIA#6058), the mHC recompute flag (NVIDIA#4531), and pg_collection.cp (dev NVIDIA#6206). - 279630e feat(dsa): share sparse attention indices across repeated MTP depths Reject repeated-layer sharing only when the resolved CUDA graph scope captures attention (graph_captures_attention). Since dev NVIDIA#6058, cuda_graph_captures_attention is the imported scope helper, so testing it directly is always true and rejects every sharing configuration, including cuda_graph_impl='none'. Also count only depth-0 indexer executions in the DSA indexer-metric normalization added by dev NVIDIA#6980. - f5ada35 feat(dsa): share latent KV across repeated MTP depths Reuse the canonical post-TP/CP latent KV from the first repeated MTP invocation while preserving its autograd graph. Keep sparse-index sharing independent and account for each component separately in FLOPs and tests. Keep the reused-KV key path aligned with dev: apply key RoPE with the dynamic-CP effective_cp_group and rope_freqs_max_seqlen (NVIDIA#5944; rope_max_seqlen_kv is undefined for unpacked inputs), and return no key for a latent-KV consumer on the NoPE path added by dev NVIDIA#8287. Adaptations to the current dev (main-to-dev sync NVIDIA#7960 and later): - GPTModel.forward and the MTP block depth loop keep dev's new code (MTP-only training, mtp_input_mask, hidden-state mixing, precomputed embedding roll, CP layout views) inside the forward-sharing lifetimes. - Dev now runs TransformerBlock full recompute through the common megatron.core.recompute.checkpointed_forward (main NVIDIA#6648), so the forward-sharing checkpoint wrapper moves there. That path also serves HybridStack, so its full recompute is wrapped too when forward sharing is enabled. The wrapper only pins what the block-level carrier holds: HybridModel opens no forward-sharing lifetime (GPTModel, TransformerBlock and the MTP block do), and hybrid layers with a CP layout plan read per-layer PackedSeqParams, so for hybrid models it often has nothing to pin. - Test doubles follow dev's MTP block/layer and recompute interfaces. Dev test doubles: three of dev's existing CPU tests build doubles that lack state this PR now reads, so they fail with AttributeError. test_absorbed_mla_contracts.py gains a config (read by is_forward_sharing_enabled in the checkpointed attention forward), test_mla_cp_runtime_cpu.py gains core_attention (read by AbsorbedMLASelfAttention.forward for cross-depth sharing), and test_mtp_roll_axes_cpu.py's block config gains mtp_repeated_layer_shared_components=None. The PR's own block double in test_forward_sharing.py gains pg_collection.tp, which dev's common recompute path passes to precision_aware_checkpoint (NVIDIA#7709), and the MTP guide spells the CP communication type as cp_comm_type=all_gather. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
…TP sharing (PR 7107) Squash of NVIDIA#7107 at head 96e0016 (same tree; stacked on PR 6472). Merging the updated PR 6472 brought one adaptation to dev c4220af: dev checkpoints the MTP layer through precision_aware_checkpoint (NVIDIA#7709), so the depth-0 shared tensors are passed to it as extra positional inputs (*shared_inputs), and the MTP layer double in test_mtp_dsa_recompute.py gains tp_group. Carry depth-0 shared tensors as explicit checkpoint outputs and consumer inputs so full recompute preserves the latent-KV autograd graph. Snapshot ordinary IndexShare state per checkpoint invocation, checkpoint only the query up-projection when latent KV is shared, and reject the unsupported latent-KV plus selective core-attention combination. Build the split Q/K primitives on dev's absorbed-MLA code: reuse the effective_cp_group that get_query_key_value_tensors already captures for checkpoint replay, select RoPE rows by its size, apply RoPE with rope_freqs_max_seqlen (NVIDIA#5944 and dev NVIDIA#6206), and keep the NoPE path from dev NVIDIA#8287. Give the test doubles what these paths now read: pg_collection.cp, use_rope, and the projection context in the MTP checkpoint tests, and core_attention in dev's microbatch-CP-group RoPE test, which reaches the split recompute dispatch. Full-iteration CUDA graphs: PR 6472 rejects mtp_repeated_layer_shared_components whenever a CUDA graph captures attention, because a per-layer graph would freeze one depth's shared state (latent KV, sparse attention indices) while the producer or a consumer runs outside it. A full-iteration graph (cuda_graph_impl='full_iteration') captures the whole forward and backward, so the producer and every consumer depth run inside the same graph in their eager order, and replay reproduces the sharing. Exempt cuda_graph_impl='full_iteration' from the check, and say in the error that only per-layer scopes that capture attention are rejected. The config test now expects that message for the transformer_engine and local cases, drops the full_iteration rejection case, and test_cross_depth_sharing_accepts_full_iteration_graph accepts each list of shared components with a full-iteration graph. The MoE-router-only case is already covered by PR 6472's test_cross_depth_sharing_accepts_graph_scopes_without_attention. The compatibility list in the MTP user guide says the same: per-layer scopes that capture attention are unsupported, MoE-only scopes and full-iteration graphs are compatible. Dev has no other whole-step graph implementation; if chunk-wise graphs (NVIDIA#5258) land, the exemption has to be revisited for them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
…es (PR 8535) Squash of NVIDIA#8535 (head ce54a6c, same patch-id). When the last pipeline stage holds only MTP layers, its local decoder block spec is empty. Both MTP spec builders take their "no decoder layers" branch before they look at experimental_attention_variant and fall back to the standard-attention layer spec, so a DSA model with an MTP-only last (virtual) stage builds its MTP layer with ordinary attention on that stage: - gpt_builders.gpt_builder tested for empty layer_specs before the experimental-attention branch, and its else branch held an experimental-attention case that could never run. - megatron.training.models.gpt.mtp_block_spec used _te_or_local_layer_spec for an empty TransformerBlockSubmodules. With experimental_attention_variant set, both builders now take the MTP layer spec from the global experimental-attention pattern (the last entry of get_transformer_layer_with_experimental_attention_variant_spec) before the empty-stage fallback, and the unreachable branch is removed. Tests cover gpt_builder on an empty decoder stage (test_gpt_model.py, which needs Transformer Engine to collect) and mtp_block_spec with an empty block spec (test_gpt_builder.py). The latter passes a real TransformerBlockSubmodules(layer_specs=[]), as its neighbouring test does. mtp_block_spec checks isinstance(spec, TransformerBlockSubmodules), so the Mock with empty layer_specs that the test used before never reached the branch under test and failed even with the fix. The main-to-dev sync NVIDIA#7960 added an experimental-attention branch to gpt_builder after the empty-stage check, which an MTP-only stage still never reaches; moving the experimental case first also absorbs that branch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
Squash of NVIDIA#8536 at head aef0af1 (same patch-id): the tail-less certificate fix, and the empty-query guard in fused_qk_topk_kpool with its tests and test_kpool_contract.py. Tail-less rows. A top-k-computing DSA layer can prove from its layout that every query row selects at least one key, and then skips the empty-row masking around sparse attention (can_prove_all_topk_rows_nonempty). KPool selection (NVIDIA#8287, dsa_indexer_kpool > 1) without dsa_indexer_kpool_always_select_tail leaves the rows before a query's first complete pool without a key (all -1), which the layout check cannot see, so such a layer claimed the certificate for rows that are empty and the masking was skipped for them. Do not certify when the layer runs KPool selection without its tail. Layers without KPool, and KPool with tail selection (the default), are unchanged. A CPU test builds a one-layer DSA layer from a local config (no index sharing, so it runs on dev alone), turns on KPool selection and runs it with a plain causal layout, with and without tail selection: without the tail the selection has empty rows and the certificate is now False (it was True); with the tail every row selects a key and the certificate stays True. Empty query rows. fused_qk_topk_kpool flattens the selected pools with pool_topk.reshape(rows, -1). When q has zero query rows, rows is 0 and the reshape cannot infer the last dimension, so the call raises "cannot reshape tensor of 0 elements into shape [0, -1]". Packed context-parallel splits can leave a rank with no query tokens, and that rank then crashes in the KPool indexer instead of returning an empty selection, although the helpers in the same path (_compute_index_scores and _kpool_topk_from_score_chunks) already handle sq == 0. Return early when q.size(0) == 0 with int32 indices of shape [batch, 0, index_topk + pool_size - 1] (or index_topk when always_select_tail is off) and float32 scores of shape [batch, 0, sk // pool_size], or None for scores when return_index_scores is False. These dtypes and widths match the non-empty output. Tests: test_kpool_causal_tail.py gains test_kpool_empty_query_keeps_fixed_width (CPU, both return_index_scores settings). The new test_kpool_contract.py covers selector-contract validation, a frozen selector without auxiliary loss, per-segment compression against an independent torch reference, score-free selection without an autograd graph, and finite zero gradients for empty absorbed rows. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
…pool origin Add two default-off options that reproduce the KPool indexer arithmetic of the Hugging Face reference implementation (huggingface/transformers v5.16.1, Glm5NextTextIndexer), and start explicit-mask pools at each sample's first valid key, as that reference does. Reproducing zai-org/GLM-5.3-Flash, the public model this reference implements, needs both options. - dsa_indexer_kpool_key_dtype_compression. Cast the FP32 pool softmax to the key dtype before multiplying and reducing, as Glm5NextTextIndexer.get_pooled_states does. False keeps the FP32 weighted reduction and BF16 output of the existing KPool implementation (NVIDIA#8287). The rounding is independent of Hadamard rotation and score-input quantization. - dsa_indexer_hf_scoring. Upcast the BF16 projected head weights before the head-count scale, scale the FP32 q@k before ReLU, and reduce heads with an FP32 matmul, in the reference's order, for both materialized and chunked scores. TransformerConfig requires KPool > 1, BF16 projection output and ReLU scoring with this option. Grad mode. KPool selection is discrete, and _kpool_topk_from_score_chunks runs under no_grad, so a caller in grad mode records no autograd state for the chunked scores. _compute_kpool_hf_scores has no grad-mode policy of its own: like the legacy _compute_index_scores it follows the caller, so materialized HF-order index scores carry gradients when a caller enables them. No caller in megatron/ does, and KPool rejects the indexer loss. Explicit-mask pool origin (always on). Without packing, fused_qk_topk_kpool pooled keys on the global grid [0:4], [4:8], ... and only filtered the pools with the mask, while the reference starts the pool grid, and places the incomplete tail, at each sample's first valid key. With left padding that is not a multiple of the pool size the two select different pools and tails. Infer each sample's first valid key on device, as the first key visible to any query row of a shared [sq, sk] or a [batch, sq, sk] mask, shift the key, gate and mask key axes left by its offset modulo the pool size, run the existing pooling, scoring, selection and tail logic, and add the phase back to the selected tokens; the final token masking uses the caller's mask. The change is an identity unless an explicit mask is not pool-aligned: a zero phase is exact, so no-mask, varlen, packed and pool-aligned explicit masks keep their selections and index-score coordinates bit-for-bit, and there is no host sync. Contiguous left and/or right padding with causal rows matches the reference when padded query rows are fully masked; masks with gaps, sliding windows or row-dependent starts keep the earlier boundary handling without a parity claim, and packed THD is unchanged. Model-level left padding stays unsupported because causal DSA layers ignore the attention mask. Tests. Add the two fields to GOLDEN_CONFIG in test_hybrid_moe_model.py so test_constructor passes. test_kpool_contract.py checks the config gates (dsa_indexer_hf_scoring rejects the precision paths it does not implement, and the compression flag is independent of score quantization), that the default compression is unchanged, the reference compression rounding exactly, packed pooling with segment resets, and the selector against an independent padding-vector transcription of the reference, including a counterexample for the global grid and per-sample padding, and, on CPU, the grad-mode placement above: the chunked selector saves no tensors for autograd in grad mode, with a positive control for the probe, and the HF-order scores follow the caller's grad mode. test_kpool_causal_tail.py adds a left-padding invariance test. The reference transcribed by the tests is huggingface/transformers v5.16.1 (commit 93c8b7b485963a10800c91f55304db6be211c2bd), src/transformers/models/glm5_next/modeling_glm5_next.py: Glm5NextTextIndexer.forward, get_visible_tokens, get_pooled_states and append_visible_tail. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
…acked-CP layouts Two context-parallel improvements for DSA: a gather-based backward for the CP key reorder, and packed-CP layout tables that are built at batch construction instead of inside the first DSA layer's forward. Inverse-permutation backward for the CP key reorder With all-gather context parallelism, DSA reorders the gathered latent K/V and the indexer K into full-sequence order with index_select(0, kv_reorder_idx). Autograd's backward for that is a zero fill plus an atomic index_add, which runs serially right before the dKV reduce-scatter (on GB300, about 4.6 ms + 1.1 ms per layer per microbatch at 1M tokens with CP64, 0.6 + 0.1 ms at 128K with CP32). The reorder is a permutation, so the backward is a plain gather with the inverse permutation. Add _PermuteRows, have the layout builders optionally return the inverse (the host builder checks bijectivity once while inverting; the device and zigzag builders get it for free), memoize it next to the reorder and use it at the K/V and indexer-K reorder sites. The result matches index_select autograd exactly up to the sign of zero. KPool selection (NVIDIA#8287) gathers the KPool gate score next to the indexer K and reorders it with the same kv_reorder_idx. The gate keeps the existing index_select: KPool selection runs without grad (KPool rejects the indexer loss, and the gate is created with requires_grad=False), so it has no backward to speed up. test_dsa_kpool_cp_reorder.py checks on CPU, at CP2 and CP4 in train and eval with a faked CP gather, that the KPool selector receives the indexer K and the gate at the global length in the same row order and without grad; it fails if the gate is not reordered. Prebuilt packed-CP layout tables Upstream NVIDIA#6206 builds DSA's packed context-parallel layout (all-gather CP query positions and the gathered-key reorder) from the host cu_seqlens that prebuild_thd_cp_partition_routes records at batch construction, but it still builds the tables lazily inside the first DSA layer's forward and memoizes them on the per-microbatch PackedSeqParams. That build cannot run inside a CUDA graph capture, and a memo hit from an eager warmup would be baked into the graph and go stale on replay, which is why NVIDIA#6206 deferred graph support for changing packed layouts. Build the tables before the forward instead: - dsa_layout: PrebuiltPackedCPLayout carries query_positions [T] and kv_reorder_idx [T * cp] together with the PackedCPLayoutGeometry (cp_size, cp_rank, local rows, cover flags) they answer. prebuild_packed_cp_layout builds it with the NVIDIA#6206 host builders, making the exact call of the memo path so the tables are bit-identical, for the realized CP group carried by PackedSeqParams. It attaches the layout bound to the host boundary lists it was built from, so metadata whose routes were recorded again never serves a stale layout. Prebuilding is idempotent for unchanged metadata and builds nothing during a capture. copy_() refreshes a layout of the same geometry in place with Tensor.copy_, keeping its storage for static graph inputs; copy_prebuilt_packed_cp_layout_() does the same for a static PackedSeqParams and also re-points its host boundaries and drops its lazy memo. Tables are enumerated by field, so an inverse reorder can be added as one field. - dsa: the three layout memo sites take the prebuilt tables in place of a memo entry when their geometry and device match the request (SP query positions, CP positions + reorder, K/V and indexer-K reorder), after validating them. Without a layout DSA keeps the lazy memo path, and on a mismatch it warns once per kind of mismatch and falls back. Inside a CUDA graph capture both cases raise, as does a layout inconsistent with its own geometry. The cover flags come from one helper shared with the prebuild. - training: prepare_packed_seq_params prebuilds the layout for DSA models right after the routes, except for raw-boundary middle PP stages whose physical row count is not known there; those detach any earlier layout. With the inverse-permutation backward of the CP key reorder, DSA's memo entry is (positions, reorder, inverse), so the carrier holds a third, required table, kv_reorder_inv_idx [T * cp]: the one field the dsa_layout item above leaves room for. It comes from the same single host-builder call with return_inverse_key_reorder=True, which inverts the reorder (and checks that it is a permutation) on the host before copying either table, and it is validated, refreshed by copy_() and compared like the other tables; validate() also checks that it inverts the reorder when the tables are on the host, where that needs no device read. Both reorder match sites return it, so a matched layout equals the memo entry bit for bit, inverse included, and _PermuteRows saves the prebuilt inverse for its backward: the refresh-after-backward rule covers that table too (the KPool gate reorder runs without grad and saves nothing). It costs one more [T * cp] int64 table (8 MiB per microbatch at 1M tokens with CP64) and 4 to 12 ms more host time per build there (1 to 512 packed sequences, measured on the host without the extra 8 MiB host-to-device copy). The CPU tests compare the inverse bit for bit and check that the K/V and indexer-K backward gathers with the prebuilt inverse, with and without sequence parallelism. The generic context_parallel_layout code is unchanged and does not import DSA. Under dynamic CP the table sizes follow the realized CP size, so they are static per CP size only. The host build moves from the first DSA layer into get_batch, and so into the batch-generator timer: it is O(cp_size * num_seqs) Python, about 220 ms per microbatch on a GB300 node at 1M tokens with CP64 and 512 packed sequences. Wiring the tables into a specific graph implementation is left to a follow-up; a slot must not be refreshed before the backward of every microbatch that read it has run. Tests: CPU bit-identity against the memo path (CP 2/4/8, every rank, zero-length entries, all-empty packs, uncovered padding rows, distinct Q/K boundaries, every cover-flag combination, SP), DSA forward consumption with builder spies, fallbacks, simulated graph capture, stale metadata, idempotent prebuilds, in-place refresh and dynamic CP. A 2-GPU test captures a CP2 packed DSA forward/backward into a CUDA graph and replays it after refreshing cu_seqlens and the prebuilt tables for another packing, with a stale-table negative control. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
Under activation recompute the first call of a checkpointed top-k DSA layer runs without grad, so it selects top-k without the indexer loss, and its backward replay reruns the -inf fill, the dense IndexerForward and the chunked TopK only to recover the selected scores for the sparse KL indexer loss: about 95 ms per top-k layer per microbatch at 1M tokens (CP64, GB300). Add dsa_indexer_reuse_checkpoint_topk (default off). The forward-sharing checkpoint wrapper marks a checkpoint's no-grad first call and its replay with an explicit CheckpointPhase. The first call also keeps the selected fp32 scores next to its published index-sorted indices, and the replay computes the sparse indexer loss from them with SparseIndexerLossFromTopKFunc (sparse attention for the LSE, then the existing sparse loss and precomputed gradients) instead of rerunning the indexer. A missing stash falls back to the previous path with a per-rank warning. The stash costs fp32 [batch, sequence, topk] per top-k layer per in-flight microbatch. Validated on GB300: deterministic mode gives bitwise-identical indexer loss and gradients for TP1/CP1 and TP2/CP2+SP; a reduced-size 16-GPU PP4/CP4 DSA model with full recompute, eager and full-iteration CUDA graph, matches the baseline loss within non-deterministic drift and no longer reruns the indexer in the replay. With a frozen indexer (dsa_indexer_loss_coeff <= 0 or unset) the flag reuses indices only: the first call keeps no scores, and the replay of a checkpointed top-k layer takes the indices that its first call published, which the checkpoint snapshot holds, and skips the indexer projections, the dense scorer and the top-k. At 1M with full recompute each top-k layer then runs the dense indexer once per microbatch, and the replay attends exactly the first call's selection. A missing entry falls back to recomputation with the same warning; a key-length or row mismatch, or a checkpoint that runs a layer twice, raises. Config validation requires the sparse indexer loss only when the loss is on, and for a frozen indexer rejects KPool and quantized indexer projections under fp8_recipe 'delayed'/'custom' or fp4_recipe 'custom', whose Transformer Engine recompute stash the replay would never consume. The loss-on path is unchanged. CPU tests cover the frozen config matrix and replay (the coefficient-0 rejection case is now accepted); a GPU test compares eager, recomputing and reusing replays. GOLDEN_CONFIG in test_hybrid_moe_model.py lists dsa_indexer_reuse_checkpoint_topk next to NVIDIA#8287's dsa_indexer_qk_proj_use_quantization. NVIDIA#8105 wraps checkpointed callables in with_current_cuda_device outside the forward-sharing wrapper; that wrapper only sets the device and does not change grad mode, so the first call is still marked FIRST_FORWARD and the replay REPLAY. KPool selection (NVIDIA#8287) forbids the indexer loss, and config rejects dsa_indexer_reuse_checkpoint_topk with dsa_indexer_kpool > 1. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
…ty top-k certificate Two index-sharing improvements for DSA: pipeline layouts whose chunks start on an index-sharing layer are rejected when the config is parsed, and index-sharing layers inherit the producer's non-empty top-k certificate. Reject pipeline layouts whose chunks start on index-sharing layers With dsa_indexer_topk_freq > 1 a DSA skip layer reuses the top-k of the most recent computing layer through the forward-sharing state of one model chunk's forward, so that computing layer has to run earlier in the same (PP rank, VPP stage) chunk. Today this is only enforced when the model is built (the per-chunk spec checks in experimental_attention_variant_module_specs and gpt_layer_specs) or at the first forward, i.e. on the GPU after process group setup, and only on the ranks that own the bad chunk while their peers continue into collectives. TransformerConfig now checks every chunk of a custom pipeline_model_parallel_layout when it is parsed: decoder layers are numbered in the layout's execution order (VPP stage major, then PP rank), MTP layers continue from num_layers + 1, and each chunk's first DSA layer (its first decoder layer, or the first MTP layer of an MTP-only chunk) must compute its own top-k. Because the numbers inside a chunk are consecutive, that is exactly the condition for every later skip layer of the chunk to find its source in the chunk. A violating layout raises a ValueError naming the chunk, the layer and its source on every rank alike, and without a GPU (the CPU argument parse of a launch already fails). For example, a 78-layer model with 7 repeated MTP layers, offset 3 and frequency 4 runs the interleaved layout Et*10|t*12|t*12|t*12|t*12|t*12|t*8|m*7L (chunk starts 1, 11, 23, 35, 47, 59, 71 and MTP 79 all compute top-k) or Et*22|t*24|t*24|t*8m*7L; with the 4-layer granularity of the rule, a one-off layout edit easily lands a chunk on a skip layer. Models without DSA, or without cross-layer sharing, are unchanged. Pipeline splits without a custom pipeline_model_parallel_layout (even splits, or num_layers_in_first/last_pipeline_stage) are not checked at parse time; they keep the model-build checks. Tests (CPU): two PP4 layouts (18/24/24/12 and 22/24/24/8 decoder layers) and two VPP2 layouts pass and match an independent chunk-start reference; PP, VPP and MTP-only chunks that start on a skip layer are rejected with the chunk named; empty and embedding-only chunks are skipped; TransformerConfig accepts a VPP2 layout with an MTP-only chunk (VPP 2, MTP not standalone) and rejects a bad one, and does not check non-DSA or topk_freq=1 configs. The MTP user guide states the rule next to the IndexShare paragraph. The user-guide sentence extends the repeated-layer IndexShare paragraph of docs/user-guide/features/multi_token_prediction.md that the repeated-MTP sharing change (NVIDIA#6472) adds, so that one docs hunk needs the repeated-MTP sharing change; the code and tests apply to dev without it. Share the non-empty top-k certificate with index-sharing layers A top-k-computing DSA layer can prove that every query row selects at least one key and then skips the empty-row masking around sparse attention. Layers that reuse its indices through cross-layer index sharing or repeated-MTP sharing could not, so they always ran the dO masked_fill clone and the in-place dq masked_fill in the sparse-attention backward (on GB300, about 0.6 ms per skip layer per microbatch at 4K/128K and 2.4 ms at 1M). Publish the producer's certificate next to its indices in both sharing payloads (always written, False included; copied in __copy__; popped with the MTP indices) and let consumers combine it with their own layout checks. A missing certificate means False. Consumers receive the producer's tensors, so the skipped masking is a value no-op and results are bitwise identical. The frozen-indexer replay of dsa_indexer_reuse_checkpoint_topk, which reuses its own first call's indices, likewise takes the certificate published with them (CPU test for a certified and an uncertified layout). The GPU test of the shared certificate runs in non-deterministic mode here; its deterministic case needs the deterministic cuDNN DSA mode and comes with it. A KPool producer without dsa_indexer_kpool_always_select_tail certifies nothing (see the preceding KPool fix) and publishes False, so its consumers inherit no certificate either. The KPool CPU test of that fix now runs on the small index-sharing config of the repeated-MTP sharing tests (layer 3 computes top-k, layer 4 reuses it) and also runs the producer's skip layer with and without the tail. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
…ts (PR 8534) Squash of NVIDIA#8534 (head 8c7cba6, same patch-id). cuDNN frontend 1.26 (NVIDIA/cudnn-frontend#316) changed two defaults that the DSA cuDNN call sites rely on. Sparse indexer backward: indexer_backward_wrapper treats top-k ids as per-batch local unless topk_indices_global=True. MCore's sparse-loss path has always materialized global B * S_k + local ids, so BSHD calls with more than one sequence per micro-batch lost or misplaced the sparse indexer gradient of every batch after the first (THD, one sequence per batch, is unaffected). Call the wrapper through _call_sparse_indexer_backward, which sets topk_indices_global=True whenever the installed wrapper names the parameter (or takes **kwargs) and keeps the historical call otherwise. Dense causal offsets: the dense indexer forward, dense attention score recompute and dense indexer backward kernels gained q_causal_offsets and, when it is not given, use a zero (top-left) q_causal_offset instead of the bottom-right alignment (seqlen_k * ratio - seqlen_q): FE <= 1.25 aligns bottom-right, FE >= 1.26 top-left. MCore's BSHD call sites pass no offsets and were written against the bottom-right contract. The single-sequence packed CP scorer in particular crops the key prefix of each query chunk so that bottom-right masking reproduces each row's absolute causal length; on a top-left frontend the same call masks top-left and the top-k reads -inf columns inside the row's length, silently changing the selection. Pass the bottom-right offset explicitly when the installed frontend is 1.26 or newer and the offset is non-zero. Older frontends and square (zero-offset) calls keep their exact call and kernel compile key. A top-left frontend whose wrapper lacks q_causal_offsets fails closed. Add CPU fakes of both frontend generations (dsa_cudnn_fe_fakes.py: BottomRightDefaultDSA for FE <= 1.25, Fe131DSA) and tests that the cropped segment scorer selects identical top-k on either default; the offset tests cover the 1.25 / 1.26 boundary. The native parity tests install fakes that model the frontend <= 1.25 wrappers, so while a fake namespace is installed they now report frontend 1.21 instead of the installed cudnn version; on a host with FE >= 1.26 the offset gate would otherwise treat those fakes as a top-left frontend and fail four tests. The real namespace keeps reporting the installed version, so the real-kernel tests still exercise the top-left path. The fake indexer_backward_wrapper of the five sparse-loss tests in test_dsa_native_parity.py (selected scores, masked query rows, TP reduction, batch-major ids for batched KV, sequence chunking) now takes a required topk_indices_global argument and asserts it is True, so they fail without the sparse-backward fix; the batched-KV test is the BSHD micro-batch > 1 case. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
Propagate deterministic mode through the split DSA hooks, canonicalize top-k cutoff ties, and use fixed-order indexer and sparse-attention backward implementations. Reuse the cuDNN deterministic backward workspace across eager and graph execution, with focused parity and capture coverage. The canonical top-k membership may select every score except NaN and -inf (_selectable_topk_scores). +inf is the largest score, not a masked one (the indexer mask value is -inf), so +inf members count toward the cutoff, an +inf cutoff keeps the lowest-index +inf copies, and -inf is never selected. Counting only finite raw members and candidates would return a row holding +inf scores short by the number of +inf members the radix kernel ranked into its top K (a 64-column row with five +inf scores and K=16 would return 11 members), and a row with more than K +inf scores empty, while the non-deterministic path keeps them. A NaN is never a member, and a slot the raw kernel spent on a NaN is not refilled. CPU tests check rows with +inf, NaN, both, more +inf than K, short rows and -inf, and random rows of +-0, +-inf, NaN and ties against an independent sort reference for both NaN placements; all five fail when +inf is not selectable. This commit stacks on the preceding "adapt cuDNN DSA call sites to cuDNN frontend >= 1.26 defaults" fix: the deterministic sparse indexer backward reaches cuDNN through its _call_sparse_indexer_backward with global top-k ids, _supports_deterministic_sparse_attention_backward uses the inspect import it adds, the new test_cudnn_deterministic_sparse_loss_canonicalizes_topk_order expects topk_indices_global=True like that fix's sparse-loss tests, and both change adjacent lines of dsa_cudnn_kernels.py. A cut onto a base without that fix has to carry these pieces. Upstream NVIDIA#7441 (7a39a74) is the CSA counterpart in csa_utils/fused_sparse_attention.py: it defines module-private helpers with the same names (_DETERMINISTIC_SPARSE_BWD_HEAD_COUNTS, _DETERMINISTIC_INDEXER_DK_CHUNK_MAX_BYTES, _get_deterministic_sparse_attention_workspace, _deterministic_indexer_dk_chunk_rows, ...) with CSA-specific bodies and budgets. These deterministic helpers stay module-private on both sides; dev's CSA module imports only four padding/alignment helpers from dsa_cudnn_kernels (_get_head_padding, _get_score_recompute_topk_alignment, _pad_attn_target_heads, _pad_sparse_backward_topk), which this stack leaves unchanged. Both gate on config.deterministic_mode or torch.are_deterministic_algorithms_enabled(). dsa_kernels filters hook kwargs with megatron.core.utils.filter_kwargs_for_callable, whose signature cache test_filter_kwargs_for_callable_inspects_signature_once in tests/unit_tests/test_utils.py already covers, so no separate signature-inspection test is added. The GPU tests of the shared non-empty top-k certificate (test_dsa_topk_nonempty_certificate.py) and of the checkpoint top-k reuse (test_dsa_checkpoint_topk_reuse.py) gain their deterministic cases. They require bitwise-identical loss and gradients with and without the shared certificate or the reused top-k; the certificate test skips when the installed frontend has no deterministic sparse-attention backward. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
Resolve cuDNN DSA entry points once cudnn.DSA imports the kernel modules behind an entry point on first attribute access. Under CUTLASS DSL 4.6.2 every @cute.kernel decoration leaves a frame that references itself (CuTeDSL.kernel keeps inspect.currentframe() in a local), and its f_back chain runs through the import into the frame that touched the entry point. MCore first touches sparse_attention_backward_wrapper and indexer_backward_wrapper inside the first backward of a training step, so those imports pinned the activations and gradients on the Python stack at that moment until the next cyclic collection. With --manual-gc that is the collection before full-iteration CUDA graph capture: on a 16-GPU scaled-down GLM-5.2 model (PP4, CP4, 16K tokens) the FE 1.31 stack held 1349 / 1397 / 0 / 3628 MB per PP stage through the eager warmup steps, all of it freed by gc.collect(). DSL 4.4.2 takes inspect.currentframe().f_back without keeping its own frame, so the FE 1.21.1 stack held nothing. Resolve every entry point MCore calls in _ensure_dsa_namespace and run one gc.collect() before it returns (skipped during stream capture), so no later call imports kernel modules and no caller frame outlives the import. Entry points a frontend lacks are left to their call sites, which already probe for them. CPU tests model the lazy namespace with a DSL-4.6-style decorator: the model pins its caller without a collection, the namespace load does not, every entry point is resolved exactly once, a missing one is tolerated, and the entry point list covers every namespace access in the module. Size top-k chunks by scratch elements and preallocate outputs cuDNN radix top-k allocates an int32 scratch of rows x 2 x num_cols per launch, and _indexer_top_k_wrapper_chunked sizes its row chunks by a scratch byte budget. The budget stays 2 GiB, 256 rows per launch at 1M keys; _topk_wrapper_chunk_rows takes it as an optional argument that defaults to _TOPK_WRAPPER_MAX_SCRATCH_BYTES. Do not add a row cap. The frontend limit is on elements (one launch may index at most 2**31 scratch elements), so cap rows by that element count as well, which binds for budgets above 8 GiB. At 1M keys an 8 GiB budget is exactly 2**31 elements, which takes FE 1.31's own split branch with the same single buffer; 2**31 - 1 would add a 17th launch. Replace the 512-row chunk alignment, which floors every budget above 512 rows to a 512-row multiple (with an 8 GiB budget, the 1524 rows of the 1M CP64 rank-42 width became 1024), with the smallest row multiple that keeps every chunk start 32-byte aligned (the compiled kernel assumes a 32-byte aligned base pointer): 1 when a row is a 32-byte multiple, at most 8 for FP32 scores. A budget that fits less than one such row group still launches one group (at most 8 FP32 or 16 BF16 rows); a smaller chunk would start every later chunk off a 32-byte boundary. Write each chunk's result into [rows, K] outputs allocated before the first launch and drop the chunk before the next launch, refusing a chunk whose dtype or size does not match instead of casting it. The frontend allocates a launch's outputs before its scratch, so with the former collect-then-torch.cat every later launch's outputs were carved out of the scratch block just freed and needed a second scratch-sized segment, whatever the budget: with an 8 GiB budget, two 8 GiB blocks in the CUDA graph pool at 1M (16.02 GiB reserved against an 8.25 GiB allocated peak). With preallocated outputs the pool holds one block (8.27 GiB) eagerly and in a graph. The CSA paths call the frontend directly with their own sizing and selection (on dev the CSA dense fallback row-chunks under its own 1 GiB score budget, NVIDIA#7811), so the budget does not apply there; the wrapper docstring says so. CPU tests cover the chunk sizes at the 2 GiB default and at 8 GiB (1M keys and the CP64 rank 21/42/63 widths of a 1M-token sequence, 128K), the int32 element limit under a 64 GiB budget, a fractional budget, the one-row minimum, the one-group minimum below a row group's scratch (FP32 and BF16), the alignment rule, and budget-sized launches with 32-byte aligned chunk starts at both budgets. The deterministic chunking test in test_dsa_native_parity.py now scores 8-column (32-byte) rows, so its one-row chunks stay aligned. The real-frontend GPU checks of the 1M launches and the scratch reservation live in the deterministic top-k GPU suite added with the tie_break path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
cuDNN frontend fills the whole dense indexer score buffer with -inf before every scorer launch. At 1M-token context on GB300 (CP64, packed THD) that is a 64 GiB fill per call on CP rank 0: 8.77 ms per call (32 FillFunctor launches at about 7.85 TB/s), about 0.84 s per step on the bottleneck pipeline stage (96 calls) and 0.98 s on the first stage; at 128K it is about 53 ms per step. Yet on the general packed THD path every reader only consumes the columns the kernel writes, and FE 1.31 writes exactly [0, causal_len) per row. Probe indexer_forward_wrapper for an explicitly named clean_logits parameter (the upstream proposal, NVIDIA/cudnn-frontend#1761; no released frontend has it yet) and require frontend >= 1.29, whose varlen radix top-k carries the out-of-bounds-lane guard (NVIDIA/cudnn-frontend#410) that keeps degenerate rows from re-reading lanes past the row. When both hold and the new dsa_indexer_skip_masked_score_fill option is set (default; disable with --no-dsa-indexer-skip-masked-score-fill for A/B), only _indexer_topk_packed_thd passes clean_logits=False. Every other scorer keeps the prefill: the dense indexer loss reads whole rows and the BSHD paths rely on it for masking. Frontends without the switch keep the prefill and log it once. Without the prefill the scorer derives each row's kernel column limit from the same segment layout it hands to cuDNN, entirely on the device, and fails closed: top-k lengths never exceed that limit, rows at or after the last cu_seqlens entry (never written by the kernel) get length zero, and the full-cover shortcut drops columns past the limit before the caller's bounds and finiteness mask. Scorer opt-ins now travel from the config to the packed THD executor in a frozen IndexerScorerOptions value threaded through the split and fused hooks; run_fused_qk_topk receives the config through the existing signature-filtered hook kwargs. run_fused_qk_topk_with_scores, which the checkpoint top-k reuse of the recomputed indexer loss added below this commit, passes the same options. CPU tests (with an Fe131CleanLogitsDSA fake that poisons the unwritten columns with NaN, +inf or large finite values) require bitwise-identical top-k indices, lengths and scores against the prefilled scorer for CP1, CP2 and CP4 layouts, padded cu_seqlens and the full-cover shortcut; they also cover the probe, the uncovered-row clamp and the absence of host synchronization. A GPU test runs a DSAttention layer with a 64-head DSA attention geometry (512+64 absorbed latent, 32x128 indexer, sparse indexer loss) through forward and backward on the real cuDNN frontend and FlashMLA kernels with the option off and on: packed THD multi-sequence packs with a dummy tail and fixed-capacity cu_seqlens at CP1 (ranked and full-cover top-k), zigzag CP4, and dynamic-CP static slots on CP2/CP1 subgroups of a CP4 world, with NaN and +inf poison in a caller-owned score buffer. In deterministic mode top-k, output, indexer loss and every gradient must be bitwise identical to the prefilled run. In non-deterministic mode (unordered radix top-k slots, atomic sparse backward) the per-row top-k sets with their scores and the lengths must be bitwise equal, and output, loss and gradients must stay within a relative L2 error of 1e-2 and a peak error of 2e-2 (1e-3 for the fp32 loss), limits that a prefilled rerun also has to meet. The new field defaults to True, so GOLDEN_CONFIG in test_hybrid_moe_model.py lists dsa_indexer_skip_masked_score_fill: True. KPool selection (NVIDIA#8287, dsa_indexer_kpool > 1) bypasses the cuDNN scorer, so the option has no effect there. NVIDIA#8287's weights cast in run_fused_qk_topk_with_loss is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
Deterministic mode rebuilt the cutoff-tie membership from full-row masks, cumsums and searchsorted over every score column in 16- or 31-row chunks. A 1M-token CP64 rank-0 call took 2221 ms that way (CUDA graph, 529 launches; 2238 ms eager) against 14.26 ms for the non-deterministic top-k with an 8 GiB scratch budget, which made deterministic runs unusable at 1M. cuDNN frontend 1.31 top-k takes tie_break=1, which keeps every radix key above the cutoff key plus the lowest-index copies of the cutoff key. The canonical rule depends on the raw result only through its selectable in-row members (every score except NaN and -inf), whose count and minimum do not depend on how ties were broken, so the tie_break=1 members already are the canonical set, and sorting them by index gives the canonical order. The one exception is a zero cutoff: radix keys rank +0.0 above -0.0 while the rule compares values. Those rows replace their zeros with the lowest-index in-row zeros of either sign through the new select_lowest_index_zero_ties, which reads only those rows: on CUDA a Triton kernel scans just the rows with a positive quota and stops once it is met, other devices use a fixed-shape PyTorch implementation. Values come from the score buffer, so -0.0 survives; padding stays -1 / finfo.min; NaN, -inf and +inf are handled exactly as in the full-row rule (_selectable_topk_scores): +inf is selected, -inf never, and a slot the raw kernel spent on a NaN is not refilled. The path runs at the non-deterministic chunk size, keeps fixed shapes and needs no host synchronization (CUDA Graph capturable). A capability probe requires an indexer_top_k_wrapper that names tie_break on frontend >= 1.31; otherwise (older frontends, **kwargs wrappers) the full-row canonicalization and its repair-sized chunks stay in use, now rounded to the same 32-byte row alignment as the other top-k chunks so that later chunk starts keep the frontend's assumed input alignment; when the repair budget fits less than one aligned row group (above about 4.2M FP32 columns), a launch takes one group. Non-deterministic calls never pass tie_break, so their compile keys are unchanged. Deterministic mode logs once per distinct reason which path runs: a warning for the full-row fallback (naming the frontend version and whether the wrapper names tie_break), a warning when Triton is missing on CUDA, and an info line otherwise. Measured through _indexer_top_k_wrapper_chunked on GB300 with FE 1.31 (CUDA graph, 8 GiB top-k scratch budget): - 1M CP64 rank 0: 2221 -> 18.8 ms on Gaussian scores; 27.1 ms on a 1/64-quantized ties set with 661 signed-zero cutoffs, where FE's own tie_break=1 launch accounts for most of the extra cost (21.90 vs 14.30 ms; repair + sort + gather add 4.6 to 5.2 ms). - 128K CP32 rank 0: 22.85 -> 1.42 ms (1.72 ms on ties). - Bitwise equal to the full-row rule on both shapes and both data sets (0 mismatching rows); graph replay equals eager and repeats are bitwise equal. CPU tests model the radix kernel (bit-pattern keys, random or lowest-index ties, shuffled output, two NaN placements) and require bitwise equality with the full-row rule on random, heavy-tie, signed-zero, +-denormal, DSA relu-head and NaN/inf rows, short rows, chunked launches, and the probe, fallback, alignment (including the one-group minimum below the repair budget) and logging. They also check that every chunk result of the preallocated-output loop is dead when the next launch starts on the non-deterministic, tie_break and full-row paths, and that a chunk dtype mismatch is rejected. The zero-tie selector is compared with a brute-force reference, the Triton kernel under TRITON_INTERPRET=1. The signed-zero cases with NaN and +inf expect the +inf member, plus an all-+inf case, and the +inf / NaN tests against the independent sort reference now run both canonicalizations. GPU tests on the installed frontend: with tie_break the wrapper must equal the full-row rule bitwise on Gaussian, heavy-tie, signed-zero, all-zero, NaN/infinity, short and garbage-tailed rows at widths that take the 16-bit, 32-bit and global-spill kernel paths, match a stable-sort reference for lowest-index ties, repeat bitwise, replay from a CUDA graph, and reach cuDNN as tie_break=1 with return_val=False (full-row path tie_break=0, non-deterministic path return_val=True); +-denormals never fill zero-tie slots, and a 2048 x 1M case with signed-zero cutoffs runs the Triton repair at full width. Without tie_break (FE < 1.31) the probe must report no support and the wrapper fall back. On both frontends, with an 8 GiB scratch budget set by the tests (the default stays 2 GiB), 2048 rows of 1M columns must launch as two 1024-row chunks, select the masked torch.topk set, peak within the budget, and reserve one scratch block eagerly and in a graph pool. The NaN/inf rows (+inf x5, NaN, both, more +inf than K, short rows, -inf) run at 64 x 70000 and 16 x 1M against the full-row rule and a sort reference; on GB300 with FE 1.31 every +inf row is complete, every NaN row is exactly one member short, and the whole deterministic GPU file passed with an 8 GiB budget in every test (47 passed, 1 skipped: the pre-1.31 fallback check). The tests that do not set the budget now run at the 2 GiB default and have not been rerun on a GPU since. The determinism kernel manifest (NVIDIA#7148) gets an entry for dsa_topk_zero_ties.py with its exact-reference and bitwise-repeat tests. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
Experimental and opt-in: dsa_indexer_precision stays 'bf16' by default, and nothing changes unless a run sets it to 'mxfp8'. Only the DeepGEMM selection scorer later in this series builds on its code (the MXFP8 option, the cuDNN MXFP8 fallback and the loss-score rule below). Measured trade-offs on GB300: - Speed: at the GLM-5.2 1M CP64 per-call shapes the MXFP8 scorer kernel saves 15.2-15.6 ms per call over FE 1.31 BF16 (15.1-15.5 ms net of the TE quantization), about 0.7 s per 1M step with dsa_indexer_reuse_checkpoint_topk (48 scorer calls on the bottleneck stage) and about 1.4 s without it. Most of the saving over the FE 1.21.1 BF16 scorer (27.3-28.0 ms) comes from the FE 1.31 BF16 upgrade. - Accuracy (not yet validated for training): recall@2048 against the BF16 scorer is 0.95-0.96 mean (p1 0.93-0.94) on Gaussian synthetic inputs and 0.92-0.95 (min 0.70) on outlier-channel inputs without a Hadamard rotation; outlier inputs with a Hadamard rotation reach 0.986-0.990. The kernel itself is exact on the dequantized inputs, so the loss comes from quantization. The recall data in this commit is synthetic; recall on real model activations and a loss A/B are required before enabling it for training. What it does: dsa_indexer_precision was read only by the DSv4 compact CSA indexer and validated only for dsv4_hybrid, so setting it to 'mxfp8' for experimental_attention_variant='dsa' was silently ignored. cuDNN frontend 1.29+ provides an MXFP8 dense indexer scorer (E4M3 Q/K with E8M0 1x32 scales, BF16 weights, 32 or 64 heads, THD with q_causal_offsets). Make the option work for 'dsa' with dsa_kernel_backend='cudnn', or fail: - Config validation requires the cuDNN backend, head_dim 128, 32 or 64 indexer heads, BF16 indexer weights, the sparse indexer loss and the fused DSA kernels (attention_backend other than unfused). The SM100+ device, the frontend's MXFP8 scorer API and Transformer Engine MXFP8 plus Triton are checked at the first scorer call with one error listing every missing piece. - The general packed THD scorer quantizes Q and the score-K (after the CP layout index_select) with indexer_quantization and calls indexer_forward_wrapper(precision='mxfp8') with per-segment padded scale prefixes. Prefixes are built on the device and the scale storage is sized from tensor shapes alone, bounded for any segment boundaries with the same token and segment counts, so nothing reads a CUDA tensor on the host or allocates a data-dependent shape under full-iteration CUDA graph capture; all storage is dispatch-local. Packs that hold a single sequence, which eager dynamic-CP runs produce, are routed to the same general scorer when their packed metadata passes its structural checks. - Every other scorer path (plain causal, single-sequence segments without packed metadata, unfused bounds, dense indexer loss) raises instead of silently scoring in BF16, and so do the reference fallbacks of a computing DSAttention layer (fused kernels disabled, no fused bounds for the mask layout, or a declined hook) and DSAIndexer.forward_with_scores, so an A/B run can never report BF16 numerics or throughput as MXFP8. Indexer loss: MXFP8 decides only which keys are selected. Weights stay BF16, the indexer loss and backward keep the BF16 Q/K, and the selected scores that feed the sparse KL (and that dsa_indexer_reuse_checkpoint_topk keeps for the checkpoint replay) are recomputed in FP32 from the BF16 indexer Q/K/weights by cuDNN FE sparse_indexer_score_recompute_wrapper (in FE 1.21.1 and 1.31; added to _CUDNN_DSA_ENTRY_POINTS and the protocol) instead of being the MXFP8 scorer's own values: - Rule (_recomputes_indexer_loss_scores): the selected scores are recomputed if and only if the scorer computes in MXFP8 (dsa_indexer_precision='mxfp8'), whatever kernel scores and whatever dtype its scores have. 'bf16' keeps the dense scorer's FP32 values byte for byte (_selected_scores_for_indexer_loss hands back the same tensors), and a future scorer computing in BF16 with BF16 output would keep its values too. Frozen-indexer and selection-only calls keep only indices and never recompute: the loss is on when the caller asks for scores (return_topk_scores), not when a kernel returns values, and a loss-on call without values raises. dsv4_hybrid's compact CSA indexer keeps its own loss semantics (NVIDIA#5992); only the DSA path changes here. - Where: _indexer_topk_packed_thd scores and selects in _indexer_topk_packed_thd_cudnn and has a single exit, so the rule (_selected_scores_for_indexer_loss) covers the ranked path and the every-column path (rows no wider than the top-k, no ranking) alike. - Payload: the selected global key ids, padded with -1 to the top-k rounded up to 128 (one kernel shape per configured top-k) and passed without topk_length, so the kernel masks -1 itself and the valid slots need not be a prefix. The result is log(max(p, FLT_MIN)), -inf at -1: the BF16-input logits minus the row's log-sum-exp over the selection, FP32 [b, sq, K] as before, so the score stash and its replay are unchanged. Every consumer applies masked_log_softmax over a subset of the selection, which is shift invariant, so given the selection the KL gradient equals the BF16 path's (the backward clamps p to FLT_MIN anyway). The KL value does too, except where p < FLT_MIN (a logit more than 87.3 below the row's log-sum-exp over the selection, which exceeds its maximum by up to ln |S|: from about 79.7 below the maximum at top-k 2048) and the attention target has mass on that slot: that term saturates near 87.3 nats while the BF16 path's keeps growing (logits [495.06, 0], target [0.5, 0.5]: 42.98 against 246.84). A loss A/B reports both arms' KL with the same clamp. - Selections the caller drops stay dropped: a slot whose MXFP8 value is not finite (a -inf prefill column, a NaN score), that lies outside its row's [starts, ends) or past the keys becomes -1 before the recompute. - What changes for this scorer: the KL predict used to be the MXFP8 values. A CPU emulation on GLM-5.2 indexer inputs (16K and 1M rows, an untrained indexer) puts the MXFP8 logit error on the selected keys at 10-17% of their spread and the predict TV against BF16 at about 0.015 per row: a fixed, input-determined bias of ~0.03 * coeff (L1) on the KL gradient that does not shrink as the indexer converges. Now MXFP8 changes only the selection, which is also what a loss A/B against BF16 isolates. - Cost [inferred, to be measured]: one recompute per loss-on scorer call. At the 1M CP64 rank-0 call it gathers ~8 GiB of K rows (partly from L2) and 0.28 TFLOP: ~1.5-3 ms per call against the scorer's 15.2-15.6 ms saving. At 128K CP32 rank 0, 2 GiB and 69 GFLOP with K L2-resident: ~0.1-0.5 ms per call. The BF16-compute paths (dsa_indexer_precision 'bf16') are unaffected. - A frontend without the entry point raises with the loss on rather than silently changing the loss. - CUDA graphs: the recompute adds no host sync, and its kernel compiles at the first eager call per top-k and head count, like the scorer and the top-k, so a full-iteration capture after the eager warm-up steps replays it. CPU tests cover config validation, scale storage bounds (random layouts and 4K, 128K CP32 and 1M CP64 packs), argument plumbing at CP1 and CP2 through a mocked frontend and quantizer (fake_quantize_indexer_mxfp8), single-sequence routing, composition with the skipped prefill, the explicit errors and the absence of host syncs (scoring, selection and the recompute, ranked and every-column). The packed THD layout helper of the clean_logits tests moves into the shared fakes module, which also gets a torch reference of the FE sparse indexer score recompute (Fe131DSA.sparse_indexer_score_recompute_wrapper, calls recorded). The loss-score rule is pinned by a precision {bf16, mxfp8} x loss {on, frozen} x {ranked, every-column} matrix (which path supplies the scores, recompute calls); MXFP8 selection with the BF16 recompute payload at CP1 and CP2; bf16 selections passed through as the same tensors, and indices-only calls keeping their indices and no scores (no recompute) even when a kernel returns values; the MXFP8 payload giving the BF16 path's KL loss and backward predict for the same selection (ternary inputs through both real paths, Gaussian ones against the BF16 path's scores of those keys); the recompute's padding to one shape, masking and FLT_MIN floor; a missing entry point; dropped selections (non-finite, below or past the row, past the keys inside a row whose bounds exceed them) never reaching the recompute; prefilled uncovered rows; and the checkpoint score stash holding the recompute and the replay consuming it. test_dsa_cudnn_mxfp8_scorer_gpu.py adds GB300 checks of top-k recall, the MXFP8 scorer's own selected values against BF16's (read with the recompute switched off), the loss payload being the BF16 logits shifted per row, sparse indexer loss and gradients against BF16 for CP1, CP2 and CP4 layouts, the recompute against log-softmax of the FE BF16 scores (CP1 pack and 128K CP32 rank 5 with 32 and 64 heads, plus an underflowing key) within an FP32 summation-order error bound computed per element from the inputs, the MXFP8 loss-on call (score, select, recompute) under set_sync_debug_mode('error') and its CUDA graph replay against eager after a data swap (32 and 64 heads), the MXFP8 path's loss and gradients equal to the BF16 path's given the same selection, plus hardware checks of the explicit bottom-right offsets and of the skipped prefill; the latter compares the prefilled and the skipped-fill top-k in deterministic mode, because the radix top-k writes its selection in an unspecified slot order otherwise. The recompute bound (an error model, not a fixed tolerance): both kernels compute s = sum_h w_h relu(q_h . k) in FP32 from BF16 inputs (products exact) in their own order. A term passes through at most n = 128 + H + 2 FP32 operations, each with relative error below u = 2^-23 (rounded or truncated, as tensor-core accumulation is), so each kernel's logit is within gamma_n * A of the exact one, A = sum_h |w_h| sum_t |q_ht k_t| (Higham, Lemma 3.1 and (3.5)); log p = s - LSE with the LSE 1-Lipschitz in the max norm, plus the log-softmax's own rounding (gamma_K + 4u (M + ln K)). The two kernels then differ per element by at most 2 gamma_n (A_i + max_j A_j) + 2 (gamma_K + 4u (M + ln K)), 0.02-0.07 on these inputs. The first GB300 run measured 1.6e-5-1.9e-5 against the earlier fixed 1e-5, uniform across heads and layouts (FP32 order noise; two FP32 orders on CPU give 0.8e-5-1.9e-5). The test prints the worst element's |log p|, logit, A and bound, and asserts as a negative control that a logit from MXFP8-cast q and k (the precision the recompute exists to avoid) exceeds the bound (26-62x in CPU emulations). IndexerScorerOptions gains only precision here and the packed THD scorer call only the MXFP8 kwargs; the code references no top-k backend or tie_break. Upstream NVIDIA#8287 added KPool selection (dsa_indexer_kpool > 1), which DSAttention runs instead of the cuDNN dense scorer. MXFP8 validation needs only the sparse indexer loss, so 'mxfp8' with KPool and a zero loss coefficient (which KPool requires) passed validation and the layer silently scored with the KPool selector, in BF16 or, with dsa_indexer_kpool_use_quantization, in E4M3 with power-of-two scales. Config validation now rejects 'mxfp8' with dsa_indexer_kpool > 1, and the KPool branch of DSAttention raises for a config changed after validation, as the other non-MXFP8 fallbacks do. NVIDIA#8287's other knobs are independent of this option: dsa_indexer_qk_proj_use_quantization picks the context of the indexer Q/K projection GEMMs, whose BF16 outputs the MXFP8 scorer quantizes itself, and NVIDIA#8287's weights.to(q.dtype) cast is a no-op for the BF16 weights MXFP8 requires. NVIDIA#8104's dsa_compact_indexer_workspace_sharing is CSA only; the DSA MXFP8 scale storage is allocated per scorer call. The dsa_indexer_precision docstring says all of this. New CPU tests cover the rejection at loss coefficients 0 and 0.01, an unchanged BF16 KPool layer, and the runtime guard in train and eval. The legacy-frontend case of the MXFP8 runtime-requirement test uses BottomRightDefaultDSA, the FE <= 1.25 fake of the cuDNN frontend >= 1.26 fix. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
… scores
Experimental and opt-in: dsa_indexer_topk_backend stays 'cudnn' by
default, and the DeepSelect builds have run on GB300 (SM 10.3) only.
New config field dsa_indexer_topk_backend ('cudnn' default | 'deepselect',
validated in TransformerConfig; 'deepselect' requires
experimental_attention_variant='dsa' and dsa_kernel_backend='cudnn').
With 'deepselect', _indexer_top_k_wrapper_chunked runs an eligible
non-deterministic top-k as one DeepSelect launch over all rows instead of
the scratch-bounded, row-chunked cuDNN radix top-k. Every top-k caller
gets the backend from the config's scorer options (packed THD general,
packed segments, score chunks and the dense BSHD path). Deterministic
mode, including torch.are_deterministic_algorithms_enabled(), always
stays on the cuDNN tie_break=1 path. FP32 scores are described first;
the last section extends the backend to BF16 scores.
Measurement setup: every cuDNN top-k comparison below was measured on
GB300 with an 8 GiB cuDNN top-k scratch budget. The default budget is
2 GiB (_TOPK_WRAPPER_MAX_SCRATCH_BYTES; configurable as
dsa_cudnn_topk_scratch_gib later in this series), which runs the cuDNN
top-k in smaller row chunks at wide rows.
Why: in a comparison of top-k kernels at GLM-5.2 1M-token CP64 shapes on
GB300, the patched DeepSelect build passed every correctness check and
took 5.92 ms against 14.30 ms for the 1M CP64 rank-0 call, bitwise
deterministic run to run and batch-independent, with no scratch
(0.002 GiB extra memory against 8.03 GiB).
New module dsa_deepselect_topk:
- Builds: the tuned ds_tune package (DeepSelect bfa4507 + 3 patches,
config f_t512_b2_2k_d3 for FP32 K <= 2048) is used if it imports, else
upstream deep_select. Both APIs are probed by named
arguments once; a missing build is a logged fallback, not an error.
- Eligibility, from shapes, strides and the start address only (no host
sync): SM 10.3, FP32 scores, stride(1) == 1, a 16 B-aligned start (the
TMA descriptor's global address; the build checks only the stride) and
a row stride that is a multiple of the build's input alignment
(1024 B), width below 2**23, a compiled config
for K <= 4096, K a multiple of 8 (contiguous preallocated outputs), K
below the width, and a per-build minimum width: for ds_tune 65536 up to
4096 rows and 81920 above (see "Row-dependent width threshold" below),
524288 for deep_select, where each was measured faster than cuDNN. Any
other call runs cuDNN; the choice is logged once per reason (a warning
when the backend is unavailable).
- SM 10.0: the builds ship sm_100a kernels (COMPILED_SM), but only GB300
has run them, so SUPPORTED_SM is {10.3} and a GB200 keeps cuDNN with a
"not validated" warning. test_dsa_deepselect_topk_gpu.py runs on every
COMPILED_SM device, so one passing GB200 run is what enabling SM 10.0
waits for.
- Contract equal to the cuDNN wrapper's: int32 [rows, K] indices in
unspecified slot order, FP32 values or None, rows no longer than K
gathered first and padded with -1 / -inf (end = per-row lengths,
idx_oob_fill_value=-1, value_oob_fill_value=-inf, sorted_index=False;
callers keep their own compaction and index sort). Only which equal
scores fill the last slots may differ, which cuDNN leaves unspecified.
- NaN: abort_when_nan_found=False, and a Triton kernel (PyTorch fallback
off CUDA) rewrites rows flagged with 0x3F3F3F3F to their first K in-row
columns with NaN values. This reaches
the sparse attention too: with return_topk_scores=True,
_indexer_topk_bshd keeps only in-bounds finite selections, so a flagged
row gets topk_length 0 (no key); without it the row attends to its
first K columns. cuDNN spends one slot on a (positive) NaN and keeps
K - 1 keys. Both are garbage-in rows; the module and config docstrings
state the difference.
- CUDA graph safe: fixed shapes, preallocated index output, no scratch,
no host sync.
Score stride: the cuDNN THD scorer allocates [rows, max_seqlen_k] (rounded
to 4 elements), which is already 1024 B-aligned for every 1M CP64 rank
width (multiples of 8192 elements) and 128K rank width (multiples of
2048 at CP32, and of 1024 at CP64). For other widths
DeepSelect would take, the packed THD path passes the scorer an out= view
of a row-padded buffer: frontend 1.31 binds it directly (a GPU test checks
the returned pointer and that the padding columns are untouched, with and
without clean_logits), so no copy is added. Widths that are not a multiple
of 4 (which the frontend would stage and copy) and older frontends keep
the default buffer and cuDNN.
Measured on GB300 (FE 1.31, through _indexer_top_k_wrapper_chunked, CUDA
graph median, K=2048, GLM-5.2 row lengths), cuDNN -> ds_tune:
- 1M CP64 rank 0 (16384 x 1048576): 14.53 -> 5.94 ms (0.41x); upstream
deep_select 7.94 ms; graph pool 8.28 -> 0.26 GiB.
- 128K (4096 x 131072): 0.564 -> 0.436 ms; deep_select 0.598 ms (slower).
- 64K: 0.322 -> 0.315 ms; 32K: 0.190 vs 0.228 ms; 16K: 0.124 vs 0.160 ms.
- On a 40%-zeros tie-heavy 1M batch: 57.2 -> 7.0 ms.
- End to end, _indexer_topk_bshd on the 1M CP64 rank-0 packed THD layout
(scorer + top-k, eager): 74.55 -> 62.43 ms; same topk_length and selected
scores, 8 of 16384 rows pick other members among exact FP32 ties.
Tests: CPU tests with fake ds_tune and deep_select builds cover
eligibility (every condition, per-build widths, 128K CP32 and 1M CP64
rank widths, the
SM 10.0 warning fallback, a misaligned start address), the
import probe, routing and fallback (deterministic, narrow, unaligned,
missing build), one launch regardless of the scratch budget, contract
equivalence with the cuDNN path on Gaussian, tie-heavy, signed-zero and
+-inf rows with short, empty and garbage-tailed rows, NaN rows, the
indices-only call, the Triton repair kernel in the interpreter, no host
sync, log-once, the packed THD padded score buffer, and config
validation. test_dsa_deepselect_topk_gpu.py checks the real builds against
cuDNN at 128K and 1M, NaN repair with untouched other rows, bitwise
repeats and batch independence, graph replay after a data swap, a
row-padded score view, the frontend's direct out= binding and the 1M CP64
rank-0 packed THD call; all passed on GB300 with ds_tune and with
deep_select. It also has a case for long rows with fewer than K finite
scores (masked -inf heads, as varlen starts produce), where the kernel's
-inf tail placeholder can tie with in-row -inf; that case has not run on
a GPU yet (its assertions were checked on CPU with the fakes).
Row-dependent width threshold
A single ds_tune width taken from one 4096-row point (65536 columns,
0.315 vs 0.322 ms) cannot be right for every call: a GB300 width sweep
through the Megatron top-k entry point shows that the break-even width
grows with the row count, so it differs between 4096-row and 16384-row
calls.
Sweep (GB300, FE 1.31, ds_tune at DeepSelect bfa4507):
_indexer_top_k_wrapper_chunked with topk_backend 'cudnn' vs
'deepselect' (threshold lowered for the measurement), K=2048, FP32
N(0,1) scores, NaN past every row end, CUDA graph replay, 5 rounds x
20 replays alternating the backend order, every output checked exactly.
12 widths 8K..128K x rows 2048/4096/16384 x CP-balanced row lengths
(rank 0: half short, half long rows; last rank: all rows long), 70
points. DeepSelect / cuDNN median time:
rows layout 57344 65536 81920 131072
2048 rank 0 1.018 0.942 0.870 0.753
2048 last 1.047 0.964 0.901 0.762
4096 rank 0 1.052 0.976 0.909 0.772
4096 last 1.088 1.007 0.943 0.787
16384 rank 0 1.146 1.088 0.938 0.812
16384 last 1.123 1.079 0.960 0.799
Below 57344 cuDNN is faster at every row count (1.01-1.68x, 1.26-1.50x
at 16K); the ratio varied by at most 0.02 between rounds. Hence ds_tune
applies from 65536 columns for calls of up to 4096 rows and from 81920
columns for larger calls (16384 rows were 1.08-1.09x at 65536).
The 4096-row limit matches a 128K-token GLM-5.2 configuration with
dynamic CP and max_seqlen_per_dp_cp_rank 4096: it realizes CP32, so
each call has 4096 query rows and rank r a width of 131072 - 2048 r, i.e.
67584..131072. With a 2048-row limit the seven ranks narrower than
81920 (25..31) would fall back to cuDNN. Only the band edges were
measured at 4096 rows: 65536 is break-even (0.976 rank-0 layout, 1.007
last-rank layout, the latter about 4 us slower) and 81920 is
0.909/0.943. Interpolating both times linearly, the last-rank layout
reaches 1.0 at about 67K columns and 0.997 at 67584, and the cuDNN
fallback would have cost those ranks up to about 5% of their top-k
time at 79872. No row count between 4096 and 16384 was measured, so
those calls keep the 81920 width.
The GLM-5.2 1M CP64 call (16384 rows, 532480..1048576 columns) is
unaffected: 5.90 vs 14.34 ms at rank 0 and 6.82 vs 13.65 ms at the
narrowest rank (same sweep, graph replay, Gaussian scores; 5.99 vs
12.89 ms on tie-heavy rows, 5.92 vs 14.55 ms on real indexer-score
rows). At 128K
(4096 x 131072) it was 0.434 vs 0.569 ms. The upstream deep_select
width (524288) is not row dependent and is unchanged.
min_score_width(kind, rows) combines MIN_SCORE_WIDTH with
MANY_ROWS_MIN_SCORE_WIDTH ({'ds_tune': (4096, 81920)});
static_fallback_reason, and with it fallback_reason and the packed THD
out= sizing, uses it and reports "score width below 81920 for ds_tune
with more than 4096 rows" as its own log-once reason. The module and
config docstrings carry these numbers. CPU tests cover the rule at the
row limit and at both width edges for ds_tune and deep_select, and check
that every 128K CP32 rank width at 4096 rows uses ds_tune; the test
fixture clears the row rule along with the width.
Layer GPU test
The DSA layer suites score at most a few thousand columns, so with
dsa_indexer_topk_backend='deepselect' every top-k there falls back to
cuDNN ("top-k covers the whole row") and the backend never runs inside a
layer. The standalone top-k file, the packed THD 1M call and benchmarks
do not cover the layer path either, hence a dedicated layer test.
test_dsa_deepselect_layer_gpu.py runs a 64-head DSAttention layer
(512+64 latent, indexer 32x128, top-k 2048, sparse indexer
loss, cuDNN kernels, packed THD with a dummy tail sequence and
fixed-capacity cu_seqlens) at score widths of 81920 to 131072 columns,
where the production width rule sends every top-k to DeepSelect:
cp1-3seq-dummy-unaligned 82000 rows x 82000 columns (row-padded out=)
cp4-3seq-dummy zigzag CP4, 32768 rows x 131072..81920
dcp2of4-slot dynamic CP2 in CP4, 57344 x 114688/86016
dcp1of4-slot-prefill dynamic CP1 in CP4, 81920 x 81920, -inf
prefill kept (the others skip it)
Each layout runs forward and backward in production (non-deterministic)
mode five times on the same layer: cuDNN, cuDNN again, cuDNN with each
row's selected slots shuffled where the top-k returns them with scores
(the sparse-loss path), DeepSelect, and DeepSelect again. A spy on
dsa_deepselect_topk checks that every top-k of a DeepSelect run used
DeepSelect with no fallback, on rows x width scores with the expected
row stride (padded only for the unaligned width), and that cuDNN runs
never reached it.
Per top-k call the lengths must be equal, and each row's selection
sorted by key must be bitwise equal with its scores, except rows that
pick different keys among those scoring exactly the row's cutoff (at
most 1% of the rows). The DeepSelect rerun must repeat its raw top-k,
slot order included, bitwise.
Float tolerances are the production-mode ones of
test_dsa_cudnn_skip_score_fill_layer (relative L2 1e-2, peak 2e-2,
indexer loss 1e-3), applied with the cutoff-tie rows in mind. At these
widths cuDNN itself picks other cutoff-tied keys from run to run (up
to 15 rows per rank and call). Such a query row attends to another key
with an unrelated attention weight, so its output and query gradient
move by O(that weight) and the two key rows trade its gradient: the
first GPU run, with the limits applied to every element, failed on the
cuDNN rerun and the shuffled-slot run themselves (peaks up to 0.073 and
0.092 of grad.query, rel L2 <= 1.1e-3) and on DeepSelect in the same
way (up to 0.14 of the output, rel L2 1.1e-3). So with tie rows the
output and query gradient meet the
limits on the other rows, the key gradient meets the relative L2 limit,
and every other gradient and the indexer loss meet them in full; the
cuDNN rerun and the shuffled-slot run must meet the same limits, and
the DeepSelect rerun (no ties) must meet them on every element.
GB300 (4 GPUs, FE 1.31, ds_tune at DeepSelect bfa4507): 13 passed in
55 s. Every layout and rank ran DeepSelect on every call.
Rows with other cutoff-tied keys per rank: cuDNN rerun 2..15, shuffled
1..14, DeepSelect 1..16. Worst rel L2 over all tensors: cuDNN rerun
1.2e-3, shuffled 1.3e-3, DeepSelect 9.2e-4. The largest reported peak
is on grad.key in every comparison (3.1e-2, 2.4e-2 and 3.1e-2); all
limited peaks are within 2e-2. Indexer loss <= 1.8e-7 relative.
DeepSelect rerun: raw top-k bitwise, worst rel L2 4.3e-4, peak 4.8e-3.
The layer helpers are shared with test_dsa_cudnn_skip_score_fill_layer.
CPU tests cover the tie-aware selection comparison, the slot shuffle
and the layout widths against the ds_tune row and width thresholds.
Integration
GOLDEN_CONFIG in test_hybrid_moe_model.py lists
dsa_indexer_topk_backend: 'cudnn'. KPool selection (NVIDIA#8287,
dsa_indexer_kpool > 1) does not run the dense top-k, so the option has
no effect there and config accepts it, like the other perf-only knobs;
the field docstring says so. The packed THD scorer passes the DeepSelect
out= kwargs after the MXFP8 precision kwargs.
The determinism kernel manifest (NVIDIA#7148 via NVIDIA#7960) registers every
Triton-bearing file, so dsa_deepselect_topk.py gets an entry naming its
GPU repeatability tests.
BF16 scores
The DeepGEMM MXFP8 selection scorer of the next commit writes BF16
logits, and DeepSelect's ds_tune build has a BF16 config
(h_t256_o2_b2_2k_d3, max_topk 2048) that reads them directly, so the
top-k backend also takes BF16 scores; no in-tree caller passes BF16
scores yet, and FP32 callers see no change.
Why: on GB300 (CUDA graph p50, K 2048, GLM-5.2 row lengths), DeepSelect
BF16 against DeepSelect FP32 took
- 4.20 vs 6.08 ms at the 1M CP64 rank-0 call (16384 rows),
- 4.97 vs 7.00 ms at 1M CP64 rank 63,
- 0.337 vs 0.469 ms at 128K CP32 rank 0 (4096 rows).
Changes:
- SUPPORTED_DTYPES = (FP32, BF16). Any other dtype is declined with
"scores must be FP32 or BF16, got <dtype>".
- Width thresholds are looked up per score dtype
(SCORE_WIDTH_TABLES_BY_DTYPE; min_score_width(kind, rows, dtype)).
BF16 has its own table: ds_tune from 16384 columns at every row count
(no many-rows step); deep_select keeps its FP32 524288. FP32 is
unchanged (ds_tune 65536, 81920 above 4096 rows).
Measured on GB300 (a BF16 width sweep: CUDA graph p50 of 20
replays, K 2048, cuDNN frontend 1.31 with the 8 GiB scratch budget,
DeepGEMM-shaped rows with unwritten tails, CP-balanced and Gaussian row
lengths): ds_tune BF16 was 1.70-2.69x faster than cuDNN's BF16 radix
top-k at every point, widths 16384..131072 by 2048..16384 rows (64
points, no selection past a row's length), e.g. 0.054 vs 0.099 ms at
2048 x 16384 and 1.62 vs 3.19 ms at 16384 x 131072. The sweep covered
2048..16384 rows at K 2048 only; the table applies 16384 to every row
count and top-k, which extrapolates to fewer rows and other K. Narrower
BF16 widths were not measured and keep cuDNN, and their decline reason
says so ("not measured below it for BF16, cuDNN is kept"); deep_select
was not in the sweep (never timed on BF16) and its reason says that it
keeps its FP32 threshold. So every 128K CP32 and 1M CP64 rank width
and every BF16 call from 16384 columns uses ds_tune. The sweep timed
the selections only (and found none past a row's length); their
exactness at the widths the table opens is checked on GPU (below). For
reference, cuDNN's BF16 radix top-k against its FP32 path (cuDNN
frontend develop 4f175f3f with the 1.31 binary, 4096 Gaussian rows):
0.43 vs 0.36 ms at 32K columns, 0.80 vs 0.62 at 64K, 0.89 vs 1.08 at
128K; 22.75 vs 14.44 ms at the 1M CP64 rank-0 call. Correctness does
not depend on the thresholds.
- The compiled config is looked up for the scores' dtype (ds_tune
config_for(dtype, K)); a build without one declines with "has no
compiled config for top-k K with BF16 scores" (a warning).
- The values output is [rows, K] in the scores' dtype, so its 32-byte
row alignment needs K a multiple of 16 for BF16 (8 for FP32, as the
int32 indices already required); other K keep cuDNN.
- score_row_stride(width, dtype) counts elements of the dtype (1024 B is
512 BF16 elements); the packed THD padded-out sizing passes FP32.
- Values come back in the scores' dtype from both backends (DeepSelect
allocates them in scores.dtype; cuDNN already did). NaN repair writes
NaN of that dtype (the Triton store casts to the pointer type).
- log_topk_backend_once keys on (requested, reason, dtype), so FP32 and
BF16 outcomes log separately; non-FP32 lines name the dtype.
- Docstrings of the module, _indexer_top_k_wrapper_chunked and
dsa_indexer_topk_backend carry the BF16 numbers and contract.
Tests (CPU, fake builds): the ds_tune fake gains BF16 configs and checks
that the config matches the dtype; the cuDNN fake ranks BF16 through its
exact FP32 copy and asserts compact rows. New cases: the BF16 config is
used; no BF16 config keeps cuDNN with one warning; BF16 uses its own
table at both width edges, at every row count and for both builds, while
FP32 keeps its thresholds; every BF16 sweep shape and every measured
BF16 shape (128K CP32, 1M CP64) uses ds_tune; the narrow-width reasons
name each dtype's evidence; a per-dtype table overrides only its dtype
(the fake-build fixture lets BF16 follow its patched FP32 tables); BF16
needs K a multiple of 16 (K 24 keeps cuDNN); the contract against the
cuDNN path on Gaussian, tie-heavy and signed-zero/inf BF16 rows;
indices-only equals the value call; NaN repair with BF16 NaN (torch
fallback and, in the Triton interpreter, both dtypes); per-dtype row
stride; the 1024-byte stride check in BF16 elements; no host sync;
logging once per reason and dtype. The FP32 eligibility case now checks
that BF16 is accepted and FP16 is not.
GPU (test_dsa_deepselect_topk_gpu.py; the 25 tests below passed on
GB300): BF16 selection is an exact top-K against
torch.topk, tie-aware (the selected value multiset equals the K largest
in-row values) at 128K and 1M on Gaussian and tie-heavy rows; NaN rows
repaired with BF16 NaN and other rows untouched; bitwise repeats, batch
independence and indices-only equal to the value call for config 1
(determinism had been shown for the FP32 config only); graph
replay after a data swap at 128K and 1M; every 128K CP32 and 1M CP64
rank width eligible. New, not yet run on a GPU: exactness from the BF16
minimum width (2048 x 16384, 16384 x 16384, 4096 x 32768 and 8192 x
65536, Gaussian and tie-heavy) on DeepGEMM-shaped rows (row stride
align(align(w, 384), 512), unwritten tails poisoned with +inf, NaN and
3e38), with a bitwise repeat.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Robin Zhang <robinz@nvidia.com>
New config field dsa_indexer_scorer_backend ('cudnn' default |
'deepgemm'). With 'deepgemm' and dsa_indexer_precision='mxfp8', the
general packed THD indexer path (any CP size) scores with DeepGEMM's
fp8_fp4_mqa_logits instead of the cuDNN MXFP8 dense scorer. Its BF16
logits go to DeepSelect's BF16 top-k (previous commit) or to cuDNN's BF16
radix top-k. The indexer loss is unchanged by the backend: as for the
cuDNN MXFP8 scorer, the selected scores it consumes are recomputed from
the BF16 Q/K/weights (the MXFP8 commit's rule), so MXFP8 changes only
which keys are selected, whichever kernel selected them.
Why (GB300, DeepGEMM 2.8.2 5081f85 with JIT sm_103f, CUDA graph p50,
K 2048, GLM-5.2 row lengths; quantize + score + top-k + index restore
per call):
- 1M CP64 rank 0: FE BF16 + DeepSelect FP32 58.1 ms, FE MXFP8 + DS FP32
43.0 ms, DeepGEMM MXFP8 + DS BF16 28.8 ms. At 128K CP32 rank 0:
2.38 / 1.84 / 1.04 ms.
- Score buffer: 64 -> 31 GiB at 1M rank 0 (here 32 GiB: the buffer is
sized by the existing host width bound, 1048576, not max(ke - ks)).
- On ternary inputs the DeepGEMM logits equal FE BF16's bit for bit on
all 8.07e9 valid elements (1M ranks, 128K and small layouts).
- Selection: overlap with FE MXFP8 0.992. Recall against the FP32
oracle is set by MXFP8 quantization: 0.960 at 1M and 0.971-0.977 at
16K on untrained-indexer inputs, the same as FE MXFP8. So this stays
experimental: it should not be enabled for training until a loss A/B
or recall on trained weights supports MXFP8 selection. MXFP4 (recall
0.72-0.90) is not exposed.
- Expected step effect [inferred]: 1M -1.41 s/step against FE BF16
(-0.68 s against FE MXFP8), 128K -0.26 s/step, minus the recompute
cost below.
Mapping: each packed THD row r (zigzag CP included) attends to keys
[starts[r], ends[r]) of the original, globally ordered K. DeepGEMM takes
exactly such ranges as cu_seq_len_k_start/end over a [Sk, D] K:
ks = starts, ke = starts + length with
length = min(ends - starts, width) on rows of a packed segment and 0 on
rows past them (cu_q[-1] at CP1, 2 * sum(len // 2cp) at CP>1), which
equals the skip-fill cuDNN path's per-row column limit and its
uncovered-row behaviour. Column t of row r is key ks + t, so the restore
is ks + t, exactly the cuDNN path's topk_indices + starts. The original
K is used: no duplicated-prefix K pool, no index_select (258 MiB less at
1M). The width is today's host bound (_packed_thd_score_width, factored
out of the cuDNN path unchanged), so nothing reads the device.
Data flow (_indexer_topk_packed_thd_deepgemm):
- K [sk, 128] is quantized once and Q [rows, H, 128] per DeepGEMM call
(row chunk) with TE MXFP8 (new
indexer_quantization.quantize_indexer_mxfp8_packed_ue8m0): the 4
logical E8M0 bytes of a 128-element row are one int32 in DeepGEMM's
packed UE8M0 layout (byte-identical to deep_gemm.utils.
per_token_cast_to_fp8(gran_k=32, use_packed_ue8m0=True)); no Triton
scale packing.
- dsa_deepgemm_scorer.mqa_logits builds the scheduler metadata from that
call's (ks, ke) (a captured kernel, never cached across calls) and
calls fp8_fp4_mqa_logits with keyword arguments; it checks the output
contract (BF16, unit columns, 1024-byte rows, an allocation of whole
padded rows) and raises naming the build otherwise.
- Top-k (_deepgemm_logits_topk), decided once per call on the tensor
that will be read: DeepSelect reads the [sq, width] logits in place
when its fallback_reason on them allows it; otherwise cuDNN's BF16
radix top-k (explicitly, so DeepSelect is not reconsidered on the wider
view) reads the compact [sq, row_stride] full-row view (it is compiled
for compact rows). DeepGEMM leaves columns past each row's length
unwritten; cuDNN frontend >= 1.29 never reads them (out-of-bounds lane
guard), and the DeepGEMM path requires that frontend (below).
- Rows: DeepGEMM's scheduler-metadata kernel keeps two words per query
block of 128/H rows in one CTA's shared memory and host-asserts that
it fits 232448 bytes, so one call takes at most 116208 rows with 32
heads and 58104 with 64 (dsa_deepgemm_scorer.max_rows_per_call;
mqa_logits raises above it). One TE MXFP8 cast takes at most
65535 * 64 rows of 128 elements (one grid y block per 64 rows, at most
65535 of them: 65535 query rows with 64 heads, 131070 with 32; on
GB300 a 65536 x 64 Q raised "CUDA Error: invalid argument"; that is
TE's generic cast kernel, TE 2.16's default: its cast-only
kernel, taken with ENABLE_CAST_ONLY=1, has 4-row blocks and no grid
check, and nothing sets that), so Q is quantized per chunk and a chunk
respects both limits (DeepGEMM's is the smaller today); one cast of K
takes up to 4194240 keys. Larger calls are quantized, scored and
selected in row chunks starting at multiples of 128 (16-byte-aligned
slices of Q, scales, weights and ks/ke); rows are independent in all
three, so only the indices and values are copied together. 128K CP32
and 1M CP64 calls (4096 or 16384 rows per rank) always take one chunk,
with no copy.
- The top-k always returns the selected values, so a checkpoint's
indices-only first call and its loss replay run the identical top-k
invocation on bitwise-identical logits. Rows no longer than K skip the
top-k; when also no scores are wanted, DeepGEMM does not run either.
Indexer loss (no DeepGEMM values; the MXFP8 commit's rule, unchanged):
- _indexer_topk_packed_thd hands whatever selection ran (DeepGEMM, or
the cuDNN MXFP8 scorer on a fallback) to the MXFP8 commit's
_selected_scores_for_indexer_loss, which recomputes the selected
scores in FP32 from the BF16 inputs if and only if
dsa_indexer_precision is 'mxfp8'. That is the scorer's compute
precision, which 'deepgemm' requires, never the backend or the dtype
of the selection values (BF16 here, FP32 for the cuDNN scorer); 'bf16'
keeps its own FP32 scores and cannot take DeepGEMM. So a DeepGEMM
call, a fallback call (including its every-key shortcut) and the cuDNN
MXFP8 scorer feed the loss the same payload, identical given the same
selection, and a BF16 / FE MXFP8 / DeepGEMM A/B compares selections
under one loss definition. Frozen and selection-only calls never
recompute (the rule reads the caller's return_topk_scores, which the
dispatcher passes on for either kernel).
- The payload (cuDNN FE sparse_indexer_score_recompute_wrapper on the
selected global keys, padded with -1 to one 128-key-aligned shape per
configured top-k, without topk_length) is log(max(p, FLT_MIN)), -inf at
-1: the BF16-input logits minus the row's LSE over the selection, FP32
[b, sq, K], consumed exactly like logits by every shift-invariant
consumer. Given the selection the KL gradient is the BF16 path's; the
KL value is too, except that a selected key with p < FLT_MIN (a logit
more than 87.3 below the row's LSE over the selection, which exceeds
its maximum by up to ln |S|: from about 79.7 below the maximum at top-k
2048) contributes at most about 87.3 nats where the BF16 path's term
keeps growing (logits [495.06, 0], target [0.5, 0.5]: 42.98 against
246.84), so a loss A/B reports both arms' KL with the same clamp.
- Selections the caller would drop stay dropped: a slot whose selection
value is not finite (a NaN-flagged DeepSelect row, a NaN key a radix
top-k ranked first, a -inf prefill column on a fallback) or that lies
outside its row is set to -1 before the recompute, so such a row or
slot ends with no key, as when the caller drops non-finite scores.
- Why not DeepGEMM's values (CPU emulation of this path on GLM-5.2
indexer inputs): on the selected set the MXFP8 + BF16-
epilogue logit error is 0.038-0.040 (10-17% of the logit spread),
predict TV 0.015-0.016 (max 0.021), KL 7.5-8.3e-4 on 16K and 1M
GLM-5.2 rows. That is a fixed, input-determined bias of
~0.03 * coeff (L1) on the KL gradient that does not shrink as the
indexer converges; DeepGEMM's BF16 epilogue alone adds only TV
0.003-0.004.
- Cost [inferred, to be measured], the same as for the cuDNN MXFP8
scorer: at 1M CP64 rank 0 the recompute gathers ~8 GiB of K rows
(partly L2) and 0.28 TFLOP: ~1.5-3 ms per call against a 29 ms saving;
a 1M run with full recompute and the checkpoint top-k reuse
(required below) recomputes once per checkpointed loss-on call,
in the checkpoint's first call (the stash the replay reuses): 48
calls on the bottleneck stage (6 top-k-computing layers x 8
microbatches), as many as without the reuse. At 128K CP32 rank 0:
2 GiB and 69 GFLOP with K L2-resident, ~0.1-0.5 ms per call against a
1.34 ms saving; all 192 calls on the bottleneck stage of a 128K run
with selective recompute (6 layers x 32 microbatches) are loss-on, so
+0.02-0.10 s against -0.26 s/step. The BF16-compute paths
(dsa_indexer_precision 'bf16') are unaffected.
- A missing recompute entry point raises with the loss on rather than
silently changing the loss.
Fallbacks (logged once per reason, a warning except deterministic):
- Config rejects: unknown value; 'deepgemm' outside experimental variant
'dsa' with dsa_kernel_backend='cudnn'; precision != 'mxfp8' (which
brings every MXFP8 requirement: fused kernels, 32/64 x 128 indexer,
BF16 weights, sparse loss, no KPool); deterministic_mode (BF16 logits
and DeepSelect do not keep lowest-index ties); and a recompute that
replays the indexer (full, or selective with 'core_attn', the default
module list) without dsa_indexer_reuse_checkpoint_topk, with the
indexer loss on or frozen (ties, below).
- At run time the call uses the cuDNN MXFP8 scorer when: deep_gemm is
missing, older than 2.8.1, still has 2.8.0's clean_logits or
logits_dtype argument, or lacks the named arguments; the device is
not SM 10.3 (SM 10.0 is not validated yet) or not this process's
DeepGEMM device (below); heads/dim/dtypes do not fit; TE MXFP8 is
missing; the cuDNN frontend is older than 1.29;
torch.use_deterministic_algorithms is on; or the JIT warm-up fails
(recorded, not retried; a CUDA out-of-memory error propagates instead
and is not recorded). Where the cuDNN MXFP8 scorer cannot run either it
raises its own errors. Non-THD layouts and the dense loss raise the
existing MXFP8 error; _check_scorer_precision_supported also rejects an
unknown backend and 'deepgemm' without MXFP8.
- The frontend rule holds for every call, not only for those DeepSelect
declines: a per-call rule would switch scorers between batches under
dynamic CP. It matches dsa_indexer_precision='mxfp8', which needs
frontend >= 1.29 anyway; on FE 1.21.1 the fallback raises the cuDNN
MXFP8 scorer's own error.
- A fallback scores in FP32, so its score buffer is twice DeepGEMM's
(64 against 32 GiB at the 1M CP64 rank-0 call); the config docstring
says so.
Ties and recompute (measured on GB300, cuDNN frontend 1.31): on
tie-heavy BF16 logits DeepSelect BF16 repeats its selection bitwise, but
cuDNN's BF16 radix top-k keeps other tied keys on a repeat (128K CP32
rank 0 and every width 4K..1M; its selections stay exact
top-K, so only keys equal to the K-th value differ). cuDNN's takes every
call DeepSelect declines (all calls with dsa_indexer_topk_backend
'cudnn', and with 'deepselect' the BF16 rows narrower than 16384). A
checkpoint replay that runs the indexer again therefore may select other
keys than its first call. That selection feeds the recomputed attention,
whose backward must belong to the first call's output, and with the
indexer loss on, the loss as well; a frozen indexer's replay without the
reuse reruns the dense indexer and top-k too. So the config rejects
'deepgemm' with full recompute or selective 'core_attn' recompute unless
dsa_indexer_reuse_checkpoint_topk is set, with the indexer loss on or
frozen (the replay then reuses the first call's selection). The check
runs after the recompute options are final (default module list,
moe_layer_recompute). Full recompute with the reuse (loss on and
frozen) and selective 'layernorm' recompute, which does not replay the
indexer, stay valid.
CUDA graphs and JIT (dsa_deepgemm_scorer):
- All decisions are host facts, stable between warm-up and capture; no
tensor is read on the host; every buffer is call-local.
- The first eager call per (device, heads, head dim, MXFP8, scheduled,
deep_gemm.get_num_sms()) runs warmup() (128 x 256); mqa_logits raises
instead of compiling or loading a kernel inside capture (and
_deepgemm_scorer_fallback_reason raises there for a cold key), so a
capture needs >= 1 eager step or an explicit warmup(). The guard reads
only a per-(device, heads, head dim) record of eager runs and warm-up
failures before any deep_gemm call: a process's first call, even
get_num_sms(), builds DeepGEMM's runtime (cuBLASLt handle, device
properties), which a capture cannot record (on GB300 the old order
raised another error from inside get_num_sms in a fresh process). A
recorded warm-up failure is per device too, so a capture falls back
like the eager steps did.
- One DeepGEMM device per process: deep_jit caches each kernel's
CUfunction per process, keyed by its source, not by device or CUDA
context (on GB300, ranks that warmed up on cuda:0 and then ran on
their own device got CUDA_ERROR_INVALID_HANDLE). The first device that
calls DeepGEMM owns it, claimed before its first deep_gemm call; a
call on another device falls back to the cuDNN MXFP8 scorer with that
reason logged once (inside capture too), and a direct mqa_logits or
warmup on it raises. Megatron runs one device per rank, so this never
triggers in training; warmup() must run after torch.cuda.set_device.
- DG_JIT_CACHE_DIR: a persistent directory shared by all ranks of an
image, pre-populated by one single-GPU warmup(); unset, deep_jit uses
$HOME/.dj and the backend line warns once. nvcc must be present at run
time.
Import and packaging:
- deep_gemm is imported lazily at the first DeepGEMM-scored call, or
reused from sys.modules, and probed once (version >= 2.8.1, both
functions with the named arguments, parsed from the pybind signature
line, neither of 2.8.0's clean_logits/logits_dtype arguments,
get_num_sms). The log line names version, file and JIT cache.
- batch_invariant_kernels imports deep_gemm at Megatron import time; its
except now also catches OSError, RuntimeError and AttributeError, so a
build for another torch cannot break Megatron's import. When a
DeepGEMM build is importable, batch-invariant mode would also pick it
up; make it importable only for runs that use this scorer.
- No new hard dependency.
Also: IndexerScorerOptions.scorer_backend (read from the config); the
host width bound _packed_thd_score_width factored out of
_indexer_topk_packed_thd_cudnn without behaviour change for 'cudnn'; the
config docstrings (dsa_indexer_scorer_backend defers the loss to
dsa_indexer_precision's rule, which names the backend and its fallback).
GOLDEN_CONFIG in test_hybrid_moe_model.py lists
dsa_indexer_scorer_backend: 'cudnn'.
Tests (CPU, test_dsa_deepgemm_scorer.py, fakes in dsa_deepgemm_fakes.py;
the FE recompute fake and the recompute helper's own tests come with the
MXFP8 commit): the import probe (missing,
ImportError/OSError/RuntimeError/AttributeError, old or unparsable
version, renamed or unnamed pybind arguments, missing functions, one
import, sys.modules reuse, batch_invariant_kernels staying importable
for each exception type); the fallback matrix with its log-once levels
(SM 10.0, no build, deterministic flag and argument, JIT failure) and
the static mismatches reaching the cuDNN MXFP8 errors; the capture guard
(cold key, eager warm-up, head count and SM count in the key, direct
mqa_logits and warmup, no deep_gemm function called before it, a failed
warm-up falling back inside capture); one DeepGEMM device per process
(another device falls back with the reason logged once, inside capture
too, direct calls raise, a warm-up that fails in the logits kernel or
before it still claims its device); output-contract violations; argument
plumbing (original K, int32 ranges from the same tensors for metadata
and logits, the host width, BF16 weights, quantized shapes, no K pool);
the TE scale int32 view; row lengths equal to the cuDNN column limits on
randomized CP1/2/4/8 layouts with padded cu_seqlens and uncovered rows;
selection as a tie-aware top-k restored to global keys on ragged
CP1/CP2/CP4 layouts for both top-k backends; ternary selection equal to
the cuDNN MXFP8 scorer's tie-aware; every-key calls without scoring;
poisoned tails never selected; frontends < 1.29 keeping the cuDNN MXFP8
scorer for every call (and FE 1.21.1 raising its MXFP8 error);
DeepSelect reading the logits in place, cuDNN reading the compact
full-row view, and a call DeepSelect declined never reaching it on the
wider view; DeepGEMM's row limit against the fake's shared-memory model,
mqa_logits refusing more rows, the chunk ranges (within TE's MXFP8 cast
limit too), K quantized once and Q per chunk, and chunked calls
selecting and scoring what one call does (CP1/CP2, ranked and every-key,
both top-k backends, no host sync); a CUDA OOM in the warm-up
propagating unrecorded; indices-only against value calls; the loss-score
rule as a precision {bf16, mxfp8} x backend {cudnn, deepgemm, deepgemm
falling back} x loss {on, frozen} x {ranked, every-key} matrix (which
kernel ran, which path supplies the loss scores, recompute calls; bf16
with deepgemm rejected); the every-key payload bitwise identical for the
cuDNN MXFP8 scorer, DeepGEMM and the fallback; the DeepGEMM loss payload
against BF16 log-softmax, a missing entry point, equal KL and gradient
inputs in _compute_sparse_indexer_loss_and_grads, the checkpoint score
stash accepted by the replay, the frozen path, fill-on uncovered rows on
the fallback, NaN keys past the repaired window and NaN keys picked by
the radix top-k (both NaN signs); no host sync; non-THD layouts; and the
config matrix, including the replay rule (rejected: full, 'core_attn',
the default modules, 'layernorm' with 'core_attn', for both top-k
backends, with the indexer loss on and frozen; accepted: the reuse with
the loss on and frozen, selective 'layernorm' (+ 'mla_up_proj', and
frozen), the 'cudnn' scorer, no recompute). Mutations of the DeepGEMM
path, the DeepGEMM module, the config, batch_invariant_kernels and the
BF16 DeepSelect rules each fail at least one of these tests.
test_dsa_deepgemm_gpu_helpers.py checks the GPU tests' comparison
helpers on CPU: the exact BF16 and FP32 spacings and the FP32 ulp
distance; G10's slot-order bound against an emulation of the
recompute's FP32 row sum under random slot orders (a row normalized by
another row's sum, or by a sum off by 2**-15, fails it); the graph
test's dkey checks against an emulation of the atomic FP32 dKV sum, the
BF16 partials and the 2-rank reduce-scatter (coherent, mixed,
zero-mean, wide-range and cancelling contributions pass; one query's
contribution moved to another key row, two swapped key rows, a 2x error
on 1% of a partial, one misrounded dkey element and a reduce-scatter of
the other rank's rows fail); the distinct-key checks, the top-k
recorder, the key-partial hook and the collected failures. Mutations of
these helpers each fail at least one of these tests.
GPU (on GB300; the GPU runs found the guard order, the TE cast limit,
the per-process kernel handles, the cuDNN BF16 radix ties, the G10
payload tolerance, the graph test's pinned selection and test defects,
all fixed here; every DeepGEMM-scored call asserts that DeepGEMM ran
with no fallback reason):
test_dsa_deepgemm_scorer_gpu.py G1-G10 (output contract; bit-exact
mapping on ternary inputs for CP1, zigzag CP2/CP4, unequal segments,
uncovered rows (CP1, and CP2, which only the DeepGEMM path takes) and a
block straddling the covered boundary, with a poisoned-block write check
that fails closed (Q and K are quantized first and the poison block has
the size of DeepGEMM's logits allocation, so the allocator must hand it
back) and tie-aware selection against FE BF16; overlap with FE MXFP8 and
recall against FE BF16 on Gaussian and outlier inputs; graph
capture/replay against eager with set_sync_debug_mode('error'),
tie-aware on the DeepGEMM logits for the cuDNN top-k, and the guard in a
fresh process, for a direct mqa_logits and for the production entry
point (_deepgemm_scorer_fallback_reason, which reaches jit_problem),
which prints any exception's type and text and checks that no DeepGEMM
device was claimed; repeatability, batch independence and indices-only
against value calls, bitwise for DeepSelect BF16 and, for the cuDNN BF16
radix, a documented negative asserting tie-only differences and
reporting whether the repeats were bitwise; NaN outcomes; loss and
gradients equal to the BF16 path given the selection, with the predict
TV of DeepGEMM's own values reported (the recompute's own parity check,
G7, is in test_dsa_cudnn_mxfp8_scorer_gpu.py with the recompute); cuDNN
BF16 radix at 4K..1M widths, repeats compared tie-aware; row chunks
against one call at 128K CP32, tie-aware for the cuDNN top-k, with the
loss payload bitwise where the slots match and, where the cuDNN top-k
returned the same keys in another slot order, within the absolute bound
the recompute's FP32 row sum, its only slot-order-dependent step, puts
on log p: 2**-18 plus 2 FP32 ulp of log p (1 ulp, 3.8e-6, measured), and
65536 rows with 64 heads in two calls; every tie-aware comparison also
requires each row's valid slots to name distinct keys).
test_dsa_deepgemm_layer_gpu.py
(each test warms DeepGEMM up after model-parallel init has set the
rank's device): checkpoint replay with the loss (reuse off and on, stash
declined) and frozen replay, both on one rank with Q and K sharing one
padded cu_seqlens tensor (the MXFP8 scorer's boundary proof) and a
pinned top-k (the first call's selection, checked against each later
call's logits, since the cuDNN BF16 radix does not repeat among ties);
production layers against the cuDNN MXFP8 scorer (CP1, zigzag CP4,
dynamic CP and a wide DeepSelect layout; both arms take the recompute);
and CUDA graph replay of a CP2 layer with the prebuilt packed-CP layout,
before and after an in-place refresh to another packing: the replay's
logits bitwise and its own cuDNN BF16 radix selection tie-aware against
eager's (no key taken twice) and, given eager's selection, its output,
dquery and dup_v_weight bitwise. dkey is checked per element: each
rank's partial (the gradient at the key's CP all-gather, hooked) within
the bound of cuDNN's FP32 atomic dKV summation order (1 BF16 ulp of the
element plus 2**-9 of its key row's largest partial, assuming that each
element's contributions add up in magnitude to at most 64 times that
partial), each run's dkey one BF16 rounding of the summed partials (the
2-rank reduce-scatter), and dkey within the partials' differences plus
that rounding. The test feeds eager's selection through buffers
allocated before the capture (allocated inside it, they came from the
graph's pool and every replay rewrote them before the layer read them),
captures with set_sync_debug_mode('error'), collects every check (from
the warm-up and the eager references through the capture to each
replay) and raises on both ranks at the next synchronization point, so
no rank waits alone in a collective (an error inside one rank's own
eager step still can), and releases the graph before the process groups
are destroyed (a live graph holding the CP collectives hung the
teardown).
test_indexer_quantization.py checks the packed UE8M0 bytes against
deep_gemm.utils. The measurements above used FE 1.31 and CUTLASS DSL
4.6.2 for every arm.
Megatron's batch_invariant dependency group pins deep_gemm ea2b7805
(2.8.0). That release's fp8_fp4_mqa_logits still takes clean_logits
(default True) and logits_dtype (default FP32), which DeepGEMM 2.8.1
(057ca59) removed when it made the SM100 logits BF16; called with this
scorer's arguments, 2.8.0 fails its BF16-weights host assert in every
warm-up. The probe therefore parses major.minor.patch, requires >= 2.8.1
and rejects a build whose fp8_fp4_mqa_logits still takes either
argument, so a 2.8.0 install falls back to the cuDNN MXFP8 scorer at
import (and the GPU tests skip) instead of failing each warm-up. The
probe tests cover 2.8.0 (the pin), a version without a patch number, a
2.8.1-labelled build with 2.8.0's signature, and 2.8.1, 2.8.2+local and
3.0 being accepted.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Robin Zhang <robinz@nvidia.com>
… budget The cuDNN radix top-k allocates a transient int32 scratch of rows x 2 x cols per launch, and _indexer_top_k_wrapper_chunked sizes its row chunks to a fixed 2 GiB budget: 256 rows per launch at 1M keys. Long rows run faster in fewer, fuller launches, but a larger budget costs transient memory on every call with long rows, so make the budget a config option and keep the 2 GiB default. dsa_cudnn_topk_scratch_gib (GiB, default 2.0, rejected unless positive and finite; --dsa-cudnn-topk-scratch-gib) reaches every cuDNN top-k of the DSA indexer as IndexerScorerOptions.topk_scratch_bytes, which _scorer_options_from_config fills for every hook. _indexer_topk_bshd passes it, next to the top-k backend, to the dense, score-chunk, single-sequence segment and general packed THD executors, and the DeepGEMM scorer passes it to the cuDNN BF16 top-k it runs when DeepSelect does not take its logits. Rows per launch stay capped by the frontend's 2**31-element limit, so budgets above 8 GiB give no larger launches, and a launch takes at least one 32-byte-aligned row group. Budgets above 2 GiB change only the executors that hand the top-k a whole score buffer, the general packed THD executor (cuDNN scores or DeepGEMM logits) and the dense indexer loss: the plain causal, bounds and single-sequence segment executors score in chunks of at most 1 GiB of FP32 scores, whose top-k the 2 GiB default already runs in one launch. Calls without a budget keep the module default, _TOPK_WRAPPER_MAX_SCRATCH_BYTES (2 GiB). Guidance, measured through _indexer_top_k_wrapper_chunked on GB300 (cuDNN frontend 1.31, k=2048): 8 GiB takes a 1M-token CP64 rank-0 call from 21.05 ms (64 launches of 256 rows) to 14.26 ms (16 launches of 1024 rows) in a CUDA graph and 22.06 -> 14.54 ms eager; that is about 0.65 s per 1M-token step if all 96 calls sit on the bottleneck stage. At 128K CP32 all 4096 rows of the rank-0 call run in one launch (0.593 -> 0.566 ms). The cost is memory: each call's transient scratch is min(budget, twice its FP32 score buffer), eagerly and inside CUDA graph pools, so with 8 GiB every whole-buffer call with more than 1 GiB of FP32 scores takes more than at the default, e.g. a 32K CP1 packed THD call (4 GiB of scores) takes 8 GiB where it takes 2 GiB. 8 GiB was checked only for 1M-token CP64 and 128K-token CP32 calls of one GB300 configuration; check the headroom of the tightest pipeline stage before raising the budget. 4K CP1 calls stay far below 1 GiB of scores. Tests: the config default is 2 GiB and equals the module default, 8 GiB is accepted and zero, negative, NaN and infinite budgets are rejected, the CLI flag parses, _scorer_options_from_config converts GiB to bytes (None without the field), the wrapper's budget argument overrides the module default, and every _indexer_topk_bshd executor (plain causal, bounds and general packed THD, with and without dense scores), the single-sequence segment executor and the DeepGEMM logits top-k launch with the given budget. The 8 GiB chunking cases pass the budget to the wrapper, the DeepGEMM GPU test stubs of _deepgemm_logits_topk take and forward the new argument, and GOLDEN_CONFIG gets the new field. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
num_floating_point_operations rejected every hybrid model with 'D' (DSA) layers in its --hybrid-layer-pattern, because only the standard-model path in transformer_flops had a DSA term. Count them on the hybrid path with the same terms instead: - Attention: a 'D' layer is absorbed MLA. Its token-linear projections equal the MLA ones, since absorption only moves the K/V up-projection to the query and output side. Core attention spans the KV latent and is scaled from dense causal pairs to the top-k pairs (_dsa_sparse_core_scale), or to the selected pools with KPool (_dsa_kpool_attention_scale). mla_attn_layer_flops gains the absorbed and sparse_core_scale arguments for this; other MLA callers are unchanged. - Indexer: _dsa_indexer_flops, executed at the layer numbers DSAttention uses for cross-layer top-k sharing: pattern positions for the main decoder and num_layers + position for each MTP depth (_num_hybrid_dsa_indexer_layers). The indexer carries its own forward/backward factor (forward only unless its KL loss is on, which KPool does not allow), so it is added after the global x3, as on the standard path. An explicit 'D' pattern now gives the same count as the equivalent standard model, with and without KPool. Still rejected: - a model-wide experimental_attention_variant='dsa' without 'D' layers, e.g. DSA over GQA through a custom '*' spec, which this does not model; - 'D' layers under another variant than None or 'dsa'. Under dsv4_hybrid, 'D' remains DSv4 attention and is unchanged; - KPool with a top-k that is not a multiple of the pool size, or with the indexer loss on, through the shared KPool helpers. This only changes the FLOPs estimate behind the logged TFLOP/s and throughput; model numerics and training are unaffected. Estimates for models without 'D' layers are unchanged. test_num_floating_point_operations.py replaces the two hybrid DSA rejection tests with the two remaining rejections, and checks that explicit 'D' layers match the standard path at pool size one (with and without MTP, indexer loss and sparse loss) and with KPool, that MTP indexer executions use physical layer numbers, and that invalid KPool settings are rejected. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
Add docs/user-guide/features/dsa.md and list it in the features toctree. The guide's Overview summarizes Megatron Core's DSA support (kernel backends, DSA over GQA, KPool, IndexShare, all-gather context parallelism, packed THD), then its Features part gives one self-contained section per feature: what it does, its options and defaults, the rules that TransformerConfig validation enforces, and who benefits. The sections cover repeated-MTP sharing (NVIDIA#6472, NVIDIA#7107), DSA layers on MTP-only pipeline stages (NVIDIA#8535), the KPool fixes (NVIDIA#8536) and HF-reference KPool scoring, the inverse-permutation CP key-reorder backward and prebuilt packed-CP layouts, checkpoint top-k reuse under activation recompute, IndexShare pipeline-layout validation and the shared non-empty certificate, the cuDNN frontend >= 1.26 call sites (NVIDIA#8534), deterministic cuDNN DSA with the tie_break top-k, cuDNN top-k launch sizing and the dsa_cudnn_topk_scratch_gib budget, the -inf score prefill skip, MXFP8 indexer precision, and the DeepSelect top-k and DeepGEMM scorer backends. Each section can move into its feature's pull request on its own. A configuration-rules table, a dependency matrix (cuDNN frontend 1.26, 1.29, 1.31 and the unreleased clean_logits switch, CUTLASS DSL, Transformer Engine, Triton, DeepSelect, DeepGEMM and their fallbacks) and the known limitations close the page, including the overlap with NVIDIA#6793, NVIDIA#7813, NVIDIA#8403 and NVIDIA#8154 and the scope of the forward sharing state. Experimental features (MXFP8 selection, DeepSelect, DeepGEMM) are marked as such. Measured numbers are limited to those stated in the public GLM-5.2 tracker NVIDIA#6944 (body as of head edebb6b) and scoped to its 256-GPU GB300 configuration. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Robin Zhang <robinz@nvidia.com>
|
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. |
This branch has not been 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.
The feature guide added by the last commit,
docs/user-guide/features/dsa.md, documents every option, its default, the rules between options, and the library versions.Commits
Stack head
a3640cd26fondevc4220af3a6. Commits 1, 2 and 4 have the same tree (1, 2) or the same patch (4) as the current heads of #6472, #7107 and #8536; commits 3 and 9 have the same patch as #8535 and #8534.mtp_repeated_layer_shared_components)dsa_indexer_reuse_checkpoint_topk)dsa_indexer_topk_freq > 1)--deterministic-mode)clean_logits(NVIDIA/cudnn-frontend#1761)tie_break=1dsa_cudnn_topk_scratch_gib(cuDNN top-k scratch budget, default 2 GiB)Validation
run_ci_test.sh: intests/unit_tests/transformer/**, 1815 passed, 63 skipped, 0 failed (dev at the base: 1743 passed). The other changed test files pass:models/test_gpt_model.py,models/test_hybrid_moe_model.py,test_checkpointing.py,test_num_floating_point_operations.py,training/models/test_gpt_builder.py, the text-generation controller tests, and the determinism kernel manifest and inference kernels.test_cudnn_indexer_topk_single_packed_cp_real_kernel_uses_bottom_right_alignment(DSA cuDNN real-kernel tests (packed THD + CP) return wrong results on GB200/SM100; masked byflaky/flaky_in_devmarkers #7105) fails at the base and passes here.Known issues and overlaps
DSAttention._get_index_share_topk_holder, which Megatron-Bridge's GLM hybrid attention overrides.main) isolates the index-share holder during recompute with a different design. [Dev] Introduce forward sharing state and enable repeated MTP sharing #6472/[Dev] Extend recompute support for repeated MTP sharing #7107 need to be reconciled with it.🤖 Generated with Claude Code