test(sdpa): fuzz per-tensor ragged token-stride gaps in test_mhas_v2 by vedaanta · Pull Request #516 · NVIDIA/cudnn-frontend

vedaanta · 2026-08-07T22:14:11Z

What

Ragged Q/K/V/O were always generated with canonical packed strides (token stride == h*d): the gaps stride machinery in random_config.py was dense-only, and #462 varies only the STATS layout. So no sweep could ever produce a tensor like a K/V view of a kv-interleaved [T, 2, H, D] buffer (token stride 2*h*d) — the exact layout torch.nn.attention.varlen users get by slicing a fused KV projection (PyTorch's own varlen suite uses it in AttentionBlock).

with_ragged_token_gap, enabled unconditionally in the fp16 ragged sweeps (no True/False dice — the packed case comes from the gap draw itself): each of Q/K/V/O independently draws a token stride of 1–4 whole tokens (gap = n*h*d, n ∈ 0..3, seeded from rng_geom_seed; deterministic through serialize/deserialize repro, explicit strides still win). n=0 is plain packed, n=1 exactly a kv-interleaved [T, 2, H, D] view, n=2 a [T, 3, H, D] QKV-interleave. Whole-token gaps keep every ragged base address in the packed layout's alignment class by construction (sub-token gaps can violate the graph API's 16-byte pointer-alignment contract — we verified an odd-element gap is illegal for every engine). The config field defaults False: fixed configs and the fp8/mxfp8 harnesses still assume packed allocations, and the cu_ragged form derives offsets internally. Harness: ragged buffers (incl. gradients) are allocated with the configured strides, and ragged offsets scale by each tensor's actual stride[2] — the same generalization #462 made for stats offsets.

Known failures this exposes (intentionally NOT masked)

The FROST THD forward engines (sdpa_fwd_prefill_sm120, sdpa_fwd_prefill_sm100_d128) claim non-packed-stride THD graphs and silently mis-address them — 100 % of O wrong (the stride ORDER is still BSHD, so the order-only layout gate passes; the THD lowerings rebuild packed views). Verified on cc 10.0 and RTX 5080 (sm120), FE develop 3f17f5b, torch nightly. Backend engines serve every gapped combination correctly, fwd and bwd (5/5 seeded repros, a bf16 training config, and a 128-test slice of the fwd ragged L0 sweep).

Per review direction, this PR ships the coverage only — no engine-side decline — so the failing configs under CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 are the standing repro set for fixing the THD lowerings to honor declared strides.

Deterministic repro (fails on FROST, passes on backend; no-gap configs unchanged on both routes):

cd test/python && CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 pytest test_mhas_v2.py::test_repro --repro "{'data_type': 'torch.float16', 'diag_align': 'cudnn.diagonal_alignment.TOP_LEFT', 'batches': 2, 'h_q': 8, 'h_k': 8, 'h_v': 8, 'd_qk': 128, 'd_v': 128, 's_q': 256, 's_kv': 256, 'seq_len_q': [192, 256], 'seq_len_kv': [192, 256], 'is_ragged': True, 'is_padding': True, 'is_infer': True, 'with_ragged_token_gap': True, 'rng_geom_seed': 1, 'rng_data_seed': 2}"

Found while integrating PyTorch's varlen_attn onto the cuDNN python API (caught by PyTorch's upstream varlen test suite).

🤖 Generated with Claude Code

Summary by CodeRabbit

Bug Fixes
- Improved handling of non-contiguous token layouts in ragged attention operations.
- Preserved configured spacing between tokens across forward and backward computations.
- Improved correctness when query, key, value, and output tensors use independent token gaps.
Tests
- Expanded randomized coverage for ragged inputs with token gaps.
- Added regression coverage for explicit stride overrides and packed-stride fallback scenarios.

coderabbitai · 2026-08-07T22:14:30Z

📝 Walkthrough

Walkthrough

Ragged SDPA configurations now generate independent Q/K/V/O token-stride gaps with defined fallback cases. Ragged allocations, gradients, and offsets use the configured strides. Regression tests cover explicit stride overrides and packed-stride fallbacks.

Changes

Ragged SDPA stride handling

