Conversation
* Add Sage FP8 BSA kernels * Fix BSA softmax equation rendering * Fix FP8 BSA split workspace fallback * Optimize Blackwell BSA forward kernels * Guard Sage FP8 exact-layout fallback
…A#710) The generic compressor kernels (compressor_sm100.py) are JIT-specialized per (ratio, head_dim, coff) and carry no ratio-dependent assumption below the documented register bound at ratio ~ 32; the {4, 128} envelope in api.py's check_support encoded validation scope, not a kernel limitation. Widen the generic-family branch to ratio in {2, 4} -- ratio=2 (coff=2) is the compressor configuration used in production training of the model family this operation serves. No kernel changes; the ratio=128 family and its routing are untouched. Validation on B200 (CC 10.0, torch 2.13, cutlass-dsl 4.5.1): - ratio=2 test cases added mirroring every ratio=4 pattern (numerics vs the fp32-intermediate eager reference with the bitwise dKV/dScore contract now asserted at ratio in {2, 4}, ragged/short/empty-segment packs, static-capacity padding, NaN-canary zero-write ownership, replay determinism, check_support acceptance/rejection): module suite 145 passed, 1 skipped (multi-GPU test on a single-GPU runner); the new cases fail with the envelope ValueError on the unpatched gate. - at ratio=2, out/dKV/dScore are bitwise-identical to the fp32 eager reference in every tested config; dAPE differs at fp32 atomic-order level only (~1e-6 relative), per the documented contract. - nsys pure-kernel at ratio=2 coff=2 (50-iter mean): fwd 4.7/10.4/13.5/ 39.6 us, bwd 11.4/25.6/27.7/69.2 us for d128/d512 x 1/3x8192-token packs -- 18-21x the eager reference region. Signed-off-by: zky <kaiyue.zhou@z.ai>
* generalize-utmastg * a * per-tensor-tmastg * fix-code-rabbit-comment
…on-contiguous, dense-compatible layouts (NVIDIA#712) * Enable FROST SDPA forward engines to write dense LSE directly to non-contiguous, dense-compatible layouts Signed-off-by: Haobin Guo <haobing@nvidia.com> * Address comments --------- Signed-off-by: Haobin Guo <haobing@nvidia.com>
* benchmark: sample the SM clock of the GPU the benchmark actually runs on The peak-MMA/SOL clock sampler indexed NVML with torch.cuda.current_device(), but NVML enumerates PHYSICAL GPUs and ignores CUDA_VISIBLE_DEVICES while torch indexes only the visible subset — so a shard pinned to GPU n via CUDA_VISIBLE_DEVICES always sampled physical GPU 0. When several single-GPU shards run side by side on a multi-GPU node, a shard whose neighbor GPU 0 has drained records GPU 0's IDLE clock as the window peak, collapsing the chart's MMA-max line and SOL% by the idle-vs-boost ratio (observed >10x too low on GB200/GB300). Map the torch index through CUDA_VISIBLE_DEVICES (index, UUID and MIG forms) before asking NVML for the handle. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * benchmark: add the Ampere (sm80) row to the peak-MMA table _FLOPS_PER_CLOCK_PER_SM had sm90/sm100/sm12x entries only, so A100 runs computed no peak_mma_tflops and their charts drew no MMA-throughput max line. A100: 312 dense BF16/FP16 TFLOPS (FP32 accumulate; 624 is the sparsity figure) = 108 SMs x 1.41 GHz x 2048 FLOPs/clk/SM. No fp8/mxfp8 entries -- Ampere has neither datapath, and those cases already record unsupported. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
* refactor-m-major-output * remove-dead-code
…#736) Add a pr-merge-requirements workflow that fails while a PR has no Milestone or is not on any Project board, so it can be made a required status check. Bot-authored PRs and PRs labeled cat-routine-update are exempt. The check queries live PR state, so a manual re-run after setting the fields is enough to turn it green. Runs as pull_request_target (fork PRs need the repo secret) without checking out PR code. The Projects lookup needs a PROJECT_READ_TOKEN repository secret, since the built-in GITHUB_TOKEN cannot read Projects v2. Also add a PR-template checkbox reminding authors to set both fields. Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
…A#739) If the PROJECT_READ_TOKEN PAT is rejected (e.g. NVIDIA enterprise forbids classic tokens with >366-day lifetime), the GraphQL call aborted the script via errexit with only a cryptic exit code. Capture the failure and surface the API error message instead. The check still fails closed either way. Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
NVIDIA#728) With pre-9.26 cuDNN headers (or _WIN32) the whole #if body compiles away and 'm' is unused; -Werror=unused-parameter then fails the pip source build (seen in containers shipping older cuDNN headers, where every 'pip install .' of current develop breaks). Mark it [[maybe_unused]]. Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
* style: apply black formatting to benchmark and frost test files pre-commit's black hook (26.3.1, line-length 160) reformats these three files; clean them up so `pre-commit run --all-files` passes in CI. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * ci: enforce pre-commit style check via GitHub Actions Runs `pre-commit run --all-files` (clang-format v21.1.6 + black 26.3.1, as pinned in .pre-commit-config.yaml) on every PR and on pushes to develop/main, on a plain ubuntu-latest runner — no GPU needed. This replaces the internal analysis:clang-format CI job and enforces the formatting contract already documented in CONTRIBUTING.md. Also fix a stale reference to ci/run_style_check_diff.sh in the frost README (that script is internal-only and superseded by pre-commit). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
… (default SM100/SM90) (NVIDIA#572) The default indexer-backward pipeline mutates attn_score/index_score in place in kernel 1 and only faults (or corrupts memory) later in the GEMM when the plan signature is inconsistent, so validation problems used to fail dirty. Validate the full signature up front, before any kernel launch: * check_support now enforces the output-dtype contract (d_index_q and d_weights bf16-only -- the kernel dW store rounds the fp32 accumulator to bf16, so an fp32 d_weights buffer cannot be produced faithfully and is rejected instead of silently receiving bf16-precision values; d_index_k accepts bf16 or fp32), the semantic shape relationships between all nine tensors, and compact-contiguous layouts (the kernels address K/dK with a hard-coded compact (D, 1) stride and the backend caches do not key the layout). * execute() re-validates the runtime tensors against the descriptors captured at plan-build time (dtype/shape/stride) so a directly-built or exported plan reused with a mismatched tensor raises a clean ValueError while the score buffers are still pristine. * IndexerBackward and the wrapper validate ranks (index_q 4D, index_k / topk_indices 3D) before deriving plan dimensions. * An fp32 d_index_k output buffer is now zeroed internally on the selected stream (the dK epilogue atomic-adds into it) on both SM100 and SM90, removing the fragile caller pre-zero contract. * The wrapper plan cache keys the output dtypes (d_index_q / d_weights / d_index_k) so the output-dtype validation cannot be skipped on a cache hit. The bf16 compute path is untouched. On SM100, dQ/dW and the in-place score-grad outputs stay bitwise-identical to the previous default backend; on SM90, dQ and the score-grad outputs are bitwise-identical, while dW sits in the pre-existing SMEM-atomicAdd jitter band (SM90 dW was never run-to-run deterministic, even on the unpatched base). dK is an fp32 atomic scatter (within jitter) on both. Signed-off-by: zky <kaiyue.zhou@z.ai>
…#684) Six independent changes to FlashAttentionDSABackwardSm100, measured on a B200 at a locked 1830 MHz, S=8192, over topk 128/512/2048 x causal 0/1. All twelve scenarios improve: D512 by 4.9-8.9%, D576 by 11.1-13.3%. dq is bit-identical to the previous kernel in every scenario (relative error exactly 0), and the launch shape is unchanged: grid 4096, block 640, 96 regs/thread, 216 KiB SMEM/CTA on D512 and 232 KiB on D576. Gather-index path. The per-tile top-k indices were read by lane 0 one row at a time into an rmem tensor, then broadcast. Lane i now reads its own row in _load_tile_topk_idx and the consumer shuffles it out, so the indices live in a single register and the next tile's can be fetched at the tail of the current iteration instead of at the head of the next one. Register budget. num_regs_load_KV 40 -> 56 clears the spills on the gather address path (all of them on D512, ~99% on D576). The per-warp counts must exactly exhaust the CTA pool, which the 640-thread launch fixes at 96 regs/thread: 128*56 + 128*128 + 256*128 + 128*40 = 61440 = 96 * 640. MMA order. dQ = K @ dS now issues before dKV = Q @ dS, which puts load_mma_K_pipeline.consumer_release ahead of the dKV GEMMs rather than after them, freeing the K buffer earlier in the iteration. S lifetime. The fenced T2R of S has fully consumed TMEM S and nothing below reads it, so its consumer_release moves up to the fence instead of trailing P's publication. This unbinds S's lifetime from P's, nothing more. dQ epilogue. store_dQ moves from Ld32x32bOp -- the last one in the file -- to the Ld16x256bOp/StMatrix form the P and dS paths already use, and the four 128-dim sub-tiles each get their own staging slot in the dead K buffer, so a store no longer waits on the previous store's SMEM read. The 64-column D576 tail stages through the dead P buffer instead of sK: sK does have spare room, but only at an offset past the four dQ sub-tiles, which would need its own swizzle/TMA-box proof. That alias is an exact fit today (cosize 4096 == 4096), so it and the sK alias both gain an assert to catch a future stage bump. test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py: 15 passed, 1 skipped. Signed-off-by: Butterfingrz <13524387014@163.com>
…VIDIA#734) * samples: skip deterministic sdpa backward test on old cuDNN versions Blackwell + cuDNN<9.18 + deterministic SDPA backward is a combination explicitly marked as unsupported at include/cudnn_frontend/node/scaled_dot_product_flash_attention.h:1385. * samples: fix invalid double destroy of child CUDA graph in cudagraphs sample In "Cuda graphs with matmul add", `cudnn_cuda_graph_new` is obtained via `cudaGraphChildGraphNodeGetGraph(cudnn_node_in_main_graph, &cudnn_cuda_graph_new)`. Per the CUDA Runtime API documentation for `cudaGraphChildGraphNodeGetGraph`: "This call does not clone the graph. Changes to the graph will be reflected in the node, and the node retains ownership of the graph." Destroying `main_cuda_graph` with `cudaGraphDestroy(main_cuda_graph)` destroys the parent graph and all embedded child graphs owned by its nodes. Calling `cudaGraphDestroy(cudnn_cuda_graph_new)` afterwards attempts to destroy an already-destroyed graph handle, returning cudaErrorInvalidValue and leaving a sticky CUDA runtime error in the process device context that corrupts subsequent tests executing on that context in monolithic test runners like Catch2. Remove the invalid `cudaGraphDestroy(cudnn_cuda_graph_new)` call and wrap destruction calls in `CUDA_CHECK`. --------- Co-authored-by: Marcin Radomski <dextero@google.com>
* add sm120 matmul support * resolve conflicts for sm120 matmul * add benchmark test file for sm120 matmul * resolve issues about nvvm.elect_sync(), nvvm.griddepcontrol(wait), and drain of TMA warp for SM120 matmul
* add det 2k * add bias * fix * fix pipeline * refactor * fix * NFC change, refactor code * add _checked_lse/sink/bias_view * NFC refactor
…cks (NVIDIA#754) _run_dsl_graph returns the stats buffer allocated by make_dense_stats, whose shape is (B, H, S, 1); the pack_gqa features/qtrim tests compared it against the (B, H, S) reference and failed on the shape check. Squeeze at the comparison, like the strided-stats test does. Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
…VIDIA#740) * frost(sdpa): derive THD token capacity from the view's element span (NVIDIA#613) The zero-host-read THD execute (NVIDIA#606/NVIDIA#608) derives the packed token extents host-side as numel() // token_stride. That is wrong on both edges for the buffers real integrations bind: - A non-packed VIEW — a K/V slice of a kv-interleaved [T, 2, H, D] record, the layout torch.nn.attention.varlen users produce by slicing a fused KV projection — holds T tokens but only T*H*D of the record's elements, so the derived extent HALVES and the TMA descriptors cut off half the tokens: silently wrong O on every such call (issue NVIDIA#613; also 40 upstream PyTorch test_varlen_attention failures through the python-API integration). - Deriving from the untyped storage instead over-claims into ALLOCATOR SLACK, which is not benign: rows between the real packed total and the extent are masked but still multiplied (P == 0 times V), so they must be FINITE — TMA zero-fill only covers rows at or beyond the extent. A slack row carrying NaN bit patterns poisons whole sequences through 0 * NaN. Fix: capacity = the largest T whose final token's ROW still fits in the buffer's own element SPAN (1 + sum((size-1)*stride)). The span is exact on both edges: flat capacity buffers give exactly their token capacity (no slack), and interleaved/gapped views give exactly T. Every row below the capacity lies in caller-provided finite elements; every row at or beyond it TMA-clips to zeros. One shared helper serves the SM100 f16 path and the SM120/FP8 _cap sites. Verified on SM100 (isolated env): the NVIDIA#613 kv-interleave repro 41% -> 0 mismatches (frost-served); test_sdpa_random_fwd_ragged_L0 5-seed slice 84/84 (no regressions); fp8 THD ragged slice green; the new deterministic regression test (fused-record K/V views vs packed binding, torch.equal) fails on develop and passes with the fix; upstream PyTorch test_varlen_attention returns from 100 pass / 69 fail to its 140 / 29 impl-identity baseline with the torch-ops stack applied on top. Fixes NVIDIA#613. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * test(sdpa): actually fuzz ragged token gaps in the randomized sweeps The seeded per-tensor token-gap draw (NVIDIA#516) lives in ExecConfig.fill_derived_fields and only fills strides left None — but RandomizationContext, which drives every test_sdpa_random_*_ragged sweep, explicitly assigned packed bshd strides in its ragged branch. Net effect: the randomized ragged fleet has NEVER bound a non-packed THD stride, and for packed buffers the numel()//token_stride capacity heuristic is exact — which is precisely why these sweeps stayed green while issue NVIDIA#613 (interleaved K/V views halving the TMA extent) shipped and had to be found through an external integration. Fix: the ragged branch leaves Q/K/V/O strides None and __call__ ends with fill_derived_fields() — one source of truth for the gap draw and its auto-packed fallbacks (cu / offset-multiplier forms NVIDIA#538, 1-byte dtypes NVIDIA#537). The head_major stats stride and the whole dense branch are untouched. Census over the fwd ragged L0 slice (84 configs): before, 0/84 drew a gap although each config's own rng_geom_seed hand-draws nonzero gaps; after, 84/84 draw gaps and ALL 84 would have failed under the old capacity formula. Verified on SM100 (cuDNN 9.26.0.33, CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1): with the NVIDIA#613 fix the gapped fwd ragged L0 slice passes 84/84 (all frost-served) — with the pre-fix adapter swapped in it fails 80/84, i.e. this wiring alone would have caught NVIDIA#613 the day the heuristic merged. bwd ragged L0 slice 158/158, identical to the unwired control on the same lib (the backend serves every gapped gradient combination); ragged_unified_L1 24/24 and offset_multiplier_unified_L1 24/24 (cu / mult forms stay packed via the existing fallbacks — 20/20 each in the offline census); the stride-override unit test still passes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * feat(sdpa): accept max_total_seq_len_q/kv on the forward SDPA node `sdpa_backward` has taken `max_total_seq_len_q/kv` since cuDNN 9.6; the forward node never did. That asymmetry is the root of a whole bug class. A ragged (THD) graph declares `(B, H, S_max, D)` plus a device-side ragged-offset tensor, so the packed token total is not expressible anywhere in the forward graph — and reading `cu_seqlens[-1]` host-side is exactly the D2H sync the zero-host-read THD execute (NVIDIA#552) exists to eliminate. The FROST forward path therefore has to INFER an upper bound on the token axis from the bound buffers' element span (NVIDIA#613/NVIDIA#706). That bound is memory-safe but loose, and looseness is not benign: rows between the real total and the extent are masked yet still multiplied (`P == 0` times V), so they must be FINITE. A caller that over-allocates and leaves the tail unwritten poisons whole tiles through `0 * NaN` (NVIDIA#624). Every framework already has this number — it is `q.shape[0]` in vLLM, SGLang, TransformerEngine, Megatron-Core, PyTorch and FlashInfer alike — and today it gets thrown away at the graph boundary. This lets callers declare it. - C++: `max_total_seq_len_q/kv` on `SDPA_attributes` with setters and serialization, mirroring `SDPA_backward_attributes`. Frontend-side only: like the backward twin it is never lowered to a backend attribute, so it cannot affect backend validation (NVIDIA#704). - Forward node validation rejects it on a non-ragged layout, mirroring backward's "only supported with packed layout". - pybind: `sdpa(..., max_total_seq_len_q=None, max_total_seq_len_kv=None)`. - FROST forward consumes it: the declared total is min'd against the buffer-derived capacity, so it can only TIGHTEN the extent, never widen it. A stale or wrong value cannot make a launch address memory the caller does not own — it can only make it address less. Both the SM100 f16 and the SM120/FP8 extent sites go through one helper. Effect on NVIDIA#624, measured on SM100 (bf16, cuDNN 9.26.0.33, FROST forced), `seq_lens=[200,150,47]` (total 397) bound into `(640, H, D)` buffers whose `[397, 640)` tail is NaN — only the tail fill differs between runs: undeclared: 201,728 NaNs in O (49.6%) declared: 0 NaNs, bit-identical to the zero-tail run Verified: new L0 regression test (asserts the clamp AND that the undeclared path still reaches the tail, so it tests the clamp rather than a benign shape); dense graph + attribute correctly rejected; the NVIDIA#613 interleaved-KV-views test and the gap-wired ragged L0 slice (84/84, all FROST-served) unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * feat(sdpa): expose max_total_seq_len_q/kv on sdpa_fp8 too Review follow-up. `PyGraph::sdpa_fp8` routes through `sdpa_internal`, so it already builds the same `SDPA_attributes` that now carries the packed totals -- only the entry point was missing them, and it hard-coded `py::none()` at the forwarding call. An FP8 THD caller therefore had no way to declare its totals even though the adapter side (`_thd_declared_total` at the SM100 f16 and SM120/FP8 extent sites) was already wired for them. Adds the two optional arguments to the declaration, the definition, the pybind binding and the docstring, and forwards them instead of `py::none()`. `sdpa_mxfp8` is deliberately left out: it does not go through `sdpa_internal` and builds `SDPA_fp8_attributes`, which has no such field, so covering it means extending that struct as well. Note the reviewer's stated motivation does not actually hold for FP8: the FP8/MXFP8 kernels already clamp their K/V descriptor extents to `cu_k[B]` device-side in `build_thd_meta_o_kv_descs_kernel`, so an unwritten K/V capacity tail is already TMA-unreachable there, and Q is the parallel dimension (a garbage Q row poisons only its own row, which is never stored). The change is still worth making for API symmetry and for exact rather than inferred extents. Test: `test_fp8_thd_declared_totals` runs the THD FP8 path with and without the declaration from the same seed and asserts O is bit-identical, plus the usual accuracy check against the reference. Verified: `test_sdpa_fwd_fp8_sm100.py` 61 passed; f16 THD tests 195 passed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * feat(sdpa): expose max_total_seq_len_q/kv on sdpa_mxfp8 too Correcting my own note on the previous commit: I claimed `sdpa_mxfp8` was out of scope because it "builds `SDPA_fp8_attributes`, which has no such field". That is wrong — `SDPA_fp8_attributes` is a type ALIAS for `SDPA_attributes` (graph_properties.h), so the field has been there all along and the only gap was the pybind entry point. `sdpa_mxfp8` does not route through `sdpa_internal`, so it needed its own declaration, definition, attribute plumbing, binding and docstring — but no struct change. The MXFP8 forward row serves THD (`thd_d_shapes` covers the d128 kernel), and the adapter side (`_thd_declared_total`) was already shared, so this completes the forward family: `sdpa`, `sdpa_fp8` and `sdpa_mxfp8` all now accept the packed totals. Still missing, and genuinely needing a struct change: the FP8/MXFP8 BACKWARD nodes. `SDPA_fp8_backward_attributes` is a distinct class (not an alias) with no such field, so `sdpa_fp8_backward` / `sdpa_mxfp8_backward` cannot take the totals while plain `sdpa_backward` has since cuDNN 9.6. Tracked separately. Test: `test_mxfp8_thd_declared_totals` runs the MXFP8 THD path with and without the declaration from the same seed and asserts O is bit-identical, plus the usual accuracy and amax checks. Verified: `test_sdpa_fwd_mxfp8_sm100.py` THD selection 10 passed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
…r extent choose_split_kv computed a split and then nothing used it. Two defects, both on the delivery path rather than in the cost model: - _split_points returned [no_split, chosen], so the chosen split landed at plan[1]. build_plans() stops at the first entry that builds and execute() runs _plan_index, so a plain build ALWAYS took the unsplit plan; the split was reachable only through select_plan or an ALL-policy autotune. Return [chosen, no_split] instead — the split leads, and no-split stays reachable behind it. - The model was fed rows_per_tile = tile_m * cga, but an SM100 d128 cluster covers TILES_Q * TILE_M * CTA_MMA Q rows on its CTA pair — twice that. The doubled tile count reads a half-empty machine as full, so the chooser under-splits or declines to split at all. Use _pack_gqa_tile_q, the helper that already answers "Q rows one grid tile covers", and the same extent every test in test_split_kv_heuristic.py already assumed. Flipping the lead moved the split into the base knob set, which exposed a third: the "a split set rides the plain scheduler" coupling lived only in the splits[1:] runner-up loop, so a LEADING split inherited the derived LPT_L2 policy on causal graphs — unbuildable on SM120, which raises on split_kv > 1 under an LPT remap. The coupling now binds whichever leg leads, and scheduler runners ride an unsplit leg. The chooser itself is unchanged, so a grid that already fills the machine still does not split. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
split_kv had two lists. choose_split_kv scored an implicit power-of-two ladder
bounded by _SPLIT_KV_MAX, while _split_points projected the winner onto
caps.split_kvs and returned usable[-1]. They agreed only because {1,2,4} was a
prefix of the ladder; on any other domain the returned split was one the model
had never scored.
Separate the two roles that field was sharing. Capabilities.split_kvs becomes
split_kv_supported, a boolean gate on whether the row wires the split path at
all — mismatch() checks it in the block that already special-cases
split_kv > 1, rather than in the uniform domain table, and it imposes no upper
bound because the kernels have none. WHICH splits are worth scoring becomes
split_kv_candidates(sm_count, kv_tiles): powers of two up to
2**ceil(log2(sm_count)), bounded by kv_tiles // _SPLIT_KV_MIN_TILES.
choose_split_kv loops exactly that list; _SPLIT_KV_MAX, max_split and the
usable[-1] snap are gone.
A split launches two kernels, so cost(s) is now two latencies summed:
cost(s) = waves(s) * (ceil(kv_tiles/s) + CTA_COST)
+ combine_waves * (s * COMBINE_COST)
combine_waves = ceil(S_q*H_q*B / sm_count), because split_combine_sm100's grid
is (S_q, H, B) — one block per output row, independent of s; only the per-block
work grows with s. Both terms are latency, so the combine cannot double-count
the parallelism the wave factor already divided out. Without it s reached the
model only through waves(s), a step function, leaving a larger split free
between wave boundaries.
_B300_FIT is re-measured for the new model.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
choose_split_kv gained a required combine_rows keyword-only argument for the combine-pass cost term; the sm120 expected-split helper was not updated to match, unlike its sm100 sibling in test_sdpa_fwd_split_kv_sm100.py.
This was referenced Aug 26, 2026
Merged
Aneureka
pushed a commit
to NVIDIA/cudnn-frontend
that referenced
this pull request
Aug 26, 2026
…758) choose_split_kv gained a required combine_rows keyword-only argument for the combine-pass cost term (#720), but the sm120 expected-split helper wasn't updated to match, unlike its sm100 sibling in test_sdpa_fwd_split_kv_sm100.py which already passes it. The squash-merge of #720 dropped the follow-up fix (yanzhuo607#1), so develop's sm120 CI (frost:rel:sdpa:sm120) is broken again: TypeError: choose_split_kv() missing 1 required keyword-only argument: 'combine_rows' Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
yanzhuo607
pushed a commit
that referenced
this pull request
Sep 8, 2026
…ks, GQA, any seqlen (NVIDIA#887) * feat(sdpa): SM100 d512 backward — GQA, causal/SWA/band masks, engine Adds the sdpa_bwd_sm100 engine and the three-stage chain behind it, reachable from the graph API (manifest slot 2). stage 1 do_dot reused unchanged from bprop_chain_f16_sm120 stage 2 S / dS bprop_d512_f16_sm100.py (cga4x1 role split) stage 3 dV / dK / dQ bprop_matmul_sm100.py (2-D (b,h) batch, 4-D TMA) Coverage: d in (256, 512] any multiple of 8 (envelope-served on 512 tiles via TMA zero-fill), f16/bf16, GQA/MQA, dense + causal top-left/bottom-right + SWA + right-band widening, BSHD and any dense layout (non-conforming io is staged through the workspace). Three bugs found and fixed while widening the masks, each of which had passed a narrower test first: * the kv loop split into TWO ranges, not three. `unmasked_lo == left` under causal-only, so the missing low-edge range only surfaced under SWA, where it dropped the band's left boundary. * the stage-3 K-trim ignored the band shift. Band widening and bottom-right both push the non-zero band past kv <= q, so the trim cut away data stage 2 had written. Threaded through as `causal_shift`. * an EMPTY stage-3 K range still ran the epilogue, storing an accumulator that `scale_d=False` never initialised — uninitialised TMEM in the output. The range is now clamped to one k-tile and the workspace zeroed under any trim, so those structurally-masked rows read a real zero. This presented as a 1-in-6 catastrophic dK/dV (cos 0.0006), and causal_tl at S_kv > S_q was passing on luck. Also: widen the frontend's Blackwell backward head-dim gate to admit the (256, 512] band, mirroring the forward precedent in sdpa_support_surface.h; make the benchmark's SOL check say when it cannot compute a peak instead of silently omitting it, and shout above 100%. Validated on Blackwell (cc 10.0, cuDNN 9.27): 23 mask x shape combos, 5 GQA ratios, 21 head dims, the graph API, the host chunk loop, and benchmark_single_sdpa --sdpa_backend cudnn_oss. Regression: forward SM100 502 passed, GEMM 4453 passed. Not yet: padding masks (needs the per-batch kv length), THD, FP8, pytest coverage, perf tuning. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat(sdpa): SM100 d512 backward — any seqlen, envelope floor, pytest coverage Completes the capability set and adds the test module. * Arbitrary S_q / S_kv, not just tile multiples. The engine rounds the COMPILE shape up (256 in q, 128 in kv), stage 2 computes the tail and masks it, and stage 3 gets a real-extent slice so the padding never reaches a GEMM's M/N/K. This is the uniform length only; a per-batch seq_len is still declined. * Capabilities gains `d_envelope_floor`. An envelope with no floor silently claimed every small head dim and padded it onto the 512-wide kernel — a d=128 graph would have taken the whole 512 MMA with a d256 flavor sitting right there. Only the adapter's check_support caught it, i.e. after eligibility, so the row was lying. Defaults 0, so the sm120/sm80 rows are unchanged. * test_sdpa_bwd_dsl_sm100.py: 33 tests, accept AND reject for every claim. The reject tests ask the row's own mismatch() rather than walking the ranked plan list, which is confounded by the backend (it serves d=128/256 backward and raises outright on some graphs). Two bugs the tests found: * do_dot writes `delta` with a row stride of ceil(S_q/128)*128, not S_q, so the buffer and the stage-2 view of it must use the same rounding. They coincide at every tile-multiple S_q, which is why only a ragged length exposed it, and the signature was exact: dV correct, dQ and dK wrong, because dV is the one gradient that does not consume do_dot. * the envelope floor above — found by the reject test, not by review. Validated on Blackwell (cc 10.0, cuDNN 9.27): 33/33 pytest x3 runs, plus 39 end-to-end checks over masks x shapes (including 500/500, 300/200, 257/129), GQA ratios, head dims and the host chunk loop. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * perf(sdpa): drop padding from the SM100 bwd workspace memset Stage 2 already writes zeros across the padded region: its kv bound is div_up(REAL S_kv, TILE_N), so the tail tile IS visited and apply_mask_chunk zeroes the columns past the real length, and the padded q rows are visited too (the grid is sized on the rounded S_q) with row_scale zeroing them. Only a MASK-SKIPPED tile is genuinely never written, so only that case needs the memset. Measured on the remaining (mask) case at B1 H128 S8192 d512 causal, 20 iters x 3: 36.6 ms with the memset vs 35.5 without, i.e. ~3% -- not the ~5 ms a bandwidth estimate suggests. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * perf(sdpa): stage-3 GEMM to the 2x2 cluster config on the d512 SM100 backward Ports the backward's stage-3 GEMM from cluster 2x1 (cgrp 256x256x64) to 2x2 (cgrp 512x512x64, cta 256x128x64). Measured at B=1 H=128 S=8192 d=512 bf16, paired alternating A/B/B/A so clock drift cannot pick a winner: no_mask 732.3 vs 707.8 TFLOPS +3.5% (4/4 rounds) causal 653.5 vs 606.0 TFLOPS +7.8% (2/2 rounds) Per-GEMM, sustained over 300 iters, the three products are now within 1% of each other at ~41.4% of the boost-clock SOL bound (was 40.1%). An earlier isolated-GEMM comparison read this as a coin flip; it was measuring the GEMM out of context and is superseded by the paired end-to-end numbers above. The config change is not a pure constants swap. Beyond the tile/cluster/stage constants it needs the epilogue's SMEM staging widened (shape 4096 -> 8192, tidx*32 -> tidx*64, alignment 64 -> 128, Swizzle(2,4,3) -> Swizzle(3,4,3)) AND the matching output TMA descriptor (box_dims 32 -> 64, s64b -> s128b). Missing only the descriptor half leaves the epilogue writing a 128 B swizzle through a 64 B descriptor: no crash, no launch failure, just garbled stores (cos ~0.006 on all three GEMMs). `fallback_cluster_shape_mnk` is pinned to None rather than the (2,1,1) upstream renders: `_host` always sizes the grid as a multiple of the preferred cluster, so the mixed-CGA fallback is unreachable here, and its paths carry coordinate and multicast logic the 2-D (b, h) batch rewrite has never exercised. Causal: the wider M tile breaks the tight-trim invariant `cgrp_tile_mnk[0] <= causal_gran` (512 > 256), so a 512-row M tile straddles two 256-row stage-2 write blocks and no per-tile K range can exclude the skipped region. The trim becomes an optimization only, and `_zero_ws` -- which `api_dsl` already sets whenever the trim is active -- is what makes the causal path correct. Both docstrings previously claimed the trim was a correctness bound; corrected, with the reason the zero-fill must not be narrowed later. `validate_bprop_causal_chain.py` now poisons only when the invariant holds and otherwise models production's zero-fill, so it tests the real contract instead of failing a config that is correct. Validated: 33/33 test_sdpa_bwd_dsl_sm100.py; stage-3 GEMMs across 5 shapes; non-causal and causal chains at two shapes each. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(sdpa): honor FP16 in the SM100 d512 backward stage-3 GEMM The `sdpa_bwd_sm100` Capabilities row claims `{HALF, BFLOAT16}`, and the adapter picks a stage-2 template per dtype -- but the stage-3 GEMM template hardcoded `cutlass.BFloat16` in seven places, including the `make_fake_compact_tensor` that fixes the compiled artifact's signature. So every FP16 graph the row accepted was admitted, compiled, and then died at EXECUTE with: ValueError: Mismatched Tensor on argument #1 when calling `_host(...)`, expected dtype=bfloat16 All 33 existing tests passed over it because all 33 ran bf16. `MatmulTemplateParams` gains `dtype_qkv`, threaded from the same `dtype_code` stage 2 already uses -- the two must agree, since stage 3 reads back the S/dS workspace stage 2 wrote. BF16 and FP16 are both 2 B/element and both take `Tcgen05MMAKind.F16`, so this is a token swap: every byte-sized constant (swizzle, box dims, SMEM staging, k-steps) is width-driven and unchanged. The output descriptor's stride shift now derives from `cd_dtype.width` like the A/B descriptors already did, instead of a literal 16. `validate_matmul_params` is the stage-3 backstop, mirroring what `make_cfg_d512` does for stage 2: ValueError, never assert, and reaching it means the Capabilities row admitted something it should not have. Tests (+5, 33 -> 38): * `test_dense` and a new `test_causal_dtypes` are parametrized over both dtypes. A frozenset capability field is that many separate claims; one accept test for the set leaves the other member unasserted. The fp16 cases were seen RED against the unfixed kernel -- that is what found this. * `test_reject_thd` and `test_reject_padding_mask` assert the decline on a REAL graph rather than only on the Capabilities field, so they exercise the analyzer too. Both INVERT when THD / per-batch seq_len land. No Capabilities field changed, so the support matrix is unaffected (Rule S2) -- the row's `{HALF, BFLOAT16}` claim is simply true now. Validated on Blackwell (cc 10.0, cuDNN 9.27): 38/38 x 3 consecutive runs. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(sdpa): default attn_scale on the SM100 backward; two more reject tests Review findings on NVIDIA#887. **attn_scale default (real bug).** `attn_scale` is optional on the graph, so `scale_softmax` arrives None when a caller omits it. The SM120 and SM80 adapters default it to `1/sqrt(d_qk)`; the SM100 adapter did not, so a graph built without `attn_scale` was admitted by the Capabilities row, passed `check_support()`, and then died in execute with TypeError: unsupported operand type(s) for *: 'NoneType' and 'float' `test_default_attn_scale` covers it, and was seen RED against the unfixed adapter (it compares against a reference that always uses 1/sqrt(d), so a wrong default fails the comparison rather than merely not raising). **test_reject_rectangular_head_dims.** `mismatch()` already declines `d_qk != d_v` for this row -- verified, no code change -- but nothing asserted it, and the C++ node validation's d512 exception is deliberately permissive about the pair, so that decline is the only thing keeping a rectangular graph away from an adapter that raises on it. Locked in. **benchmark SOL message.** `_peak_mma_tflops` can be unavailable for two different reasons -- no modelled dense-MMA rate for the dtype/arch, or no sampled clock -- and the note blamed a missing `pynvml` in both cases, sending the reader after the wrong one for e.g. `--data_type float`. It now names which. Two further review comments were checked and NOT acted on: * clang-format on `scaled_dot_product_flash_attention.h` -- the pre-commit CI job (clang-format 21, the repo's pinned hook) passes on this diff. * moving the parametrized sweeps off `L0` -- the sibling SM120 backward suite runs 119 L0 cases; this module is 40 and runs in ~52 s. Validated on Blackwell (cc 10.0, cuDNN 9.27): 40/40. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(sdpa): an unserved d512 backward must decline as GRAPH_NOT_SUPPORTED Widening `test_sdpa_random_bwd_L0`'s head-dim sweep to 512 turned 91 of its 384 cases red (baseline: 289 passed, 0 failed). None of them was a numerics bug -- all 91 were this branch reporting "nobody serves this graph" as the WRONG KIND of error. Chain: admitting d in (256, 512] at the node's validation (so the frontend-only FROST engine can claim the band) means such a graph now reaches plan creation. When the caller also asks for something no engine serves -- deterministic, say, which `sdpa_bwd_sm100` declines and the backend has no plan for in this band -- `override_heuristics_query()` pins a backend engine id (5 / 17) and bypasses the heuristics query entirely. The pinned engine config then fails to finalize with `CUDNN_STATUS_NOT_SUPPORTED`, which `_CUDNN_CHECK_CUDNN_ERROR` folds into `CUDNN_BACKEND_API_FAILED` -- reaching Python as a bare `RuntimeError`. The type is load-bearing, not cosmetic: every SDPA harness here skips on `cudnnGraphNotSupportedError` and FAILS on anything else, so "unsupported" became "the backend broke" for any caller trying to feature-detect. Two fixes, narrow then general: * `override_heuristics_query()` returns `{-1, {}}` for the band. That override exists to pin a specific backend bprop kernel, and the backend has no plan in (256, 512] at all, so pinning one there cannot succeed. Declining it lets the heuristics query run, return no configs, and report GRAPH_NOT_SUPPORTED -- which is already what happened for band graphs that did NOT trip the override. * `create_engine_config` maps a finalize `CUDNN_STATUS_NOT_SUPPORTED` to `GRAPH_NOT_SUPPORTED` instead of `CUDNN_BACKEND_API_FAILED`. Any other status stays an API failure. NOT_SUPPORTED from the backend means precisely that this engine cannot serve this operation graph, at every call site, so folding it into a generic API error loses the only distinction a caller can act on. `test_unserved_band_graph_declines_as_not_supported` asserts the exception TYPE (not its message) and was seen RED against the unfixed headers, reproducing the original `RuntimeError: ... CUDNN_STATUS_NOT_SUPPORTED` exactly. `test_sdpa_random_bwd_L0` then gets the widening this was blocking: `d_qk_max`/`d_v_max` 256 -> 512 and `(512, 512)` added to the high-probability list. Note the coverage trade: the generator now spends part of its 384-case budget on a band that this suite cannot execute (FROST engines are opt-in and it does not enable them), so those cases SKIP -- 200 passed / 184 skipped, against 289 / 95 before. It exercises the frontend admission and decline path for the band; the executing coverage lives in `sdpa/frost/test_sdpa_bwd_dsl_sm100.py`. Validated on Blackwell (cc 10.0, cuDNN 9.27): * `test_sdpa_random_bwd_L0` widened: 91 failed -> 0 failed. * full `test_mhas_v2.py` L0: 2090 passed, 796 skipped, 0 failed. * `test_sdpa_bwd_dsl_sm100.py`: 41/41. * C++ build (`-Werror -Wall -Wextra -Wpedantic`): clean. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
yanzhuo607
pushed a commit
that referenced
this pull request
Sep 14, 2026
…e binds fp4 and scales as uint8, as FlashInfer does (NVIDIA#1035) * pygraph: the declaration supplies the dtype too; the conformance suite binds fp4 and scales as uint8, as FlashInfer does FlashInfer's cuDNN fp4 GEMM binds the packed fp4 data as uint8 and, with it, the e4m3 scale blob viewed as uint8 (gemm_base: a_descale.view( torch.uint8) when a.dtype == uint8). The conformance suite bound torch's float4_e2m1fn_x2 instead, so it passed where FlashInfer's own tests failed on the frost plan: the pack carried the buffer's uint8 into the kernel, which was compiled for float4_e2m1fnx2 ("Mismatched Tensor on argument #1"). With frost off, and with the stock frontend, the same tests pass -- the rule was incomplete, not the kernel. The declaration is the contract, dtype included: describe_from now reads a buffer whose storage slots are as wide as the declaration's AS the declared dtype (spelled the way torch exports it: fp4 is code 17, 4 bits, 2 lanes), whether or not its extents are re-described; a narrower or wider buffer keeps its own dtype and is never re-described. A bare address gets the same spelling (set_operand takes lanes). The suite's fp4 cases bind uint8 by default and keep one x2-view case, so both spellings of the same bytes are proven to run. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * variant pack: an fp4 slot names its dtype float4_e2m1fn_x2 and is one byte wide OperandBuffer.dtype fell through to "code17_4" and element_size() (bits / 8) was 0 for a 4-bit, 2-lane slot, so nbytes was 0 for a non-empty fp4 buffer. Both now follow the storage slot, as describe_from and read_buffer_extent already did. (CodeRabbit review on NVIDIA#1035.) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * variant pack: the fp4 x2 name needs two lanes; the doc names every buffer that keeps its own dtype dtype_name spelled (17, 4) as float4_e2m1fn_x2 whatever the lane count; a public set_operand / override_operand call with lanes=1 is not torch's x2 storage and now reads as its raw code with an _x<lanes> suffix when lanes differ from one. The doc said only an undersized buffer keeps its own dtype; a buffer whose slots are not as wide as the declaration's does too. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
|
This PR has been inactive for 3 weeks and is now marked stale. It will be closed in 3 more weeks unless there is new activity (a comment, a push, or removal of the stale label). |
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
13 participants
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.
Rebases yanzhuoc/split-kv-primary onto latest develop (was 3 commits behind; rebase was clean, no conflicts) and fixes the resulting sm120 CI failures.
sm120 failures
frost:rel:sdpa:sm120(job 412040217) failed 6 tests intest/python/sdpa/frost/test_sdpa_fwd_split_kv_sm120.py::_expected_splitwith:choose_split_kvgained a requiredcombine_rowskwarg on this branch (combine-pass cost term). The sm120 test's_expected_splithelper wasn't updated to pass it, unlike its sm100 sibling intest_sdpa_fwd_split_kv_sm100.py(which already doescombine_rows=b * h_q * s_q). This mirrors that same fix for the sm120 helper:combine_rows=api.batch_size * api.h_q * api.s_q_max.🤖 Generated with Claude Code