frost(sdpa): derive THD token capacity from the view's element span (fixes #613)

[vedaanta]

Problem (issue #613)

The zero-host-read THD execute (#606/#608) derives the packed token extents host-side as numel() // token_stride. That is wrong on both edges for the buffers real integrations bind:

1. Non-packed views halve. 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 (41% mismatches on the frost(sdpa): THD zero-host-read extents break non-packed views; unit decode is extent-sensitive (regression from #606, also in #608) #613 repro). Through the PyTorch python-API integration this now breaks 40 upstream test_varlen_attention cases, because the engines recently started claiming tiny-window THD configs they previously declined.
2. Storage-derived extents poison through allocator slack. The obvious repair (derive from the untyped storage) over-claims into allocator slack — and that is not benign, which is the subtle part: rows between the real packed total and the extent are masked but still multiplied (P(=0) × V), so they must be finite; TMA zero-fill only covers rows at or beyond the extent. One slack row carrying NaN bit patterns poisons whole sequences through 0 × NaN (reproduced: batch-wide NaNs in the ragged sweeps).

Fix

Capacity = the largest T whose final token's row still fits in the buffer's own element span (1 + Σ (size−1)·stride). The span is exact on both edges:

• flat capacity buffers → exactly their token capacity (span = numel, no slack);
• interleaved / gapped views → exactly T (the last token needs only its own row footprint, not a full record span).

Every row below the capacity lies in caller-provided finite elements; every row at or beyond it TMA-clips to zeros. One shared _thd_capacity helper serves the SM100 f16 sites and the SM120/FP8 _cap sites (packed contract included).

Verification (SM100, isolated worktree + venv, no shared JIT cache)

Check	develop	this PR
#613 kv-interleave repro (test_repro, seeded)	41.1% O mismatches, frost-served	pass
new deterministic regression test (fused-record K/V views vs packed binding, torch.equal)	fail	pass
test_sdpa_random_fwd_ragged_L0 5-seed slice (84 tests)	84/84	84/84 — no regressions
fp8 THD ragged slice	green	green
upstream PyTorch test_varlen_attention (with the torch-ops stack on top)	100 pass / 69 fail	140 pass / 29 fail — the long-standing impl-identity baseline

The regression test lives next to the #606 suite (test_dsl_sm100_thd_interleaved_kv_views) and encodes both failure modes in its docstring.

Fixes #613.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes

  • Improved THD attention handling for non-contiguous tensor views and buffers with unused capacity.
  • Corrected token extent calculations for Q/O and K/V data, improving result reliability.

• Tests

  • Added coverage for interleaved, strided K/V views and ragged offsets.
  • Verified that strided inputs produce the same results as equivalent contiguous inputs.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 22369f54-8441-4428-a0ca-7ca8a4fe7f9d

📥 Commits

Reviewing files that changed from the base of the PR and between 60d0b81 and bd7a8f2.

📒 Files selected for processing (1)

• test/python/sdpa/random_config.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

📝 Walkthrough

Walkthrough

The change adds addressable THD capacity calculation for SM100 and SM120. THD packing uses it for Q/O and K/V extents. Ragged strides use derived generation. A regression test validates interleaved strided K/V views.

Changes

THD capacity correction

Layer / File(s)	Summary
Addressable THD capacity integration python/cudnn/sdpa/fwd/api_dsl.py	_thd_capacity derives token capacity from tensor geometry. SM100 and SM120 THD paths use it for Q/O and K/V extents.
Derived ragged stride generation test/python/sdpa/random_config.py	Ragged Q/O and K/V strides remain unset during randomization. fill_derived_fields() generates them before returning the configuration.
Interleaved K/V regression coverage test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py	The test executes THD SDPA with interleaved K/V views, doubled token strides, and ragged offsets. It compares the output with contiguous bindings.

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

Merge Risk: ⚪ Minimal · up to bd7a8

The change corrects THD token-capacity handling for packed and interleaved views, with the supplied regression and compatibility checks passing; no actionable merge-blocking risk remains beyond normal checks and review.

Suggested reviewers: anerudhan

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Linked Issues check	⚠️ Warning	The helper addresses view-span under-counting, but the PR does not implement the issue's required kernel-side handling of extent-sensitive unit decoding.	Update unit and tile decoding to use device ragged metadata, or make extra extent rows safe, and add slack and tile-boundary coverage.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the THD token-capacity fix and references issue #613.
Description check	✅ Passed	The description clearly covers the problem, fix, issue, and verification, but omits explicit template headings for affected area and testing.
Out of Scope Changes check	✅ Passed	The code and test changes support issue #613 by fixing THD capacity derivation and enabling regression coverage for non-packed ragged views.
Docstring Coverage	✅ Passed	Docstring check was indeterminate for this PR — some files could not be analyzed in time. Not blocking.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[vedaanta]

Added bd7a8f2 — test(sdpa): actually fuzz ragged token gaps in the randomized sweeps — the answer to "why didn't the fuzzer catch this?".

The coverage seam. The seeded token-gap draw (#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. So the randomized ragged fleet has never bound a non-packed THD stride, and for packed buffers numel() // token_stride is exact — no failure surface. Census over the fwd ragged L0 slice (84 configs): 0/84 drew a gap, although each config's own rng_geom_seed hand-draws nonzero gaps (e.g. [2,1,1,2]); the pre-fix adapter passes all 84, frost-served. That is exactly how this bug shipped and had to be found through an external integration.

The commit makes the ragged branch leave Q/K/V/O strides None and ends __call__ with fill_derived_fields() — one source of truth for the gap draw and its auto-packed fallbacks (cu / offset-multiplier forms #538, 1-byte dtypes #537). Head-major stats stride and the dense branch are untouched.

Red-team evidence (SM100, cuDNN 9.26.0.33, CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1):

suite (wired fleet)	pre-fix adapter	with this PR's fix
fwd ragged L0 slice (84, all frost-served)	80/84 FAIL	84/84 pass

i.e. this wiring alone would have caught #613 the day the capacity heuristic merged. Post-fix matrix, all green:

• fwd ragged L0 slice: 84/84 (census: 84/84 now gapped, all 84 vulnerable under the old formula)
• bwd ragged L0 slice: 158/158 — identical to the unwired control on the same lib (backend serves every gapped gradient combination; bwd FROST engines decline THD, so routing is unchanged)
• ragged_unified_L1: 24/24 · offset_multiplier_unified_L1: 24/24 (cu / mult forms stay packed via the existing fallbacks, 20/20 each in an offline census)
• test_ragged_token_gap_stable_under_stride_overrides: pass

Note: the bwd baseline requires a backend with the dSink+ragged-stats fixes (9.26 here); on a 9.24 lib the bwd ragged sweep fails for that unrelated pre-existing backend reason with or without this commit.

[Anerudhan]

Run CI.
Moving to 1.29

[vedaanta]

Follow-up filed as #718. While reasoning about how much of the capacity contract this PR actually buys, I confirmed the other half is still open: #706 guarantees every row below capacity is caller-owned, but not that it is initialized. The f16 THD paths have no equivalent of the FP8 packed-total clamp, so an over-allocated K/V buffer with an unwritten tail still poisons O — measured 49.6% NaN at seq_lens=[200,150,47], CAP=640, and exactly 0 when the extent equals the packed total. Repro and the port plan (the FP8 build_thd_meta_o_kv_descs_kernel already does all of it) are in #718.

Nothing here changes this PR — the span fix is still needed and still correct, since the clamp only helps tensors that get one. Flagging the linkage so the two are reviewed together.