Layer / File(s)	Summary
Ragged stride configuration `test/python/sdpa/random_config.py`	`compute_packed_strides` accepts token gaps. `ExecConfig` enables ragged token gaps by default. Derived Q/K/V/O strides use deterministic per-tensor gaps when eligible and packed strides for defined fallback cases.
Ragged stride allocation and validation `test/python/sdpa/fp16.py`, `test/python/test_mhas_v2.py`	Ragged tensors, gradients, and offsets use actual token strides. Regression tests validate default gaps, explicit stride overrides, and packed-stride fallbacks.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related issues

test(sdpa): fuzz head-axis stride gaps in the ragged sweeps #540: Extends related ragged SDPA stride-gap coverage to head-axis gaps.

Possibly related PRs

NVIDIA/cudnn-frontend#526: Focuses on preserving and using declared non-compact ragged token strides in SDPA tensors and kernels.

Suggested reviewers: anerudhan

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the SDPA test change and the addition of per-tensor ragged token-stride gap coverage.
Description check	✅ Passed	The description provides a detailed summary, rationale, compatibility impact, known failures, and reproducible testing information, but it omits the repository template headings and checklist.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@test/python/sdpa/random_config.py`:
- Line 77: Update the shape unpacking in the affected function to avoid binding
the unused b value, while preserving the existing h, s, and d bindings used by
the function.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 825b6793-276e-417f-a1b4-ed1a66b4b956

📥 Commits

Reviewing files that changed from the base of the PR and between 3f17f5b and 1552036.

📒 Files selected for processing (5)

python/cudnn/sdpa/fwd/engines.py
python/cudnn/sdpa/graph_analyzer.py
test/python/sdpa/fp16.py
test/python/sdpa/random_config.py
test/python/test_mhas_v2.py

coderabbitai

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

test/python/sdpa/random_config.py (1)
211-233: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Keep gap assignment stable for partially specified strides.

_gap_rng is consumed only when gap_fn runs. Therefore, setting stride_q explicitly shifts the generated gaps for K/V/O, even when rng_geom_seed is unchanged. This makes per-tensor layouts depend on which earlier strides were overridden.

Draw all four gap values before the None checks, then assign fixed values to Q/K/V/O. Add a regression test for a configuration with only stride_q specified.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@test/python/sdpa/random_config.py` around lines 211 - 233, The gap RNG in the
stride initialization flow must not depend on which strides are explicitly
provided. In the _gapped/gap_fn logic, draw four gap values for Q, K, V, and O
before the stride_q/stride_k/stride_v/stride_o None checks, then use those fixed
values only for missing strides while preserving explicit values. Add a
regression test covering a configuration with only stride_q specified and verify
the generated K/V/O layouts remain stable.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@test/python/sdpa/random_config.py`:
- Around line 211-233: The gap RNG in the stride initialization flow must not
depend on which strides are explicitly provided. In the _gapped/gap_fn logic,
draw four gap values for Q, K, V, and O before the
stride_q/stride_k/stride_v/stride_o None checks, then use those fixed values
only for missing strides while preserving explicit values. Add a regression test
covering a configuration with only stride_q specified and verify the generated
K/V/O layouts remain stable.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8536b67d-a680-4967-b103-b71bcbc208ac

📥 Commits

Reviewing files that changed from the base of the PR and between 63ad9a6 and 60b9a43.

📒 Files selected for processing (2)

test/python/sdpa/random_config.py
test/python/test_mhas_v2.py

🚧 Files skipped from review as they are similar to previous changes (1)

test/python/test_mhas_v2.py

The FROST THD forward lowerings rebuilt packed (1, T, H, D) views with hardcoded (t*h*d, h*d, d, 1) strides, so a THD tensor declaring a wider token stride — e.g. a K/V view of a kv-interleaved [T, 2, H, D] buffer (token stride 2*h*d), the layout torch.nn.attention.varlen users get by slicing a fused KV projection — was claimed and silently mis-addressed (100% of O wrong on both sdpa_fwd_prefill_sm120 and sdpa_fwd_prefill_sm100_d128; backend engines serve the same graphs correctly). Fix: adapter-level normalization, the THD analogue of the dense path's compact-BSHD normalization. New SdpaFwdDsl helpers read each tensor's DECLARED (token, head, elem) strides from its descriptor; a packed declaration stays exactly the zero-copy view it was, and a non-packed one gathers into packed scratch carved from the workspace (Q/K/V) or scatters back from it after the launch (O). Worst-case scratch (b * s_max * h * d per non-packed tensor) is added to scratch_workspace_bytes, so per-execute allocations stay zero on the FROST dispatch path. The t_kv == 0 O-zeroing paths now also zero through the declared view. No kernel changes. Verified (torch nightly cu132, FE develop + PR NVIDIA#516's fuzz tests): - previously-failing gapped repros now pass WITH the frost engines serving, on cc 10.0 (sm100 engine) and RTX 5080 (sm120 engine) - 128-test slice of the fwd ragged L0 sweep passes under CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 - packed configs unchanged (zero-copy path untouched) Follow-ups (tracked, not in this change): - native kernel-side stride support (TMA descriptors / address math over declared strides) to remove the gather/scatter copies - fp8/mxfp8 THD: inherit the normalization when that path is wired; their test harness also does not honor configured strides yet - cu_seq_len (cu_ragged) form derives offsets internally assuming packed - future FROST THD backward engines must adopt the declared-stride contract from day one - head-axis stride gaps are honored by the gather view but not fuzzed Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Second stage on top of the adapter normalization: the kernels now express declared strides NATIVELY, so gapped THD views run zero-copy. - compile() on all five f16 kernels (sm120, sm100 d128/d192_d128/d256/ d512) takes optional caller-declared (batch, seq, head, elem) strides per tensor (lru cache-key); None keeps the compact specialization bit-for-bit. Strided fakes are built with make_fake_tensor, validated against the TMA 16-byte global-stride rule. - SM120: the K/V TMA layouts read the tensor's strides instead of recomputing packed ones (Q/O offset math was already layout-driven); the entry validator accepts padded, 16-byte-granular BSHD storage (compact = the equality special case). - SM100: the Q/K/V/O TMA descriptors are built from the tensor views, so declared strides flow in unchanged; the THD O-descriptor builder steps per-batch bases by O's declared seq-axis stride (o_tensor.stride[1]) instead of the packed n_qh*d_v. - Adapters bind declared-stride views natively when every stride is a 16-byte multiple (x8 elements at 2 bytes/elem) and fall back to the packed gather/scatter normalization otherwise — the normalization stays as the safety net for sub-granularity strides. Verified (torch nightly cu132, ToT develop + PR NVIDIA#516's fuzz tests): - gapped seeded repros pass with the frost engines serving natively on cc 10.0 (sm100 d128) and RTX 5080 (sm120) - 128-test fwd ragged L0 sweep slice: 128/128 on cc 10.0 (all four sm100 flavors reached), 84/84 slice on sm120 - ex-ops suite (incl. kv-interleaved views) 11/11 on both boxes - dense fwd L0 sweep slice: 182 passed / 8 skipped (dense compile paths pass no strides -> unchanged compact specialization) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…nels; decline what TMA cannot express A THD tensor may declare a wider token stride than the packed h*d — e.g. a K/V view of a kv-interleaved [T, 2, H, D] buffer (token stride 2*h*d), the layout torch.nn.attention.varlen users produce by slicing a fused KV projection. The THD lowerings rebuilt packed (1, T, H, D) views with hardcoded strides, so such graphs were claimed and silently mis-addressed (100% of O wrong on both sdpa_fwd_prefill_sm120 and the sm100 flavors; caught by PR NVIDIA#516's fuzz coverage and PyTorch's own varlen suite). Native support, no fallback (AGENTS Hard Rule 2): - compile() on all five f16 kernels (sm120, sm100 d128/d192_d128/d256/ d512) takes optional caller-declared (batch, seq, head, elem) strides per tensor (lru cache-key); None keeps the compact specialization bit-for-bit. Strided fakes via make_fake_tensor, validated against the TMA 16-byte global-stride rule. - SM120: kv_tma_desc reads the tensor's strides instead of recomputing packed ones (Q/O offset math was already layout-driven); the entry validator accepts padded 16-byte-granular BSHD storage (compact = the equality special case). - SM100: the Q/K/V/O TMA descriptors are built from the tensor views, so declared strides flow in unchanged; the THD O-descriptor builder steps per-batch bases by O's declared seq-axis stride (o_tensor.stride[1]). - Adapters bind declared-stride (1, T, H, D) views directly. What TMA cannot express is REJECTED in check_support (NotImplementedError naming the offending strides), so the Router falls back to an engine that honors the declaration: non-innermost-contiguous head dim, or token/head strides that are not multiples of 8 elements (sub- granularity strides also violate the graph API's pointer-alignment contract for the backend, so declining is correct, not conservative). - The SM120 FP8 THD path (NVIDIA#509) keeps the packed contract for now: non-packed declarations are declined (_thd_check_strides_packed); extending native strides there is tracked as a follow-up. Verified (torch nightly cu132, ToT develop + PR NVIDIA#516's fuzz tests): gapped seeded repros pass with the frost engines serving natively on cc 10.0 (sm100) and RTX 5080 (sm120); 128-test fwd ragged L0 sweep slice green on cc 10.0 (all four sm100 flavors) and 84-test slice on sm120; ex-ops suite incl. kv-interleaved views 11/11 on both; dense fwd slice 182 passed (dense compile paths pass no strides -> unchanged); packed THD configs bit-for-bit unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

@Aneureka

…6 fwd kernels (#526) * frost(sdpa): native THD stride support in the SM100/SM120 f16 fwd kernels; decline what TMA cannot express A THD tensor may declare a wider token stride than the packed h*d — e.g. a K/V view of a kv-interleaved [T, 2, H, D] buffer (token stride 2*h*d), the layout torch.nn.attention.varlen users produce by slicing a fused KV projection. The THD lowerings rebuilt packed (1, T, H, D) views with hardcoded strides, so such graphs were claimed and silently mis-addressed (100% of O wrong on both sdpa_fwd_prefill_sm120 and the sm100 flavors; caught by PR #516's fuzz coverage and PyTorch's own varlen suite). Native support, no fallback (AGENTS Hard Rule 2): - compile() on all five f16 kernels (sm120, sm100 d128/d192_d128/d256/ d512) takes optional caller-declared (batch, seq, head, elem) strides per tensor (lru cache-key); None keeps the compact specialization bit-for-bit. Strided fakes via make_fake_tensor, validated against the TMA 16-byte global-stride rule. - SM120: kv_tma_desc reads the tensor's strides instead of recomputing packed ones (Q/O offset math was already layout-driven); the entry validator accepts padded 16-byte-granular BSHD storage (compact = the equality special case). - SM100: the Q/K/V/O TMA descriptors are built from the tensor views, so declared strides flow in unchanged; the THD O-descriptor builder steps per-batch bases by O's declared seq-axis stride (o_tensor.stride[1]). - Adapters bind declared-stride (1, T, H, D) views directly. What TMA cannot express is REJECTED in check_support (NotImplementedError naming the offending strides), so the Router falls back to an engine that honors the declaration: non-innermost-contiguous head dim, or token/head strides that are not multiples of 8 elements (sub- granularity strides also violate the graph API's pointer-alignment contract for the backend, so declining is correct, not conservative). - The SM120 FP8 THD path (#509) keeps the packed contract for now: non-packed declarations are declined (_thd_check_strides_packed); extending native strides there is tracked as a follow-up. Verified (torch nightly cu132, ToT develop + PR #516's fuzz tests): gapped seeded repros pass with the frost engines serving natively on cc 10.0 (sm100) and RTX 5080 (sm120); 128-test fwd ragged L0 sweep slice green on cc 10.0 (all four sm100 flavors) and 84-test slice on sm120; ex-ops suite incl. kv-interleaved views 11/11 on both; dense fwd slice 182 passed (dense compile paths pass no strides -> unchanged); packed THD configs bit-for-bit unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * docs(agents): Hard Rule 2 — serve the declared layout natively or decline, never adapt Closes the loophole Rule 1's letter leaves open: adapter-side normalization copies that make an unsupported layout runnable. Workspace carving does not legitimize a data-tensor copy (the carve exemption is for metadata and dead-slot dummies), the dense path's grandfathered normalization is not a license for new ones, and whatever check_support accepts the kernel must address natively. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * frost(sdpa): review hardening — validate runtime THD buffers, decline overlapping strides - _thd_view validates the runtime buffer against its declaration before reinterpreting storage: dtype/device must match and the base address must be 16-byte aligned (TMA global-address rule / assumed_align=16); as_strided already rejects views past the underlying allocation. - _thd_check_strides_native additionally requires covering (non-overlapping) strides — head >= d, token >= heads*head — matching the SM120 kernel's is_layout_supported, so sub-dense declarations are declined at check_support instead of failing at the per-execute compile (or racing on O writes on SM100). - Kernel _fake_bshd guards: the head dim must be innermost-contiguous; d256/d512 validate the O stride at BPE_O (the O storage dtype byte size). - Clearer SM120 layout-rejection message (the entry validator accepts padded storage now; the text still demanded compact). The THD host-prep stream binding flagged in the same review round is a pre-existing issue (#476) and is split into a separate PR. Addresses CodeRabbit review feedback on #526. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * frost(sdpa): make the THD native-stride gate dtype-aware (16 // itemsize) The gate hardcoded the TMA 16-byte global-stride rule as 8 elements, the f16/bf16 case. It lives in the shared base class, so express the quantum in the tensor's own element units — 8 at 2 B/elem, 16 at 1 B/elem (fp8), 4 at 4 B/elem — per descriptor, so mixed-precision declarations check each tensor at its own dtype. No behavior change for the f16 paths this PR enables; the fp8 native-stride follow-up (#537) inherits the correct quantum for free. Suggested by @Aneureka in review. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Ragged Q/K/V/O were always generated with canonical packed strides (token stride == h*d): the gaps stride machinery was dense-only, and PR NVIDIA#462 varies only the STATS layout. So no sweep could ever produce a tensor like a K/V view of a kv-interleaved [T, 2, H, D] buffer (token stride 2*h*d) — the layout torch.nn.attention.varlen users get by slicing a fused KV projection. New with_ragged_token_gap knob in both ragged random sweeps: each of Q/K/V/O independently draws token stride h*d + gap, gap in {0, 8, 64, roundup8(h*d)} seeded from rng_geom_seed (deterministic through serialize/deserialize repro; explicit strides still win). Gaps stay multiples of 8 elements so ragged base addresses keep the packed layout's alignment class (the graph API requires 16-byte-aligned pointers). Harness: ragged buffers (incl. gradients) are allocated with the configured strides, and ragged offsets scale by each tensor's actual stride[2] — the same generalization NVIDIA#462 made for stats offsets. KNOWN FAILURES this exposes (intentionally not masked): the FROST THD forward engines (sdpa_fwd_prefill_sm120, sdpa_fwd_prefill_sm100_d128) claim non-packed-stride THD graphs and silently mis-address them (100% of O wrong; the stride ORDER is still BSHD, so the order-only layout gate passes). Backend engines serve every gapped combination correctly, fwd and bwd. The failing configs under CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 are the repro set for fixing the THD lowerings to honor declared strides. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…the fp16 ragged sweeps Per review: drop the True/False dice — the fp16 ragged sweeps enable the knob unconditionally, and the packed case comes from the gap draw itself. Each of Q/K/V/O independently draws gap = n*h*d tokens, n in 0..3 (n=0 packed, n=1 a kv-interleaved [T, 2, H, D] view, n=2 a [T, 3, H, D] QKV-interleave). Whole-token gaps keep every ragged base address in the packed layout's alignment class by construction, replacing the previous multiples-of-8-elements rule. The config field stays default-False: fixed configs and the fp8/mxfp8 harnesses still assume packed allocations, and the cu_ragged form derives offsets internally. Verified: backend engines pass 5/5 seeded repros, a bf16 training config, and a 128-test slice of the fwd ragged L0 sweep; FROST THD fwd engines keep failing gapped draws (the intended standing repro set). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…ding The gap RNG was consumed lazily, one draw per stride left unspecified — so explicitly pinning e.g. stride_q shifted the gaps K/V/O derive from the same rng_geom_seed. Draw all four values up front in fixed Q/K/V/O order and apply each only where the stride is missing: per-tensor layouts are now a function of the seed alone. The all-defaults path (the sweeps) consumes draws in the same order as before, so existing seeded repros reproduce identical strides. Adds a GPU-free regression test pinning stride_q and asserting the K/V/O strides match the all-defaults derivation (seed chosen so the old lazy behavior visibly shifts two of the three gaps). Also renames the unused batch binding in compute_packed_strides (RUF059). Addresses CodeRabbit review feedback on NVIDIA#516. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…t expressible Drop the per-sweep with_ragged_token_gap opt-in: the field now defaults True, so EVERY ragged config fuzzes its per-tensor token strides. fill_derived_fields falls back to packed automatically where a gap is not yet expressible or handled — cu / offset-multiplier forms bind offsets as cu (x multiplier) and cannot declare a token gap (NVIDIA#538), and the fp8/mxfp8 harnesses (1-byte data types) allocate assuming packed (NVIDIA#537) — so mixed sweeps (ragged + cu_ragged in one RandomChoice) gap exactly the draws that support it. Explicit strides are never touched (the gap only fills strides left None), so recorded repro dicts reproduce exactly. The regression test also locks in the new semantics: default-on for plain ragged, packed for the three fallback forms. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

vedaanta · 2026-08-11T07:10:51Z

Rebased onto the merged develop (now includes #526's native declared-stride support and #533's SM120 backward) and addressed the review findings:

Gap assignment stable under partial stride overrides (outside-diff finding) — fixed in 17be556. The gap RNG was consumed lazily, one draw per stride left unspecified, so pinning e.g. stride_q shifted the gaps K/V/O derive from the same rng_geom_seed. All four gaps are now drawn up front in fixed Q/K/V/O order and applied only where the stride is None — per-tensor layouts are a function of the seed alone. The all-defaults path consumes draws in the same order as before, so existing seeded repros reproduce identical strides. The requested regression test is included (test_ragged_token_gap_stable_under_stride_overrides, GPU-free): the seed is chosen so the old lazy behavior visibly shifts two of the three remaining gaps.

Unused b binding (RUF059) — fixed in 17be556.

Follow-up simplification (746ef17, per review discussion): the per-sweep with_ragged_token_gap=RandomChoice({True : 1}) opt-in is gone — the field defaults True, so every ragged config fuzzes its token strides, including sweeps that never opted in (e.g. fwd_ragged_unified_L1). fill_derived_fields auto-falls-back to packed exactly where a gap is not yet expressible or handled — cu / offset-multiplier forms (#538) and 1-byte (fp8/mxfp8) data types (#537) — so mixed sweeps (ragged + cu_ragged in one RandomChoice) gap only the draws that support it. Explicit strides are never touched, and recorded repro dicts serialize resolved strides, so replays are exact. The regression test locks in these semantics too.

Verification on SM100 (removing the sweep knob shifts the RNG draw sequence, so all sweep configs are a fresh sample): regression test green, fwd ragged slice 84/84, fp8 ragged slice green (auto-packed fallback confirmed); bwd ragged slice 125 passed / 87 skipped / 0 failed.

vedaanta · 2026-08-11T07:12:21Z

@cudnn-ci-bot run frost

cudnn-ci-bot · 2026-08-11T07:12:49Z

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-516-746ef17
Pipeline: 62096606
Targets: frost

coderabbitai

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@test/python/sdpa/random_config.py`:
- Around line 125-139: Prevent affected FROST forward engines from being
selected for configurations with non-packed ragged strides before enabling
with_ragged_token_gap by default. Add the eligibility decline in the
engine-selection or validation path used by fill_derived_fields, and preserve
packed-layout eligibility; if no such guard exists, change with_ragged_token_gap
back to False until the FROST fix is available.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ef9eaf0f-4613-4510-8826-715549a16668

