frost(sdpa): native THD declared-stride support in the SM100/SM120 f16 fwd kernels - #526
Conversation
vedaanta
added
cat-bug
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughTHD forward execution now preserves supported Q/K/V/O strides for SM100 and SM120 kernels. Support checks reject unsupported TMA layouts. Compilation specializes on declared strides, and output descriptors use declared output strides. ChangesTHD declared-stride support
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant StrideValidation
participant SDPACompiler
participant SDPAKernel
Caller->>StrideValidation: Provide THD tensor stride declarations
StrideValidation->>StrideValidation: Validate TMA-compatible strides
StrideValidation->>SDPACompiler: Construct native Q, K, V, and O views
SDPACompiler->>SDPAKernel: Launch stride-specialized kernel
SDPAKernel->>Caller: Write results through declared strides
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@python/cudnn/sdpa/fwd/api_dsl.py`:
- Line 1085: Rename the ambiguous O variable to o_kernel in both THD execution
paths, including the assignment from _thd_normalized and all subsequent kernel
and scatter calls that reference it. Ensure both occurrences are updated
consistently without changing execution behavior.
- Around line 348-367: Update the THD preparation flow around _thd_normalized to
resolve current_stream before the SM120 zero-KV branch and normalization, then
execute all THD metadata copies, cumsum, gathers, zeroing, LSE/sink writes, and
output scatter within _torch_stream_context. Ensure both packed and declared THD
paths, including normalization and scatter-back operations, remain bound to the
supplied current_stream before launching the kernel.
🪄 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: e9fd5230-aec0-4fbc-a2da-ac16f293d4aa
📒 Files selected for processing (1)
python/cudnn/sdpa/fwd/api_dsl.py
68eeabb to
f53d3e2
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (4)
python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py (1)
1722-1733: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueValidate the O stride with
CFG.BPE_O.
_fake_bshdalways usesCFG.BPEfor the 16-byte check.fake_ois built withOUT_STORAGE_DTYPE, and this flavor tracks a separateCFG.BPE_O(the envelope check at Line 1718 uses it). WhenBPE_Odiffers fromBPE, the O stride check computes the wrong byte stride and can accept a descriptor that violates the TMA rule.Today the THD path only passes
o_stridefor f16/bf16, whereBPE_O == BPE, so no current configuration is broken. Parameterize the byte size so the guard stays correct if an FP8 O reaches this path.♻️ Proposed fix
- def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE): + def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): if stride is None: return cute.runtime.make_fake_compact_tensor(dtype, shape, stride_order=(3, 2, 1, 0), assumed_align=16) for axis in (1, 2): # seq/head global strides feed TMA: 16-byte rule - if (stride[axis] * CFG.BPE) % 16 != 0: - raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={CFG.BPE} (TMA global-stride rule)") + if (stride[axis] * bpe) % 16 != 0: + raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={bpe} (TMA global-stride rule)") return cute.runtime.make_fake_tensor(dtype, shape, tuple(stride), assumed_align=16) fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) fake_k = _fake_bshd((_fake_batch, skv, kh, d_qk), k_stride) fake_v = _fake_bshd((_fake_batch, skv, kh, d_v), v_stride) - fake_o = _fake_bshd((_fake_batch, sq, qh, d_v), o_stride, dtype=OUT_STORAGE_DTYPE) + fake_o = _fake_bshd((_fake_batch, sq, qh, d_v), o_stride, dtype=OUT_STORAGE_DTYPE, bpe=CFG.BPE_O)🤖 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 `@python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py` around lines 1722 - 1733, Update _fake_bshd to accept a byte-size parameter, defaulting to CFG.BPE for Q, K, and V stride validation, and pass CFG.BPE_O when constructing fake_o so the output stride check uses OUT_STORAGE_DTYPE’s actual byte size.python/cudnn/sdpa/fwd/api_dsl.py (1)
2016-2020: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAlign the
_bindnative-stride gate withis_layout_supported.
_bindaccepts a declared stride whenes == 1 and ts % 8 == 0 and hs % 8 == 0. The SM120 kernel additionally requireshs >= head_dimandts >= heads * hsinis_layout_supported(prefill_f16_sm120.pyLines 173-176). A sub-dense declaration that satisfies the adapter gate reachescute.compileand fails there with the message "must use compact BSHD storage", which does not describe the real cause.Add the covering checks to
_bindso an unsupported layout takes the packed-normalization fallback instead of failing at compile.♻️ Proposed gate alignment
def _bind(buf, desc, tokens, gather): (ts, hs, es), packed = self._thd_declared(desc) - if packed or (es == 1 and ts % 8 == 0 and hs % 8 == 0): + h, d = desc.shape[1], desc.shape[3] + if packed or (es == 1 and ts % 8 == 0 and hs % 8 == 0 and hs >= d and ts >= h * hs): return self._thd_view(buf, desc, tokens), None return self._thd_normalized(buf, desc, tokens, carver, gather=gather)🤖 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 `@python/cudnn/sdpa/fwd/api_dsl.py` around lines 2016 - 2020, Update _bind’s native-stride condition to also require hs >= head_dim and ts >= heads * hs, matching is_layout_supported’s SM120 requirements. Ensure layouts failing either covering check use _thd_normalized instead of _thd_view, while preserving the existing packed and valid-stride behavior.python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py (1)
1253-1255: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winImprove the layout rejection message.
is_layout_supportednow accepts padded layouts, but the error text still says "must use compact BSHD storage". A caller that supplies a sub-dense or misaligned declared stride reads a message that names the wrong requirement.♻️ Proposed message update
for name, tensor in (("Q", q), ("K", k), ("V", v), ("O", o)): if cutlass.const_expr(not self.is_layout_supported(tensor.shape, tensor.stride)): - raise ValueError(f"{name} must use compact BSHD storage") + raise ValueError( + f"{name} must be BSHD with the head dim innermost-contiguous and " + f"non-overlapping head/seq strides that are 16-byte multiples; " + f"got stride {tensor.stride} shape {tensor.shape}" + )🤖 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 `@python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py` around lines 1253 - 1255, Update the ValueError raised in the is_layout_supported validation loop to describe the actual supported layout requirement rather than requiring “compact BSHD storage.” Ensure the message accurately covers padded layouts while still indicating that unsupported sub-dense or misaligned strides are rejected.python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py (1)
2085-2091: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy lift
_fake_bshdis duplicated five times and the copies validate different things. Each kernel module defines its own_fake_bshd. The four SM100 copies check the 16-byte rule on axes 1 and 2 but never check that the innermost stride is 1; the SM120 copy performs no validation at all. Every consumer assumes the head dim is innermost-contiguous, so a non-unitstride[3]produces an invalid TMA descriptor with no error. Extract one shared helper (for example next tobuild_o_descs_kernelinthd_sm100.py) that takes the byte-per-element as a parameter, and add thestride[3] == 1guard there.
python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py#L2085-L2091: replace the local_fake_bshdwith the shared helper and passCFG.BPE.python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py#L2172-L2178: replace the local_fake_bshdwith the shared helper and passCFG.BPE.python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py#L1722-L1728: replace the local_fake_bshdwith the shared helper; passCFG.BPEfor Q/K/V andCFG.BPE_Ofor O.python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py#L1924-L1930: replace the local_fake_bshdwith the shared helper; passCFG.BPEfor Q/K/V andCFG.BPE_Ofor O.python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py#L1432-L1435: replace the unvalidated local_fake_bshdwith the shared helper so SM120 rejects the same malformed strides as SM100.🤖 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 `@python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py` around lines 2085 - 2091, Duplicate _fake_bshd implementations use inconsistent stride validation; extract one shared helper near build_o_descs_kernel in thd_sm100.py that accepts BPE, validates axes 1 and 2 for 16-byte alignment, and requires stride[3] == 1. Replace the local helpers in python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py#L2085-L2091 and prefill_d192_d128_f16_sm100.py#L2172-L2178, passing CFG.BPE; in prefill_d256_f16_sm100.py#L1722-L1728 and prefill_d512_f16_sm100.py#L1924-L1930, pass CFG.BPE for Q/K/V and CFG.BPE_O for O; replace the unvalidated helper in prefill_f16_sm120.py#L1432-L1435 with the shared helper.
🤖 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.
Nitpick comments:
In `@python/cudnn/sdpa/fwd/api_dsl.py`:
- Around line 2016-2020: Update _bind’s native-stride condition to also require
hs >= head_dim and ts >= heads * hs, matching is_layout_supported’s SM120
requirements. Ensure layouts failing either covering check use _thd_normalized
instead of _thd_view, while preserving the existing packed and valid-stride
behavior.
In `@python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py`:
- Around line 2085-2091: Duplicate _fake_bshd implementations use inconsistent
stride validation; extract one shared helper near build_o_descs_kernel in
thd_sm100.py that accepts BPE, validates axes 1 and 2 for 16-byte alignment, and
requires stride[3] == 1. Replace the local helpers in
python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py#L2085-L2091 and
prefill_d192_d128_f16_sm100.py#L2172-L2178, passing CFG.BPE; in
prefill_d256_f16_sm100.py#L1722-L1728 and prefill_d512_f16_sm100.py#L1924-L1930,
pass CFG.BPE for Q/K/V and CFG.BPE_O for O; replace the unvalidated helper in
prefill_f16_sm120.py#L1432-L1435 with the shared helper.
In `@python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py`:
- Around line 1722-1733: Update _fake_bshd to accept a byte-size parameter,
defaulting to CFG.BPE for Q, K, and V stride validation, and pass CFG.BPE_O when
constructing fake_o so the output stride check uses OUT_STORAGE_DTYPE’s actual
byte size.
In `@python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py`:
- Around line 1253-1255: Update the ValueError raised in the is_layout_supported
validation loop to describe the actual supported layout requirement rather than
requiring “compact BSHD storage.” Ensure the message accurately covers padded
layouts while still indicating that unsupported sub-dense or misaligned strides
are rejected.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5e4e250b-6bd8-4241-a1f6-d1279f589622
📒 Files selected for processing (7)
python/cudnn/sdpa/fwd/api_dsl.pypython/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.pypython/cudnn/sdpa/fwd/kernels/thd_sm100.py
There was a problem hiding this comment.
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 `@python/cudnn/sdpa/fwd/api_dsl.py`:
- Around line 357-361: Update _thd_view to validate that buf.data_ptr() is
16-byte aligned before calling as_strided, rejecting misaligned runtime buffers.
Add a regression test covering a one-element FP16/BF16 D-offset view and verify
it is rejected before TMA view creation.
🪄 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: 7c2146ee-fbf9-44e3-a83c-f91b6cad9b31
📒 Files selected for processing (1)
python/cudnn/sdpa/fwd/api_dsl.py
…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>
…line, 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>
2414f33 to
046f1ea
Compare
vedaanta
changed the title
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/cudnn/sdpa/fwd/api_dsl.py (1)
41-68: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winRequire the exact unit S-scale pair on both FP8 paths.
A reciprocal pair is not equivalent to
(1.0, 1.0)because it changes E4M3 quantization rounding and underflow. SM100 currently accepts these pairs. SM120 also foldsdescale_sand passesscale_sto the kernel. Reject every pair exceptdescale_s == 1.0andscale_s == 1.0before either launch.Proposed fix
-def _require_reciprocal_s_scales(descale_s: float, scale_s: float) -> None: +def _require_unit_s_scales(descale_s: float, scale_s: float) -> None: - product = descale_s * scale_s - if abs(product - 1.0) > 1e-3: + if descale_s != 1.0 or scale_s != 1.0: raise NotImplementedError( - f"per-tensor FP8: this kernel converts P unscaled, so it can only serve a reciprocal " - f"descale_s*scale_s == 1; got {descale_s} * {scale_s} = {product}" + "per-tensor FP8 only supports descale_s == 1.0 and scale_s == 1.0" )- _require_reciprocal_s_scales(_scalar(descale_s), _scalar(scale_s)) + _require_unit_s_scales(_scalar(descale_s), _scalar(scale_s)) ... ds, ss = _scalar(descale_s), _scalar(scale_s) + _require_unit_s_scales(ds, ss)Based on learnings, SM100 and SM120 FP8 kernels require the exact unit pair because reciprocal values alter quantization behavior.
Also applies to: 1380-1385, 2105-2111
🤖 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 `@python/cudnn/sdpa/fwd/api_dsl.py` around lines 41 - 68, Update _require_reciprocal_s_scales to accept only the exact pair descale_s == 1.0 and scale_s == 1.0, rejecting every reciprocal or otherwise non-unit combination before either FP8 kernel launch. Remove the product-based tolerance check and report both received values in the existing NotImplementedError path; apply the same validation wherever the corresponding SM100 and SM120 launch flows invoke this guard.Source: Learnings
🤖 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 `@python/cudnn/sdpa/fwd/api_dsl.py`:
- Around line 407-411: Update _thd_view and its execute() call path to validate
the runtime buf before as_strided: require matching dtype and device, verify the
expected THD shape and declared strides/contiguity, and confirm the requested
view remains within accessible storage. Reject mismatches without implicit
conversions or allocations, then create the view only after validation.
---
Outside diff comments:
In `@python/cudnn/sdpa/fwd/api_dsl.py`:
- Around line 41-68: Update _require_reciprocal_s_scales to accept only the
exact pair descale_s == 1.0 and scale_s == 1.0, rejecting every reciprocal or
otherwise non-unit combination before either FP8 kernel launch. Remove the
product-based tolerance check and report both received values in the existing
NotImplementedError path; apply the same validation wherever the corresponding
SM100 and SM120 launch flows invoke this guard.
🪄 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: 60a2198f-02ba-435c-a1d0-40dee3cd752a
📒 Files selected for processing (3)
python/cudnn/AGENTS.mdpython/cudnn/sdpa/fwd/api_dsl.pypython/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py
🚧 Files skipped from review as they are similar to previous changes (2)
- python/cudnn/AGENTS.md
- python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py
…untime THD buffers, decline overlapping strides - Bind the THD host prep (metadata upload, O-descriptor zeroing, dummy-sink zeroing on SM100; the shared _thd_pack prep on SM120) to the LAUNCH stream via _torch_stream_context: the execute-time handle may carry a stream that is not torch's current, and the prep must be ordered against the kernel that consumes it. - _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). Addresses CodeRabbit review feedback on NVIDIA#526. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Addressed the remaining CodeRabbit findings (review-body comments without inline threads) in f13de3c — the THD host-prep stream binding from the same review round is pre-existing (#476) and split into a separate PR: Adapter gate vs kernel validator (the SM120 rejection message. Now states the actual contract (head dim innermost-contiguous, non-overlapping seq/head strides in 8-element multiples, compact or padded) and includes the offending shape/stride.
FP8 Verification for this commit: gate/view validation unit-checked (packed / whole-token gap / 8-elem gap pass; 4-elem gap, overlapping head, sub-dense token, dtype mismatch, misaligned base all decline or reject), 84/84 ragged fp16 slices + 11/11 ex-ops + gapped-stride repro green on SM100. |
… 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 (NVIDIA#476) and is split into a separate PR. Addresses CodeRabbit review feedback on NVIDIA#526. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
abd6efa to
f13de3c
Compare
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-526-f13de3c |
Aneureka
left a comment
Aneureka
left a comment
There was a problem hiding this comment.
Should we add tests to protect this fix?
| (ts, hs, es), _ = self._thd_declared(desc) | ||
| h, d = desc.shape[1], desc.shape[3] | ||
| self._not_implemented_error_if( | ||
| es != 1 or ts % 8 != 0 or hs % 8 != 0 or hs < d or ts < h * hs, |
There was a problem hiding this comment.
It seems we should replace 8 with 16 // desc.dtype.itemsize to make this function dtype-aware, since this is shared in the base class, although this PR targets FP16 kernels.
There was a problem hiding this comment.
Good catch — done in 1c563ca. The quantum is now 16 // desc.dtype.itemsize, computed per descriptor inside the loop (so a mixed-precision declaration checks each tensor at its own dtype): 8 elements at 2 B/elem (f16/bf16), 16 at 1 B/elem (fp8), 4 at 4 B/elem. No behavior change for the f16 paths this PR enables — and the fp8 native-stride follow-up (#537) now inherits the correct quantum for free.
Unit-checked across the three itemsizes (f16 +4-elem gap still declines; fp32 +4 = 16 B now passes; fp8 +8 = 8 B declines, +16 passes), and the SM100 suites re-ran green (84/84 ragged, 11/11 ex-ops).
…ize) 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 (NVIDIA#537) inherits the correct quantum for free. Suggested by @Aneureka in review. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-526-1c563ca |
The THD execute paths do torch host work before the kernel launch — the [seq_kv | cu_q | cu_k] metadata allocation + D2H length reads + one-shot H2D upload (SM100 and the shared SM120 _thd_pack), the per-sequence O-descriptor buffer (SM100), the dummy-sink buffer (SM100), and the cached seq_q dummy's first-use allocation (SM120). These enqueued on torch's CURRENT stream while the kernel launches on the stream carried by the execute-time handle (ExecutionContext.stream): when the two differ, the prep and the kernel race. Run the prep inside _torch_stream_context (the same helper the fp8/mxfp8 amax paths already use), and resolve the launch stream BEFORE _thd_pack in both SM120 callers. Allocations happen inside the context too, so caching-allocator blocks are stream-tagged to the stream that uses them. Pre-existing since the FROST engines landed (NVIDIA#476); split out of the NVIDIA#526 review round to keep that PR scoped to native THD stride support. Only direct graph-API users with an explicit handle stream are affected — the PyTorch integration launches on torch's current stream. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The THD execute paths do torch host work before the kernel launch — the [seq_kv | cu_q | cu_k] metadata allocation + D2H length reads + one-shot H2D upload (SM100 and the shared SM120 _thd_pack), the per-sequence O-descriptor buffer (SM100), the dummy-sink buffer (SM100), and the cached seq_q dummy's first-use allocation (SM120). These enqueued on torch's CURRENT stream while the kernel launches on the stream carried by the execute-time handle (ExecutionContext.stream): when the two differ, the prep and the kernel race. Run the prep inside _torch_stream_context (the same helper the fp8/mxfp8 amax paths already use), and resolve the launch stream BEFORE _thd_pack in both SM120 callers. Allocations happen inside the context too, so caching-allocator blocks are stream-tagged to the stream that uses them. Pre-existing since the FROST engines landed (NVIDIA#476); split out of the NVIDIA#526 review round to keep that PR scoped to native THD stride support. Only direct graph-API users with an explicit handle stream are affected — the PyTorch integration launches on torch's current stream. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The THD execute paths do torch host work before the kernel launch — the [seq_kv | cu_q | cu_k] metadata allocation + D2H length reads + one-shot H2D upload (SM100 and the shared SM120 _thd_pack), the per-sequence O-descriptor buffer (SM100), the dummy-sink buffer (SM100), and the cached seq_q dummy's first-use allocation (SM120). These enqueued on torch's CURRENT stream while the kernel launches on the stream carried by the execute-time handle (ExecutionContext.stream): when the two differ, the prep and the kernel race. Run the prep inside _torch_stream_context (the same helper the fp8/mxfp8 amax paths already use), and resolve the launch stream BEFORE _thd_pack in both SM120 callers. Allocations happen inside the context too, so caching-allocator blocks are stream-tagged to the stream that uses them. Pre-existing since the FROST engines landed (NVIDIA#476); split out of the NVIDIA#526 review round to keep that PR scoped to native THD stride support. Only direct graph-API users with an explicit handle stream are affected — the PyTorch integration launches on torch's current stream. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The THD execute paths do torch host work before the kernel launch — the [seq_kv | cu_q | cu_k] metadata allocation + D2H length reads + one-shot H2D upload (SM100 and the shared SM120 _thd_pack), the per-sequence O-descriptor buffer (SM100), the dummy-sink buffer (SM100), and the cached seq_q dummy's first-use allocation (SM120). These enqueued on torch's CURRENT stream while the kernel launches on the stream carried by the execute-time handle (ExecutionContext.stream): when the two differ, the prep and the kernel race. Run the prep inside _torch_stream_context (the same helper the fp8/mxfp8 amax paths already use), and resolve the launch stream BEFORE _thd_pack in both SM120 callers. Allocations happen inside the context too, so caching-allocator blocks are stream-tagged to the stream that uses them. Pre-existing since the FROST engines landed (NVIDIA#476); split out of the NVIDIA#526 review round to keep that PR scoped to native THD stride support. Only direct graph-API users with an explicit handle stream are affected — the PyTorch integration launches on torch's current stream. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…THD compile keys (#552) (#543) * frost(sdpa): bind the THD host prep to the launch stream The THD execute paths do torch host work before the kernel launch — the [seq_kv | cu_q | cu_k] metadata allocation + D2H length reads + one-shot H2D upload (SM100 and the shared SM120 _thd_pack), the per-sequence O-descriptor buffer (SM100), the dummy-sink buffer (SM100), and the cached seq_q dummy's first-use allocation (SM120). These enqueued on torch's CURRENT stream while the kernel launches on the stream carried by the execute-time handle (ExecutionContext.stream): when the two differ, the prep and the kernel race. Run the prep inside _torch_stream_context (the same helper the fp8/mxfp8 amax paths already use), and resolve the launch stream BEFORE _thd_pack in both SM120 callers. Allocations happen inside the context too, so caching-allocator blocks are stream-tagged to the stream that uses them. Pre-existing since the FROST engines landed (#476); split out of the #526 review round to keep that PR scoped to native THD stride support. Only direct graph-API users with an explicit handle stream are affected — the PyTorch integration launches on torch's current stream. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * frost(sdpa): drop the redundant zero-fill of the never-read O-descriptor scratch (SM100) A per-execute fill kernel on the THD execute hot path initialized the per-sequence O TMA-descriptor buffer, whose contents provably do not matter: the kernel's builder pass copies every qword of each sequence's slot from the base descriptor (then patches address/extent) before the fence and before any consumer read — stale workspace bytes never survive to a read. The +16-qword tail is never read at all. The fill dates to the original FROST landing (#476) as belt-and-braces. Rule 1: no adapter-side fills on the execute hot path. (The matching dummy-sinks fill removal is split into its own PR.) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * frost(sdpa): compile THD with dynamic token extents — plan-time-only compile keys The THD execute paths keyed the per-shape kernel compile on the packed token totals (sq=t_q, skv=t_kv, and max_sq on SM120). Under continuous batching the totals change every step, so the lru_cache degenerated into a fresh multi-second cute.compile per execute (issue #552's worst leg). - Kernel modules (SM100 d128/d192_d128/d256/d512 f16, SM120 f16/fp8): under THD the fake tensors' token extents are cute.sym_int symbols (one per ragged group — Q/O/LSE share t_q, K/V share t_kv) and the batch stride is rebuilt symbolically (the real view's batch stride is t * token_stride, a runtime value that never steps at batch extent 1). sq/skv are ignored under THD; SM100's _host reads the runtime totals from the dynamic tensor shapes. SM120's max_sq moves from a compile parameter to a runtime __call__ argument that sizes the per-sequence grid; trace-time shape checks compare only statically-known modes. - Adapter: the THD compile key is now derivable from the graph declaration alone, so compile() builds the artifact at PLAN time (the "thd-deferred" sentinel remains only for the unwired SM100 fp8 THD) and the execute paths' lru-cached compile calls are guaranteed hits; a shared _thd_compile_kwargs() keeps the two call sites identical. The all-KV-zero clamp's swapped K/V strides mint their own entry. - The D2H .tolist() round-trip still feeds the metadata upload, the ragged views' extents and the exact grid — removing it (and the CUDA- graph capture blocker) needs the plan-time-max grid + device cu_seqlens redesign tracked in #552. - New regression tests (SM100 + SM120) prove one compiled artifact serves different packed totals, checking numerics per total and asserting zero cache misses across executes. Verified on SM100 (B200-class): 487 passed / 4 skipped across the f16 dense+THD flavors, fp8, mxfp8, graph-level THD and sdpa op suites. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * docs(frost): AGENTS.md Hard Rules 4 (plan-time-only compile keys) and 5 (launch-stream-ordered execute) Renumbered after #570 landed Rule 3 (no D2H reads on execute); these two complement it. Rule 4 codifies issue #552's compile-key lesson: never key a kernel compile on runtime data values — runtime extents compile dynamic (cute.sym_int), runtime launch scalars are call arguments, derived values (batch strides computed from totals) count as leaks, and with a plan-time-only key the compile belongs at plan time with a cache-miss regression test guarding the execute path. Rule 3 bans the read that feeds such a key; Rule 4 bans the key itself. The SM80 _compile_cached (#493) is flagged as the known open cleanup. Rule 3's THD known-violation entry is updated: the compile-side half is done (dynamic token extents), so t_q/t_kv now reach the host only for the metadata upload, ragged view extents and the launch grid. Rule 5 codifies this PR's stream-binding fix: every torch operation on the execute path (H2D uploads, buffer resets, allocator calls, post-kernel consumers) is ordered on the launch stream via _torch_stream_context, never implicitly on torch's current stream. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
3 participants
Problem
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 stride2*h*d), the layouttorch.nn.attention.varlenusers get by slicing a fused KV projection. The FROST THD forward 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 architectures, while the backend engines serve the identical graphs correctly. Caught by #516's fuzz coverage and by PyTorch's own varlen test suite.Fix: native stride support, no fallback
compile()on all five f16 kernels (sm120, sm100 d128 / d192_d128 / d256 / d512) takes optional caller-declared per-tensor(batch, seq, head, elem)strides as part of the lru cache key.Nonekeeps the compact specialization bit-for-bit. Strided fakes are built withmake_fake_tensorand validated against the TMA 16-byte global-stride rule.kv_tma_descreads the tensor's strides instead of recomputing packed ones (Q/O offset math was already layout-driven). The entry validator now accepts padded, 16-byte-granular BSHD storage — compact is the equality special case.from_view, 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 packedn_qh*d_v.check_support(NotImplementedErrornaming the offending strides), so the Router falls back to an engine that honors the declaration. Not expressible = head dim not innermost-contiguous, or token/head strides 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.This lands together with AGENTS Hard Rule 2 (second commit): serve the declared layout natively or decline — never adapt; no hidden kernel launches on the execute path.
Verification (torch nightly cu132, ToT develop, with #516's fuzz tests)
test_varlen_vs_sdpa, incl. the kv-packedAttentionBlockcases)Turns #516's red tests green and #517's
test_thd_kv_packed_viewsgreen under FROST.Follow-ups (tracked as issues)
cu_seq_len(cu_ragged) form assumes packed strides🤖 Generated with Claude Code