frost(sdpa): THD zero-host-read extents break non-packed views; unit decode is extent-sensitive (regression from #606, also in #608)

[vedaanta]

Summary

#606's zero-host-read THD execute derives the packed token extents from a host-side buffer heuristic (t_q = min(q_buf.numel() // q_ts, o_buf.numel() // o_ts), same for KV). Two independent problems follow, one breaking non-packed views, one showing the kernel is extent-sensitive in a way no host heuristic can satisfy. Found while restacking the torch-ops work (#517/#554) on #608; SM100 (merged) and #608's SM120 path (buf.numel() // ts as well) both carry it.

Bug 1 — numel() // token_stride under-claims non-packed views

A K/V slice of a kv-interleaved [T, 2, H, D] buffer (the torch.nn.attention.varlen fused-projection layout, #526's headline case) holds T tokens but only T*H*D of the record's elements, so the derived extent halves. Same for any declared-stride gap view. Result: TMA extents cut off half the tokens — silently wrong O.

Repro (fails on develop, FROST route; backend route passes):

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}"

→ 41.1% O mismatches served by frost:sdpa_fwd_prefill_sm100_d128. (The always-on stride fuzzing from #516 catches this — the ragged sweeps are partially red on develop depending on the seed's gap draws and buffer binding style.)

A storage-derived exact-addressable capacity ((storage_elems_from_offset - row_footprint) // ts + 1) fixes every view case (verified: repro 41% → 0, kv-interleave probes bit-exact) — but breaks base-buffer bindings, because of bug 2.

Bug 2 — the kernel's unit decode is extent-sensitive; "over-claim is safe" does not hold

Forcing t_q for a failing sweep config (seq_len_q=[17,174,665,911], real total 1767, harness buffer capacity 1792, CGA_TILE_M=256):

forced t_q	result
1767 (exact)	pass
1792 (= 7×256 boundary)	pass
1793	fail, ~50% O mismatches
1794	fail
1856 (64-aligned!)	fail

Empirical law: corruption whenever ceil(t_q / CGA_TILE_M) > ceil(real_total / CGA_TILE_M) — the extent feeds the unit/tile decode (or something equivalent), so an over-claimed extent that crosses a tile boundary shifts real tiles onto wrong (batch, row) coordinates. The in-tree suites stay green only because the harness allocates packed_token_capacity (64-rounded) buffers, and a 64-rounding can never cross a 256 boundary.

Why no host-side formula can fix this

The adapter must choose t_q ∈ [real_total, real_total rounded up to the SAME tile count] — but the real total is device-only by #552's own design. numel//ts breaks views; exact-addressable capacity breaks any binding with slack past a tile boundary. The fix has to be kernel-side: make the unit decode (and anything else consuming the token extent) read the device metadata (cu[n_batch]) instead of the descriptor extent, so the extent can be a plain safe upper bound. Alternatively the envelope/decode contract needs to tolerate extra dead tile rows.

Affected

• SM100 (merged, frost(sdpa): THD execute with zero host reads on SM100 — device-built metadata, plan-time envelope grid, CUDA-graph capturable (issue #552) #606): _execute_thd capacity sites in api_dsl.py.
• SM120 (frost(sdpa): SM120 THD zero-host-read execute — device-built metadata, declared-S_q envelope grid, CUDA-graph capturable (issue #552) #608, open): same heuristic (return buf.numel() // ts, and the token-major LSE cap).
• Non-packed THD views through any consumer (the torch-ops path in sdpa: add torch custom ops cudnn::sdpa_fwd / cudnn::sdpa_bwd #517/torch: put torch.sdpa on the cuDNN Python API (CUDNN provider + dense backward) #554 binds views and is currently xfail'd on the kv-interleaved cases pending this fix).

Happy to pair on this — I have the probes, the forced-t_q knob diff, and the storage-capacity patch staged.

[vedaanta]

Blast-radius update (develop @ aded990, v1.28.0): this now breaks 40 upstream PyTorch tests that previously passed.

test_varlen_attention.py through the python-API integration: 69 failed / 100 passed, up from the long-standing 29-benign baseline. All 40 new failures are test_varlen_vs_sdpa window cases (window_size 0-5 x dtypes x gqa) whose inputs come from a fused-projection AttentionBlock — i.e. kv-interleaved non-contiguous THD views, exactly this issue's layout class. Two confirmations: contiguifying q/k/v before the varlen ops returns the suite to exactly 140/29, and the same configs pass with FROST disabled (backend route).

Why they regressed now: the engines recently started claiming tiny-window THD configs they previously declined (the right-band widening / KV-split / heuristics work), so configs that used to fall through to the backend now route to FROST and hit the extent bug. Every capability-coverage improvement widens this hole until the unit decode reads the device metadata instead of the token extent.

Isolated-probe note: contiguous THD + window_left in {1,2,3,6,64,128,-1} matches the backend bit-for-bit on SM100 — the windows themselves are fine; it is strictly the non-packed-view extents.

[vedaanta]

Fix up: #706 — root cause turned out sharper than the issue's original theory. The "extent-sensitive unit decode" law from the first forced-extent matrix dissolved once isolated properly: the decode was already metadata-driven; the real second failure mode is that masked KV rows still multiply (P(=0) x V), so rows between the real total and the extent must be finite — TMA zero-fill only covers rows at/beyond the extent. numel-derived extents halve interleaved views (failure mode 1), and storage-derived extents reach into allocator slack whose NaN bit patterns poison sequences through 0 x NaN (failure mode 2 — this is what looked like tile-boundary sensitivity).

The fix derives capacity from the view's element span (1 + sum((size-1)*stride)), exact on both edges, adapter-only, one shared helper across SM100 f16 and SM120/FP8 sites. Full verification matrix in the PR — including the 40 upstream PyTorch varlen tests returning to baseline.

[vedaanta]

For the record, since the PR that carried Fixes #613 was closed rather than merged: the fix landed on develop in 3631ecb (the #740 squash, which folded in the stacked #706 commits). Verified there — _thd_capacity with the element-span formula wired at all three extent sites (SM100 t_q/t_kv, SM120/FP8 _cap), the test_dsl_sm100_thd_interleaved_kv_views regression test, and the random_config.py fuzzer wiring that makes the randomized ragged sweeps actually draw non-packed token strides — the coverage seam that let this ship unnoticed.