📥 Commits

Reviewing files that changed from the base of the PR and between 60b9a43 and 746ef17.

📒 Files selected for processing (2)

test/python/sdpa/random_config.py
test/python/test_mhas_v2.py

coderabbitai · 2026-08-11T07:13:04Z

+    # Each ragged tensor (Q/K/V/O and gradients) independently draws a token
+    # stride of 1-4 whole tokens (gap = n*h*d, n in 0..3, seeded from
+    # rng_geom_seed): n=0 is the plain packed case, n=1 is exactly a view of
+    # an interleaved [T, 2, H, D] buffer (the layout
+    # torch.nn.attention.varlen users produce by slicing a fused KV
+    # projection), n=2 a [T, 3, H, D] QKV-interleave, and so on. Whole-token
+    # gaps keep every ragged base address in the packed layout's alignment
+    # class by construction (sub-token gaps can violate the graph API's
+    # 16-byte pointer-alignment contract — an odd-element gap is illegal for
+    # every engine). Default True: every ragged config fuzzes its layouts.
+    # fill_derived_fields auto-falls-back to packed where a gap is not yet
+    # expressible or handled: cu / offset-multiplier forms (#538) and the
+    # fp8/mxfp8 harnesses (#537). Configs with explicit strides are
+    # unaffected (the gap only fills strides left None).
+    with_ragged_token_gap: bool = True


🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift

Do not default-enable a layout that FROST executes incorrectly.

The PR objective states that FROST THD forward engines address non-packed strides incorrectly and return wrong O values. Line 139 enables these layouts for all applicable ragged configurations, but no engine-side decline prevents FROST selection.

Add an eligibility decline for affected FROST engines before this default is enabled. Otherwise, keep with_ragged_token_gap disabled by default until the engine fix is available.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/python/sdpa/random_config.py` around lines 125 - 139, Prevent affected FROST forward engines from being selected for configurations with non-packed ragged strides before enabling with_ragged_token_gap by default. Add the eligibility decline in the engine-selection or validation path used by fill_derived_fields, and preserve packed-layout eligibility; if no such guard exists, change with_ragged_token_gap back to False until the FROST fix is available.

Aneureka

LGTM.

… 9.26+) The stats (LSE) tensor stride was previously pinned to BHSD in the random test generator with a TODO comment, because the SM80 flash-bprop codegen used packed seq-stride addressing in the stats loads and silently produced wrong results for any other layout (NVBug 6057616). The backend fix landed in cuDNN dev as MR !4147 and will ship in 9.26. The stats layout is DRAWN unconditionally — so one seed derives identical shapes and Q/K/V/O strides on every backend version, preserving the layouts-are-a-function-of-the-seed-alone property NVIDIA#516 established — and APPLIED only when backend_version() >= 92600; older backends fall back to the packed BHSD default (old behavior), and full layout coverage engages automatically once CI deploys 9.26. Also wires the --implementation CLI option through the random fwd/bwd/bias L0 tests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

… in tests (#304) * test: randomize softmax-stats strides in bwd SDPA random tests (cuDNN 9.26+) The stats (LSE) tensor stride was previously pinned to BHSD in the random test generator with a TODO comment, because the SM80 flash-bprop codegen used packed seq-stride addressing in the stats loads and silently produced wrong results for any other layout (NVBug 6057616). The backend fix landed in cuDNN dev as MR !4147 and will ship in 9.26. The stats layout is DRAWN unconditionally — so one seed derives identical shapes and Q/K/V/O strides on every backend version, preserving the layouts-are-a-function-of-the-seed-alone property #516 established — and APPLIED only when backend_version() >= 92600; older backends fall back to the packed BHSD default (old behavior), and full layout coverage engages automatically once CI deploys 9.26. Also wires the --implementation CLI option through the random fwd/bwd/bias L0 tests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix: reject non-BHSD softmax-stats strides for cuDNN < 9.26 (NVBug 6057616) The SM80 (and SM100 dBias) backward kernels in cuDNN < 9.26 ignore the declared strides of the Stats tensor and address it as packed BHSD, silently producing wrong gradients for any other layout. Guard both the forward and backward paths in the C++ cuDNN backend: - sdpa_support_surface.h: reject a non-ragged Stats *output* with non-BHSD strides when building a forward graph on cuDNN < 9.26, so the error surfaces at graph construction time rather than at bprop. - CompositeSDPABackwardNode::pre_validate_node(): reject a non-ragged Stats *input* with non-BHSD strides on cuDNN < 9.26 (the authoritative check; the forward check above is early-warning only). Also fixes a pre-existing formatting issue in the same region of scaled_dot_product_flash_attention.h: missing space in if(, mismatched continuation indent, and a tab in the closing brace. test: randomize softmax-stats strides unconditionally in random_config.py (the FE now rejects the broken configs on < 9.26 rather than silently corrupting; tests skip via GRAPH_NOT_SUPPORTED instead of producing wrong gradients). Verified 378/378 bwd L0 pass on cuDNN 9.26 (H100). --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

vedaanta added cat-bug Reports of incorrect behavior, crashes, regressions, or unexpected results. orig-nv-eng Reported or requested by NVIDIA engineering. mod-frost labels Aug 7, 2026

vedaanta mentioned this pull request Aug 7, 2026

sdpa: add torch custom ops cudnn::sdpa_fwd / cudnn::sdpa_bwd #517

Open

coderabbitai Bot reviewed Aug 7, 2026

View reviewed changes

Comment thread test/python/sdpa/random_config.py Outdated

coderabbitai Bot mentioned this pull request Aug 7, 2026

frost(sdpa): support ragged stats for SM100 frost sdpa forward engine #512

Merged

3 tasks

vedaanta force-pushed the vagarwalla/thd-kv-token-stride branch from 1552036 to 63ad9a6 Compare August 8, 2026 05:16

vedaanta changed the title ~~sdpa(frost): decline non-packed THD K/V token strides; fuzz them in test_mhas_v2~~ test(sdpa): fuzz per-tensor ragged token-stride gaps in test_mhas_v2 Aug 8, 2026

coderabbitai Bot reviewed Aug 8, 2026

View reviewed changes

vedaanta mentioned this pull request Aug 8, 2026

frost(sdpa): native THD declared-stride support in the SM100/SM120 f16 fwd kernels #526

Merged

vedaanta force-pushed the vagarwalla/thd-kv-token-stride branch from 60b9a43 to 1cc22af Compare August 10, 2026 06:33

vedaanta force-pushed the vagarwalla/thd-kv-token-stride branch from 1cc22af to 9b57eab Compare August 10, 2026 18:29

vedaanta and others added 4 commits August 10, 2026 23:32

vedaanta force-pushed the vagarwalla/thd-kv-token-stride branch from 9b57eab to 746ef17 Compare August 11, 2026 07:10

coderabbitai Bot reviewed Aug 11, 2026

View reviewed changes

vedaanta requested a review from Aneureka August 11, 2026 07:25

Aneureka approved these changes Aug 11, 2026

View reviewed changes

vedaanta merged commit e002833 into NVIDIA:develop Aug 11, 2026
1 check passed

This was referenced Aug 12, 2026

SM80 (A100) SDPA: FROST engines + cudnn.sdpa adapters #493

Merged

fix: reject non-BHSD softmax-stats strides on cuDNN < 9.26; randomize in tests (NVBug 6057616) #304

Merged

Navigation Menu

Uh oh!

test(sdpa): fuzz per-tensor ragged token-stride gaps in test_mhas_v2 - #516

test(sdpa): fuzz per-tensor ragged token-stride gaps in test_mhas_v2#516
vedaanta merged 4 commits into
NVIDIA:developfrom
vedaanta:vagarwalla/thd-kv-token-stride

vedaanta commented Aug 7, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Aug 7, 2026 •

edited

Loading

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Uh oh!

vedaanta commented Aug 11, 2026 •

edited

Loading

Uh oh!

vedaanta commented Aug 11, 2026

Uh oh!

cudnn-ci-bot commented Aug 11, 2026

Uh oh!

coderabbitai Bot left a comment

Uh oh!

coderabbitai Bot Aug 11, 2026

Uh oh!

Aneureka left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

vedaanta commented Aug 7, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Known failures this exposes (intentionally NOT masked)

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

vedaanta commented Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vedaanta commented Aug 11, 2026

Uh oh!

cudnn-ci-bot commented Aug 11, 2026

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot Aug 11, 2026

Choose a reason for hiding this comment

Uh oh!

Aneureka left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

vedaanta commented Aug 7, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Aug 7, 2026 •

edited

Loading

vedaanta commented Aug 11, 2026 •

edited

Loading