SM80 (A100) SDPA: FROST engines + cudnn.sdpa adapters - #493
Conversation
egilliam-nv
added
cat-feature
|
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:
📝 WalkthroughWalkthroughAdded experimental SM80 SDPA forward and backward kernels, standalone APIs, FROST engine integration, graph analysis, validation tests, exports, and documentation. The implementation supports FP16/BF16 inputs, multiple masks, GQA/MQA, packed inputs, optional features, and deterministic backward execution. ChangesSM80 SDPA support
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant SdpafwdSm80
participant SM80Kernel
participant FROST
Caller->>SdpafwdSm80: submit forward tensors and features
SdpafwdSm80->>SM80Kernel: validate, compile, and launch cached kernel
SM80Kernel-->>SdpafwdSm80: return output, LSE, and optional statistics
FROST->>SdpafwdSm80: lower eligible graph execution
SdpafwdSm80-->>FROST: copy normalized outputs and statistics
Possibly related PRs
Suggested labels: 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: 20
🧹 Nitpick comments (17)
docs/fe-oss-apis/attention/sdpa_fwd_sm80.md (1)
76-80: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueUpdate the stale SM80 envelope comment in
python/cudnn/sdpa/fwd/api_sm80.py. The forward adapter supports qwen(256, 256)and routes it toprefill_d256_f16_sm80; the four documented 256-dimensional claims are correct. Update the comment at lines 315–317, which still says the envelope ends at dsv3(192, 128).🤖 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 `@docs/fe-oss-apis/attention/sdpa_fwd_sm80.md` around lines 76 - 80, Update the stale SM80 envelope comment near the forward adapter’s kernel routing so it includes qwen (256, 256) and reflects the full (256, 256) envelope; preserve the existing prefill_d256_f16_sm80 routing. The references in docs/fe-oss-apis/attention/sdpa_fwd_sm80.md lines 76-80, docs/fe-oss-apis/overview.md lines 30-31, docs/operations/Attention.md lines 722-725, and llms.txt line 31 require no direct changes because their 256-dimensional claims are already correct.python/cudnn/sdpa/graph_analyzer.py (1)
714-727: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused unpack at line 719.
b,h,s_, anddare never used. Ruff reports RUF059 for all four. The function readst.shapedirectly at lines 723-724.♻️ Proposed cleanup
- b, h, s_, d = t.shape strides = t.stride()🤖 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/graph_analyzer.py` around lines 714 - 727, Remove the unused b, h, s_, and d unpacking from to_bshd_physical; retain the direct t.shape and stride-based layout checks and existing permutation behavior unchanged.Source: Linters/SAST tools
test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py (2)
223-226: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReplace the assigned lambda with a
def.Ruff reports E731 for line 223.
♻️ Proposed cleanup
- phys = lambda t: t.permute(0, 2, 1, 3).contiguous().permute(0, 2, 1, 3) + def phys(t): + return t.permute(0, 2, 1, 3).contiguous().permute(0, 2, 1, 3) +🤖 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/frost/test_sdpa_sm80_frontend_integration.py` around lines 223 - 226, Replace the lambda assigned to phys with a named def helper that performs the same permute, contiguous, and inverse-permute operations; keep the existing calls to phys unchanged.Source: Linters/SAST tools
111-116: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the reverse direction that the docstring promises.
The docstring says the check runs both ways. The test only asserts that the backward table rejects forward facts. The new forward-table rejection at
python/cudnn/sdpa/fwd/engines.pylines 188-189 has no assertion here. Line 116 also indexes the result ofmismatchdirectly; ifmismatchreturnsNonethe test fails with aTypeErrorinstead of a clear assertion message.💚 Proposed test additions
def test_direction_cross_rejection(): """The fwd engine must reject sdpa_backward graphs and vice versa.""" g, *_ = _build_fwd_graph() facts = dataclasses.replace(ga.analyze(g), device_cc=(8, 0)) assert engines_fwd.mismatch(_FWD_CAPS, facts) is None - assert "direction" in engines_bwd_sm80.mismatch(engines_bwd_sm80.CAPABILITIES, facts) + bwd_reason = engines_bwd_sm80.mismatch(engines_bwd_sm80.CAPABILITIES, facts) + assert bwd_reason is not None and "direction" in bwd_reason + + bwd_facts = dataclasses.replace(facts, is_backward=True) + fwd_reason = engines_fwd.mismatch(_FWD_CAPS, bwd_facts) + assert fwd_reason is not None and "forward-only" in fwd_reason🤖 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/frost/test_sdpa_sm80_frontend_integration.py` around lines 111 - 116, Update test_direction_cross_rejection to also assert that engines_fwd.mismatch rejects backward-graph facts with a direction-related result, covering the forward-table rejection described by the docstring. Avoid indexing mismatch results directly; assert the returned value is non-None before checking that it identifies the direction mismatch, so failures remain clear.python/cudnn/sdpa/bwd/engine.py (1)
94-98: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMirror the descriptive offset lookup from
fwd/engine.py.
_ID_OFFSETS[ENGINE_NAME]raises a bareKeyErrorwhen a future backward opset lands without an offset entry.python/cudnn/sdpa/fwd/engine.pylines 147-148 raise a message that names the file and the never-reuse rule. Keep the two families consistent.♻️ Proposed change
def FrostSdpaBwdEngines() -> List[FrostSdpaBwdEngine]: """The SDPA-backward engine family, in preference order.""" from .engines import ENGINE_NAME - return [FrostSdpaBwdEngine(ENGINE_NAME, _ID_OFFSETS[ENGINE_NAME])] + if ENGINE_NAME not in _ID_OFFSETS: + raise KeyError(f"engine {ENGINE_NAME!r} has no engine-id offset; allocate the next free one in engine._ID_OFFSETS (never reuse)") + return [FrostSdpaBwdEngine(ENGINE_NAME, _ID_OFFSETS[ENGINE_NAME])]🤖 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/bwd/engine.py` around lines 94 - 98, Update FrostSdpaBwdEngines to validate ENGINE_NAME before indexing _ID_OFFSETS, matching the descriptive lookup and error behavior in FrostSdpaFwdEngines. Raise an explicit error that names the backward engine offset source and states the offset must never be reused when the entry is missing; preserve the existing engine construction for valid entries.test/python/fe_api/sdpa/test_sdpa_fwd_sm80.py (1)
37-67: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove or correct the unreachable non-causal window branch.
Lines 94-95 force
is_causal = Trueformask == "swa", so theelif window_size[0] >= 0branch at lines 59-63 never runs. That branch also implements a symmetric window ((i - j).abs() <= W), while the kernel keeps a left-only band[q-W, q]. If a future parametrization adds a non-causal window case, the reference silently disagrees with the kernel.The causal branch anchors the diagonal bottom-right (
j <= i + (s_kv - s_q)) while the wrapper is called with the defaultcausal_bottom_right=False. The current shapes uses_q == s_kv, so both alignments coincide; the reference will diverge if a case withs_q != s_kvis added.Lines 40-42 also unpack
b,d_qk, andd_vwithout using them (Ruff RUF059).🤖 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/fe_api/sdpa/test_sdpa_fwd_sm80.py` around lines 37 - 67, Update _ref_sdpa to match the kernel’s non-causal SWA semantics with a left-only window, and remove the unreachable symmetric-window branch unless non-causal window cases are supported by the wrapper. Align causal masking with causal_bottom_right=False rather than bottom-right anchoring, while preserving current equal-length behavior. Remove unused shape unpacking variables to satisfy Ruff RUF059.Source: Linters/SAST tools
python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.py (1)
117-121: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe default head dims contradict this kernel's supported shape.
This module serves the qwen flavor at
d_qk = d_v = 256, butDEFAULT_D_QKandDEFAULT_D_Vare both 128. The adapter always passes explicit values from_FLAVOR_DIMS, so the defaults are unreachable throughcudnn.sdpa. A direct caller offorward()with a genuine d=256 tensor and default kwargs hits the assert on Line 2141 (V D (256) must equal compile-time d_v=128).Set both defaults to 256 to match the module's scope.
🤖 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_sm80.py` around lines 117 - 121, Update DEFAULT_D_QK and DEFAULT_D_V in the prefill kernel configuration to 256, matching the qwen kernel’s supported head dimensions and allowing direct forward() calls with default arguments to accept d=256 tensors.python/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.py (3)
1845-1878: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winUnbounded compile cache keyed on problem shape. Both kernel modules decorate
_compile_cachedwith@lru_cache(maxsize=None)and includeB,H,H_kv,SQ,SKV, andDin the key. Every distinct problem shape therefore retains a compiled binary for the process lifetime, in both host memory and loaded device modules. A serving workload with varying batch size or sequence length grows this cache without bound. The docstring above each decorator justifies caching but does not address the missing size limit.
python/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.py#L1845-L1878: set a boundedmaxsizeon the@lru_cachedecorator for_compile_cached, sized to the expected number of live shape variants.python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.py#L1828-L1861: apply the same boundedmaxsizeto this module's_compile_cached.Consider also removing
is_even_mnandis_even_kfrom the key, since both are derived fromSQ,SKV, andD, which are already key entries.🤖 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_sm80.py` around lines 1845 - 1878, The _compile_cached caches are unbounded across both kernel modules; set a finite, expected-live-variant maxsize on the decorators in python/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.py:1845-1878 and python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.py:1828-1861. Also remove is_even_mn and is_even_k from both _compile_cached cache keys if they are deterministically derived from SQ, SKV, and D, preserving all other key parameters.
2150-2153: 🩺 Stability & Availability | 🔵 Trivial | 🏗️ Heavy liftInput validation relies entirely on bare
assert. Bothforward()entry points validate dtype, rank, head-dim divisibility, GQA head ratios, mask tokens, and optional-tensor shapes with bareassertstatements. Python removes everyassertunder-OorPYTHONOPTIMIZE. The kernel then computes GMEM base pointers and row strides directly from the unvalidated shapes, so a malformed input reaches device code that reads and writes out of bounds instead of raising.These are library entry points reachable from
cudnn.sdpa, so the optimization flag is set by the embedding application, not by this code.
python/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.py#L2150-L2153: convert the shape, dtype, and divisibility checks in thisforward()to explicitraise ValueError/raise NotImplementedErrorstatements.python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.py#L2126-L2129: apply the same conversion to this module'sforward().🤖 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_sm80.py` around lines 2150 - 2153, The forward() validation in python/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.py:2150-2153 and python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.py:2126-2129 must not rely on bare assert statements. Replace every dtype, shape/rank, divisibility, GQA ratio, mask-token, and optional-tensor validation in both forward() implementations with explicit ValueError or NotImplementedError raises, preserving the existing validation rules and messages sufficiently to reject malformed inputs even under optimized Python execution.
330-406: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftThe two SM80 prefill kernels duplicate most of their logic.
prefill_d256_f16_sm80.pystates in its docstring that it differs fromprefill_f16_sm80.pyonly in the K+V prefetch pipeline. The rest is near-identical: the scheduler decode block, the causal/SWA/padded mainloop bound computation, the interior/boundary split, the mask pre-pass, the online-softmax reductions, the sink finalize, the LSE and score-stat stores, the STG.128 epilogue,_sdpa_host,_compile_cached, andforward. Together that is roughly 2000 duplicated lines. Every future correctness or performance fix must be applied twice, and the stalecp.asynccomments already flagged in each file show the two copies drifting apart.
python/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.py#L330-L406: extract the scheduler decode block, the mainloop bound computation, the softmax body, and the epilogue into shared helpers undercudnn/frost/tile_dsl/, and call them from this module.python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.py#L303-L379: call the same shared helpers, and keep only the symmetric K+V prefetch pipeline local to this module.If the fork is deliberate for the current release, record that decision and the planned convergence point in both module docstrings.
🤖 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_sm80.py` around lines 330 - 406, Eliminate the duplicated SM80 prefill implementation by extracting the scheduler decode logic, mainloop bounds, softmax body, and epilogue into shared helpers under cudnn/frost/tile_dsl/. Update python/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.py#L330-L406 and python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.py#L303-L379 to call those helpers, leaving only the symmetric K+V prefetch pipeline local to the d256 module. If the fork must remain for this release, instead document the deliberate divergence and planned convergence point in both module docstrings.python/cudnn/sdpa/fwd/api_sm80.py (1)
525-528: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCache the padded V buffer instead of rebuilding it per call.
pad_vis true wheneverhead_dim_v < flavor_d_v._pad_last_dimthen allocates a zero tensor and runstorch.cat(...).contiguous(), which allocates a second buffer and copies the whole V tensor. Both allocations repeat on everyexecute()call.Allocate one padded buffer at
compile()time, sized(B, S_kv, H_kv, flavor_d_v), and copy only the[..., :head_dim_v]slice per call. That removes the per-call allocations and the padding zero-fill.🤖 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_sm80.py` around lines 525 - 528, Update the padding flow around pad_v and the execute/compile lifecycle to allocate and retain one reusable V buffer during compile(), shaped (B, S_kv, H_kv, flavor_d_v), when head_dim_v is smaller than flavor_d_v. In each execute() call, copy V into the buffer’s [..., :head_dim_v] slice and pass the cached buffer onward, replacing _pad_last_dim and avoiding per-call allocation and zero-fill.test/python/fe_api/sdpa/test_sdpa_bwd_sm80.py (2)
113-116: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse dtype-appropriate tolerances.
Lines 114-116 apply
rtol=atol=3e-2to both fp16 and bf16. bf16 carries roughly 8 mantissa bits against fp16's 11, so one tolerance is loose for fp16 and marginal for bf16 atd=256. Select the tolerance fromdtype, as the coding guidelines require.Based on the guideline "Compare test results against a reference implementation using existing reference-module patterns and dtype-appropriate tolerances."
🤖 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/fe_api/sdpa/test_sdpa_bwd_sm80.py` around lines 113 - 116, Update the backward-output assertions in the SDPA test to select rtol and atol based on the active dtype, using tighter tolerance for fp16 and appropriately looser tolerance for bf16. Apply the selected dtype-specific tolerances consistently to dq_tensor, dk_tensor, and dv_tensor comparisons.Source: Coding guidelines
63-69: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove the 48-case sweep off
L0.The parametrization is 4 head-dim pairs x 3 masks x 2 GQA shapes x 2 dtypes = 48 cases at
B=2, S=512, dup to 256. Each distinct configuration JIT-compiles its own CuTe-DSL kernel, so this is a slow suite. The coding guidelines requireL0tests to stay fast and large parameter sweeps to move to a higher level.Keep a small
L0subset, such as llama fp16 across the three masks, and mark the full sweepL2or higher.Based on the guideline "Mark every new Python test with a level from
L0throughL4; keepL0tests fast and place large parameter sweeps at higher levels."🤖 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/fe_api/sdpa/test_sdpa_bwd_sm80.py` around lines 63 - 69, Reduce the L0 coverage in test_sdpa_bwd_sm80_wrapper to a small representative subset, such as the llama fp16 configuration across all three mask values, and move the complete d_qk/d_v, mask, gqa, and dtype parametrization sweep to L2 or higher. Ensure every resulting test configuration retains an explicit test-level marker.Source: Coding guidelines
python/cudnn/sdpa/bwd/api_sm80.py (1)
118-132: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAlign the
_pick_flavordocstring with thed_qk >= d_vrejection.The docstring at Lines 121-123 and the comment at Lines 237-238 state that a
d_qk < d_vcase is padded up to an equal-d flavor. Line 239 rejects that case with aValueErrorinstead. Update the comments so the stated behavior matches the check.Also applies to: 237-239
🤖 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/bwd/api_sm80.py` around lines 118 - 132, The `_pick_flavor` docstring and the related comments near the `d_qk >= d_v` validation describe padding rejected inputs to an equal-d flavor, which contradicts the actual `ValueError` behavior. Update those comments to state that `d_qk < d_v` is rejected, while preserving the existing flavor selection and validation logic.python/cudnn/sdpa/bwd/kernels/bprop_d64_f16_sm80.py (1)
684-688: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
**_ignoredmakes the feature gate the only correctness barrier.
backwardaccepts and discards every unknown keyword.api_sm80.executetherefore relies entirely on_d64_fast_path_eligible(Lines 110-114 ofapi_sm80.py) to reject feature requests, and that check is a hand-maintained deny-list. A new feature kwarg added tobprop_f16_sm80.backwardand forwarded by the adapter, but not added to that list, produces silently wrong gradients rather than an error.Make the kernel reject unknown keywords instead of swallowing them. The adapter's
inspect.signaturefilter already drops kwargs the kernel does not declare, so an explicit rejection here only catches genuine gate misses.♻️ Proposed change
scale: Optional[float] = None, do_dot: Optional[torch.Tensor] = None, **_ignored, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """SDPA backward for head-dim 64, fp16/bf16. BSHD in/out; ``lse`` natural-log [B,H,S_q]. Returns (dQ, dK, dV).""" + _requested = {k: v for k, v in _ignored.items() if v not in (None, False, 0, "none")} + assert not _requested, f"bprop_d64_f16_sm80 implements plain dense MHA only; unsupported: {sorted(_requested)}" assert Q.dtype in (torch.float16, torch.bfloat16)🤖 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/bwd/kernels/bprop_d64_f16_sm80.py` around lines 684 - 688, Update bprop_f16_sm80.backward to reject unexpected keyword arguments instead of accepting them through **_ignored; remove the catch-all parameter while preserving the explicitly supported scale and do_dot arguments. Keep the adapter’s inspect.signature filtering unchanged so only genuine feature-gate misses reach this validation.python/cudnn/sdpa/bwd/kernels/bprop_config_gptoss_sm80.py (1)
32-50: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOnly
D_QKandD_Vare consumed today; the tile fields are not.
api_sm80._FLAVOR_DIMSreadsCFG.D_QKandCFG.D_V. Nothing passesTILE_KVorTILE_Qtobprop_f16_sm80.backward, which defaults both frombprop_config_llama_sm80(tile_q=64). So the docstring claim at Lines 34-35 that these fields configure the shared path at d=64 does not hold for any current caller. Either wire the flavor tiles through the adapter or narrow the docstring to state the fields are aspirational.🤖 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/bwd/kernels/bprop_config_gptoss_sm80.py` around lines 32 - 50, Update the configuration documentation around Cfg to accurately reflect current behavior: since api_sm80._FLAVOR_DIMS only consumes D_QK and D_V and no caller forwards TILE_KV or TILE_Q to bprop_f16_sm80.backward, remove the claim that these tile fields configure the shared path or explicitly label them aspirational. Do not imply the existing tile values are active configuration unless the adapter is updated to pass them through.python/cudnn/sdpa/bwd/engines.py (1)
88-89: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFix the two Ruff findings so lint stays clean.
Line 88 uses
!= True, which Ruff reports as E712. Line 241 leaves__all__unsorted, which Ruff reports as RUF022.♻️ Proposed fixes
- if facts.is_backward != True: + if not facts.is_backward: return "forward/backward direction mismatch"Apply this diff at Line 241:
-__all__ = ["ENGINE_NAME", "Capabilities", "CAPABILITIES", "mismatch", "analyze_for", "probe", "build"] +__all__ = ["CAPABILITIES", "Capabilities", "ENGINE_NAME", "analyze_for", "build", "mismatch", "probe"]🤖 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/bwd/engines.py` around lines 88 - 89, Fix both Ruff findings in the affected module: update the `facts.is_backward` check to use the idiomatic boolean form instead of comparing with `True`, and reorder the `__all__` entries into Ruff’s required sorted order without changing its exports.Source: Linters/SAST tools
🤖 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 `@docs/fe-oss-apis/attention/sdpa_bwd_sm80.md`:
- Around line 75-76: Validate the THD path before invoking _thd_backward: when
cum_seqlen_q_tensor is set, reject any non-null bias_tensor, rope_freqs, or
block_mask with an error. Preserve the existing dense path and ensure these
unsupported arguments cannot be silently dropped.
In `@docs/fe-oss-apis/attention/sdpa_fwd_sm80.md`:
- Around line 23-28: Complete the Requirements sections in
docs/fe-oss-apis/attention/sdpa_fwd_sm80.md lines 23-28 and
docs/fe-oss-apis/attention/sdpa_bwd_sm80.md lines 26-30 by documenting the exact
ctm DSL setup and runtime condition, or remove the dangling “and run with”/“and”
text if no further requirement applies.
- Around line 18-21: Update the provenance reference in the SDPA forward
documentation to point to python/cudnn/sdpa/fwd/kernels/__init__.py, matching
the forward kernel package described on the page.
In `@fe_cuda_13.3.cfg`:
- Line 1: Replace the absolute developer-local target of the fe_cuda_13.3.cfg
symbolic link with a repository-relative configuration target, or add and commit
the referenced configuration file so the link resolves in CI and other
checkouts.
In `@python/cudnn/sdpa/bwd/api_sm80.py`:
- Around line 531-549: Before the THD dispatch in the wrapper containing
_thd_backward, reject unsupported bias_tensor, block_mask, rope_freqs,
seq_kv_lens, seq_len_q, and max_s_q inputs using the same validation behavior as
the generic kernel’s THD guard; preserve supported alibi, sinks, and
deterministic forwarding, and ensure max_s_q is not silently accepted in either
dispatch path.
In `@python/cudnn/sdpa/bwd/kernels/bprop_config_llama_sm80.py`:
- Around line 6-7: Update the stale kernel name in the docstrings: in
python/cudnn/sdpa/bwd/kernels/bprop_config_llama_sm80.py lines 6-7, replace
bprop_sdpa_f16_sm80.py with bprop_f16_sm80.py; make the same name correction in
python/cudnn/sdpa/bwd/kernels/bprop_d64_f16_sm80.py line 41, with no other
changes.
In `@python/cudnn/sdpa/bwd/kernels/bprop_f16_sm80.py`:
- Around line 1750-1754: Enforce the documented RoPE constraint in backward() by
rejecting requests where has_rope is true and d_qk exceeds 128, before kernel
launch. Keep the existing seq_kv_lens and sequence-alignment validations intact,
and ensure invalid RoPE configurations cannot reach the sDQ allocation path in
the kernel.
- Around line 1735-1749: Use the unpacked D_K in the shape validation near the
existing d_qk/d_v assertions: require K’s last dimension to equal Q’s last
dimension before kernel indexing proceeds. Keep the existing dimension envelope
checks unchanged and remove the unused-variable condition by making D_K part of
this consistency check.
In `@python/cudnn/sdpa/fwd/api_sm80.py`:
- Around line 612-623: Update _thd_forward to resolve a None scale_softmax from
the original, pre-padding d_qk before Q/K are passed through _pad_last_dim,
using the same 1/sqrt(d_qk) default expected by the kernel. Preserve explicitly
supplied scale_softmax values and pass the resolved scale unchanged through the
later forward call.
- Around line 566-585: Update the kernel forward path used by execute so output
buffers are supplied or reused instead of allocating O and LSE with torch.zeros
on every call. Add optional BSHD-compatible O and LSE buffer parameters to the
relevant kernel forward symbol, pass cached scratch buffers or caller-provided
storage from the adapter, and preserve the existing transpose/copy behavior for
BHSD user tensors and score-stat outputs.
- Around line 314-317: Update the head-dimension envelope comment near the
validation logic to state that the envelope includes qwen and reaches D_QK=256
and D_V=256, matching _FLAVOR_DIMS and the module docstring; leave the
validation behavior unchanged.
- Around line 705-724: Update the THD branch in the forward execution flow,
before calling _thd_forward, to reject non-default or unsupported arguments:
rope_freqs, block_mask, scale_output != 1.0, scheduler, seq_kv_lens, and
seq_len_q. Raise a clear error for these combinations, while preserving the
existing max_s_q validation and _thd_forward call for supported inputs.
- Around line 388-403: Update the mask-token resolution near self.mask_token so
the SWA branch is entered only when swa_left >= 0, matching _thd_forward; handle
a nonnegative swa_right without a left window or causal mode as unsupported,
while preserving the existing causal right-window validation. Remove the
unreachable negative-value check for swa_window_runtime after its max(0,
swa_left) assignment.
In `@python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.py`:
- Around line 1038-1046: The synchronization comment above cp_async_wait in the
iteration-top block incorrectly describes a wait count of 2; update it to
document cp_async_wait(3) draining one oldest group from four pending groups,
including the iter-0 behavior, consistent with the module docstring and the
actual synchronization sequence.
In `@python/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.py`:
- Around line 16-30: Update the module docstring diagram to describe the shipped
configuration: 8 warps/256 threads per CTA and TILE_M=128 with each warp owning
16 rows, while retaining the existing SMEM figures. Replace the outdated K/V
layout note near the swizzle documentation to state that the layout uses the
implemented XOR swizzle, consistent with load_tile_2d(..., swizzle=True) and
swizzle_xor_128b.
- Around line 1020-1067: Update the comments in the mainloop block around the
prefetch and wait logic: describe the reverse predecessor predicate using the
actual `kv_iter > kv_left` condition and avoid the obsolete successor formula,
change the wait explanation to state that three pending groups and
`cp_async_wait(2)` drain only the oldest K group, and relabel the incorrect
“Step 4.5” heading to match the surrounding step sequence. Leave the wait counts
and implementation unchanged.
In `@python/cudnn/sdpa/fwd/kernels/sdpa_config_gptoss.py`:
- Around line 6-14: Clean up the module docstring in sdpa_config_gptoss.py by
removing the draft `(tile_m=128, tile_n=64, num_warps=8)` statement and “wait,
NO” reversal. State only the pinned `(tile_m=128, tile_n=64, num_warps=4)` point
and its M_BLOCKS/register-budget rationale, while preserving the driver-side
causal scheduling override note.
In `@python/cudnn/sdpa/graph_analyzer.py`:
- Around line 306-318: The backward canonicalization in the dimension-analysis
function must keep K/V strides aligned with any dimension swaps: whenever k_dim
or v_dim is transposed from (B, H, D, S) to (B, H, S, D), swap the corresponding
entries in strides["k"] or strides["v"] as well. Remove the redundant inner
k_dim unpack and retain the existing canonicalized dimension assignments.
In `@test/python/fe_api/sdpa/test_sdpa_bwd_sm80.py`:
- Around line 50-56: Update the causal reference in the test around the
keep-mask construction to match the wrapper’s default top-left alignment: remove
the s_kv - s_q offset from the causal diagonal and the window_left bound. Keep
the existing causal and sliding-window behavior unchanged for equal and unequal
sequence lengths.
In `@test/python/fe_api/sdpa/test_sdpa_fwd_sm80.py`:
- Around line 123-129: Guard the import of sdpa_fwd_wrapper_sm80 in
test_sdpa_fwd_sm80_padded_row_lse_trim with the same ImportError-to-skip
behavior used by test_sdpa_fwd_sm80_wrapper and
test_sdpa_fwd_sm80_check_support_rejections, or add an equivalent module-level
skip for missing cutlass/ctm so the test does not error without the optional
CuTe-DSL dependency.
---
Nitpick comments:
In `@docs/fe-oss-apis/attention/sdpa_fwd_sm80.md`:
- Around line 76-80: Update the stale SM80 envelope comment near the forward
adapter’s kernel routing so it includes qwen (256, 256) and reflects the full
(256, 256) envelope; preserve the existing prefill_d256_f16_sm80 routing. The
references in docs/fe-oss-apis/attention/sdpa_fwd_sm80.md lines 76-80,
docs/fe-oss-apis/overview.md lines 30-31, docs/operations/Attention.md lines
722-725, and llms.txt line 31 require no direct changes because their
256-dimensional claims are already correct.
In `@python/cudnn/sdpa/bwd/api_sm80.py`:
- Around line 118-132: The `_pick_flavor` docstring and the related comments
near the `d_qk >= d_v` validation describe padding rejected inputs to an equal-d
flavor, which contradicts the actual `ValueError` behavior. Update those
comments to state that `d_qk < d_v` is rejected, while preserving the existing
flavor selection and validation logic.
In `@python/cudnn/sdpa/bwd/engine.py`:
- Around line 94-98: Update FrostSdpaBwdEngines to validate ENGINE_NAME before
indexing _ID_OFFSETS, matching the descriptive lookup and error behavior in
FrostSdpaFwdEngines. Raise an explicit error that names the backward engine
offset source and states the offset must never be reused when the entry is
missing; preserve the existing engine construction for valid entries.
In `@python/cudnn/sdpa/bwd/engines.py`:
- Around line 88-89: Fix both Ruff findings in the affected module: update the
`facts.is_backward` check to use the idiomatic boolean form instead of comparing
with `True`, and reorder the `__all__` entries into Ruff’s required sorted order
without changing its exports.
In `@python/cudnn/sdpa/bwd/kernels/bprop_config_gptoss_sm80.py`:
- Around line 32-50: Update the configuration documentation around Cfg to
accurately reflect current behavior: since api_sm80._FLAVOR_DIMS only consumes
D_QK and D_V and no caller forwards TILE_KV or TILE_Q to
bprop_f16_sm80.backward, remove the claim that these tile fields configure the
shared path or explicitly label them aspirational. Do not imply the existing
tile values are active configuration unless the adapter is updated to pass them
through.
In `@python/cudnn/sdpa/bwd/kernels/bprop_d64_f16_sm80.py`:
- Around line 684-688: Update bprop_f16_sm80.backward to reject unexpected
keyword arguments instead of accepting them through **_ignored; remove the
catch-all parameter while preserving the explicitly supported scale and do_dot
arguments. Keep the adapter’s inspect.signature filtering unchanged so only
genuine feature-gate misses reach this validation.
In `@python/cudnn/sdpa/fwd/api_sm80.py`:
- Around line 525-528: Update the padding flow around pad_v and the
execute/compile lifecycle to allocate and retain one reusable V buffer during
compile(), shaped (B, S_kv, H_kv, flavor_d_v), when head_dim_v is smaller than
flavor_d_v. In each execute() call, copy V into the buffer’s [..., :head_dim_v]
slice and pass the cached buffer onward, replacing _pad_last_dim and avoiding
per-call allocation and zero-fill.
In `@python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.py`:
- Around line 117-121: Update DEFAULT_D_QK and DEFAULT_D_V in the prefill kernel
configuration to 256, matching the qwen kernel’s supported head dimensions and
allowing direct forward() calls with default arguments to accept d=256 tensors.
In `@python/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.py`:
- Around line 1845-1878: The _compile_cached caches are unbounded across both
kernel modules; set a finite, expected-live-variant maxsize on the decorators in
python/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.py:1845-1878 and
python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.py:1828-1861. Also remove
is_even_mn and is_even_k from both _compile_cached cache keys if they are
deterministically derived from SQ, SKV, and D, preserving all other key
parameters.
- Around line 2150-2153: The forward() validation in
python/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.py:2150-2153 and
python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.py:2126-2129 must not rely
on bare assert statements. Replace every dtype, shape/rank, divisibility, GQA
ratio, mask-token, and optional-tensor validation in both forward()
implementations with explicit ValueError or NotImplementedError raises,
preserving the existing validation rules and messages sufficiently to reject
malformed inputs even under optimized Python execution.
- Around line 330-406: Eliminate the duplicated SM80 prefill implementation by
extracting the scheduler decode logic, mainloop bounds, softmax body, and
epilogue into shared helpers under cudnn/frost/tile_dsl/. Update
python/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.py#L330-L406 and
python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.py#L303-L379 to call those
helpers, leaving only the symmetric K+V prefetch pipeline local to the d256
module. If the fork must remain for this release, instead document the
deliberate divergence and planned convergence point in both module docstrings.
In `@python/cudnn/sdpa/graph_analyzer.py`:
- Around line 714-727: Remove the unused b, h, s_, and d unpacking from
to_bshd_physical; retain the direct t.shape and stride-based layout checks and
existing permutation behavior unchanged.
In `@test/python/fe_api/sdpa/test_sdpa_bwd_sm80.py`:
- Around line 113-116: Update the backward-output assertions in the SDPA test to
select rtol and atol based on the active dtype, using tighter tolerance for fp16
and appropriately looser tolerance for bf16. Apply the selected dtype-specific
tolerances consistently to dq_tensor, dk_tensor, and dv_tensor comparisons.
- Around line 63-69: Reduce the L0 coverage in test_sdpa_bwd_sm80_wrapper to a
small representative subset, such as the llama fp16 configuration across all
three mask values, and move the complete d_qk/d_v, mask, gqa, and dtype
parametrization sweep to L2 or higher. Ensure every resulting test configuration
retains an explicit test-level marker.
In `@test/python/fe_api/sdpa/test_sdpa_fwd_sm80.py`:
- Around line 37-67: Update _ref_sdpa to match the kernel’s non-causal SWA
semantics with a left-only window, and remove the unreachable symmetric-window
branch unless non-causal window cases are supported by the wrapper. Align causal
masking with causal_bottom_right=False rather than bottom-right anchoring, while
preserving current equal-length behavior. Remove unused shape unpacking
variables to satisfy Ruff RUF059.
In `@test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py`:
- Around line 223-226: Replace the lambda assigned to phys with a named def
helper that performs the same permute, contiguous, and inverse-permute
operations; keep the existing calls to phys unchanged.
- Around line 111-116: Update test_direction_cross_rejection to also assert that
engines_fwd.mismatch rejects backward-graph facts with a direction-related
result, covering the forward-table rejection described by the docstring. Avoid
indexing mismatch results directly; assert the returned value is non-None before
checking that it identifies the direction mismatch, so failures remain clear.
🪄 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: 3dad10f3-1f50-4f7a-8f70-6fb6f7556cd7
📒 Files selected for processing (34)
docs/fe-oss-apis/attention/sdpa_bwd_sm80.mddocs/fe-oss-apis/attention/sdpa_fwd_sm80.mddocs/fe-oss-apis/overview.mddocs/operations/Attention.mdfe_cuda_13.3.cfgllms.txtpython/cudnn/engines/engine_ids.pypython/cudnn/engines/manifest.pypython/cudnn/sdpa/__init__.pypython/cudnn/sdpa/bwd/__init__.pypython/cudnn/sdpa/bwd/api_sm80.pypython/cudnn/sdpa/bwd/engine.pypython/cudnn/sdpa/bwd/engines.pypython/cudnn/sdpa/bwd/kernels/__init__.pypython/cudnn/sdpa/bwd/kernels/bprop_config_gptoss_sm80.pypython/cudnn/sdpa/bwd/kernels/bprop_config_llama_sm80.pypython/cudnn/sdpa/bwd/kernels/bprop_d64_f16_sm80.pypython/cudnn/sdpa/bwd/kernels/bprop_f16_sm80.pypython/cudnn/sdpa/fwd/__init__.pypython/cudnn/sdpa/fwd/api_sm80.pypython/cudnn/sdpa/fwd/engine.pypython/cudnn/sdpa/fwd/engines.pypython/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.pypython/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.pypython/cudnn/sdpa/fwd/kernels/sdpa_config_dsv3.pypython/cudnn/sdpa/fwd/kernels/sdpa_config_gptoss.pypython/cudnn/sdpa/fwd/kernels/sdpa_config_llama.pypython/cudnn/sdpa/fwd/kernels/sdpa_config_qwen.pypython/cudnn/sdpa/graph_analyzer.pytest/python/fe_api/sdpa/test_sdpa_bwd_sm80.pytest/python/fe_api/sdpa/test_sdpa_fwd_sm80.pytest/python/sdpa/frost/test_sdpa_graph_analyzer.pytest/python/sdpa/frost/test_sdpa_sm80_frontend_integration.pytest/python/sdpa/frost/test_sdpa_sm80_stream_respect.py
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
Rebased onto current Structural (from the rebase): Review fixes: every finding is addressed except two left open deliberately — the per-execute O/LSE allocation (disclosed in the description as the workspace-carving follow-up) and the analyzer stride-swap comment (already correct in this revision; replied inline). Highlights: the THD default-scale-after-padding bug is fixed (real correctness issue — thanks), both THD paths now reject dense-only features loudly, Re-verified on A100: fe_api suites 103/103, frost/analyzer/router suites 137/137, full |
There was a problem hiding this comment.
Actionable comments posted: 4
♻️ Duplicate comments (1)
python/cudnn/sdpa/bwd/kernels/bprop_f16_sm80.py (1)
1735-1739: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick windO and O head dims are still unvalidated.
Line 1738 now asserts
D_K == D, which closes the K side.d_vstill comes fromV.shape[3]alone. Every dO/O GMEM offset is computed fromd_v, so a caller that passes dO or O with a different last dim gets mis-indexed reads with no error.test/python/fe_api/sdpa/test_sdpa_bwd_sm80.pyLine 192 callsgen.backwarddirectly, so the adapter's shape validation does not cover this entry point.🐛 Proposed fix
assert d_qk >= d_v, f"bprop: d_qk ({d_qk}) must be >= d_v ({d_v})" + assert dO.shape[3] == d_v and O.shape[3] == d_v, f"dO/O head dim must match V's d_v ({d_v})"🤖 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/bwd/kernels/bprop_f16_sm80.py` around lines 1735 - 1739, Validate that the last dimensions of dO and O match V.shape[3] before assigning or using d_v in the backward kernel setup. Update the shape handling near d_qk and d_v so mismatches fail immediately, while preserving the existing K/Q head-dimension assertion and using the validated V dimension for GMEM offsets.
🧹 Nitpick comments (4)
test/python/fe_api/sdpa/test_sdpa_bwd_sm80.py (1)
116-119: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse dtype-appropriate tolerances.
Lines 117-119 apply
rtol=3e-2, atol=3e-2for bothtorch.float16andtorch.bfloat16. The bf16 mantissa is 8 bits shorter, so one shared tolerance sized for bf16 leaves the fp16 case under-constrained. Select the tolerance fromdtype.♻️ Proposed refactor
- # fp16 backward accumulates over S; scale tolerance accordingly. - torch.testing.assert_close(out["dq_tensor"].to(torch.float32), dq_ref, rtol=3e-2, atol=3e-2) - torch.testing.assert_close(out["dk_tensor"].to(torch.float32), dk_ref, rtol=3e-2, atol=3e-2) - torch.testing.assert_close(out["dv_tensor"].to(torch.float32), dv_ref, rtol=3e-2, atol=3e-2) + # The backward accumulates over S; scale the tolerance per input dtype. + tol = 1e-2 if dtype == torch.float16 else 3e-2 + for key, ref in (("dq_tensor", dq_ref), ("dk_tensor", dk_ref), ("dv_tensor", dv_ref)): + torch.testing.assert_close(out[key].to(torch.float32), ref, rtol=tol, atol=tol)As per coding guidelines: "Compare test results against a reference implementation using existing reference-module patterns and dtype-appropriate tolerances."
🤖 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/fe_api/sdpa/test_sdpa_bwd_sm80.py` around lines 116 - 119, Update the three assert_close calls in the SDPA backward test to derive rtol and atol from the active dtype, using tighter tolerances for torch.float16 and the existing broader tolerances for torch.bfloat16. Apply the selected dtype-specific values consistently to dq_tensor, dk_tensor, and dv_tensor comparisons.Source: Coding guidelines
python/cudnn/sdpa/graph_analyzer.py (2)
862-875: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
TensorDescis not resolvable in the annotation.Line 862 annotates the return type as
"TensorDesc", but the name is imported inside the function body on Line 863. The string annotation stays lazy at runtime, so no exception occurs.typing.get_type_hintsand static analysis cannot resolve it, and Ruff reportsF821. Add aTYPE_CHECKINGimport at module scope.♻️ Proposed module-scope addition
from typing import TYPE_CHECKING if TYPE_CHECKING: from cudnn.api_base import TensorDesc🤖 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/graph_analyzer.py` around lines 862 - 875, Add the module-scope TYPE_CHECKING import and conditionally import TensorDesc so the string return annotation in tensor_desc_from_ir can be resolved by static analysis and typing.get_type_hints. Keep the runtime-local import in tensor_desc_from_ir only if it is still required for constructing the return value.Source: Linters/SAST tools
802-810: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReplace the unused unpack with an explicit rank check.
Line 802 unpacks
b, h, s_, dand uses none of them. The unpack is the only thing that enforces the documented rank-4 contract, which makes the intent unclear and tripsRUF059. State the rank check directly.♻️ Proposed refactor
- b, h, s_, d = t.shape + if t.dim() != 4: + raise ValueError(f"to_bshd_physical: expected a rank-4 BHSD tensor; got rank {t.dim()}") strides = t.stride()🤖 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/graph_analyzer.py` around lines 802 - 810, In the shape-analysis logic around the stride/order checks, replace the unused `b, h, s_, d = t.shape` unpack with an explicit validation that `t` has rank 4. Preserve the existing BSHD detection and permutation behavior after enforcing this documented contract.Source: Linters/SAST tools
test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py (1)
224-227: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse a
definstead of a lambda assignment.Ruff reports
E731on Line 224. A named function also documents the BHSD-logical to BSHD-physical conversion the SM80 adapter requires.♻️ Proposed refactor
- phys = lambda t: t.permute(0, 2, 1, 3).contiguous().permute(0, 2, 1, 3) + def phys(t): + """BSHD-physical copy of a BHSD-logical tensor.""" + return t.permute(0, 2, 1, 3).contiguous().permute(0, 2, 1, 3) +🤖 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/frost/test_sdpa_sm80_frontend_integration.py` around lines 224 - 227, Replace the `phys` lambda assignment with a named local function that performs the same permute-contiguous-permute conversion, documenting the BHSD-logical to BSHD-physical transformation required by the SM80 adapter. Keep the calls to `phys` and all subsequent tensor handling unchanged.Source: Linters/SAST tools
🤖 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/bwd/api_sm80.py`:
- Around line 433-446: In _thd_backward, resolve a None scale_softmax using the
original unpadded d_qk before _pick_flavor and the Q/K padding logic, matching
the forward THD path’s ordering. Preserve explicitly supplied scale_softmax
values and continue forwarding the resolved scale to bprop_f16_sm80.backward.
- Around line 266-279: Update the mask-token resolution in the shown
initialization logic to enter the non-causal SWA branch only when swa_left >= 0,
matching the forward adapter and _thd_backward behavior. Explicitly reject
non-causal configurations with swa_right >= 0 but no left window, while
preserving the existing causal handling and "none" result for window_size=(-1,
0).
In `@test/python/fe_api/sdpa/test_sdpa_fwd_sm80.py`:
- Around line 70-75: Move the full parameterized matrix associated with the SDPA
forward test out of the L0 suite by changing its suite marker to L2 or higher.
Preserve the existing parameter combinations and add a separate small
representative smoke case under L0 only if needed.
- Around line 40-42: In the shape unpacking near the SDPA forward test, rename
the unused variables b, d_qk, and d_v to underscore-prefixed names while
preserving the existing h_q, s_q, h_kv, and s_kv assignments and behavior.
---
Duplicate comments:
In `@python/cudnn/sdpa/bwd/kernels/bprop_f16_sm80.py`:
- Around line 1735-1739: Validate that the last dimensions of dO and O match
V.shape[3] before assigning or using d_v in the backward kernel setup. Update
the shape handling near d_qk and d_v so mismatches fail immediately, while
preserving the existing K/Q head-dimension assertion and using the validated V
dimension for GMEM offsets.
---
Nitpick comments:
In `@python/cudnn/sdpa/graph_analyzer.py`:
- Around line 862-875: Add the module-scope TYPE_CHECKING import and
conditionally import TensorDesc so the string return annotation in
tensor_desc_from_ir can be resolved by static analysis and
typing.get_type_hints. Keep the runtime-local import in tensor_desc_from_ir only
if it is still required for constructing the return value.
- Around line 802-810: In the shape-analysis logic around the stride/order
checks, replace the unused `b, h, s_, d = t.shape` unpack with an explicit
validation that `t` has rank 4. Preserve the existing BSHD detection and
permutation behavior after enforcing this documented contract.
In `@test/python/fe_api/sdpa/test_sdpa_bwd_sm80.py`:
- Around line 116-119: Update the three assert_close calls in the SDPA backward
test to derive rtol and atol from the active dtype, using tighter tolerances for
torch.float16 and the existing broader tolerances for torch.bfloat16. Apply the
selected dtype-specific values consistently to dq_tensor, dk_tensor, and
dv_tensor comparisons.
In `@test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py`:
- Around line 224-227: Replace the `phys` lambda assignment with a named local
function that performs the same permute-contiguous-permute conversion,
documenting the BHSD-logical to BSHD-physical transformation required by the
SM80 adapter. Keep the calls to `phys` and all subsequent tensor handling
unchanged.
🪄 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: 7104d76e-62a5-4c59-8f3e-0ff39504872d
📒 Files selected for processing (33)
docs/fe-oss-apis/attention/sdpa_bwd_sm80.mddocs/fe-oss-apis/attention/sdpa_fwd_sm80.mddocs/fe-oss-apis/overview.mddocs/operations/Attention.mdfe_cuda_13.3.cfgllms.txtpython/cudnn/engines/manifest.pypython/cudnn/sdpa/__init__.pypython/cudnn/sdpa/bwd/__init__.pypython/cudnn/sdpa/bwd/api_sm80.pypython/cudnn/sdpa/bwd/engine.pypython/cudnn/sdpa/bwd/engines.pypython/cudnn/sdpa/bwd/kernels/bprop_config_gptoss_sm80.pypython/cudnn/sdpa/bwd/kernels/bprop_config_llama_sm80.pypython/cudnn/sdpa/bwd/kernels/bprop_d64_f16_sm80.pypython/cudnn/sdpa/bwd/kernels/bprop_f16_sm80.pypython/cudnn/sdpa/fwd/__init__.pypython/cudnn/sdpa/fwd/api_sm80.pypython/cudnn/sdpa/fwd/engine.pypython/cudnn/sdpa/fwd/engines.pypython/cudnn/sdpa/fwd/kernels/__init__.pypython/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.pypython/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.pypython/cudnn/sdpa/fwd/kernels/sdpa_config_dsv3.pypython/cudnn/sdpa/fwd/kernels/sdpa_config_gptoss.pypython/cudnn/sdpa/fwd/kernels/sdpa_config_llama.pypython/cudnn/sdpa/fwd/kernels/sdpa_config_qwen.pypython/cudnn/sdpa/graph_analyzer.pytest/python/fe_api/sdpa/test_sdpa_bwd_sm80.pytest/python/fe_api/sdpa/test_sdpa_fwd_sm80.pytest/python/sdpa/frost/test_sdpa_graph_analyzer.pytest/python/sdpa/frost/test_sdpa_sm80_frontend_integration.pytest/python/sdpa/frost/test_sdpa_sm80_stream_respect.py
🚧 Files skipped from review as they are similar to previous changes (19)
- docs/fe-oss-apis/overview.md
- docs/fe-oss-apis/attention/sdpa_bwd_sm80.md
- python/cudnn/sdpa/fwd/kernels/sdpa_config_dsv3.py
- docs/operations/Attention.md
- python/cudnn/sdpa/bwd/kernels/bprop_config_llama_sm80.py
- python/cudnn/sdpa/fwd/init.py
- python/cudnn/sdpa/fwd/kernels/sdpa_config_gptoss.py
- fe_cuda_13.3.cfg
- docs/fe-oss-apis/attention/sdpa_fwd_sm80.md
- python/cudnn/sdpa/bwd/init.py
- test/python/sdpa/frost/test_sdpa_graph_analyzer.py
- python/cudnn/sdpa/fwd/kernels/sdpa_config_llama.py
- python/cudnn/sdpa/fwd/kernels/sdpa_config_qwen.py
- python/cudnn/sdpa/fwd/engine.py
- python/cudnn/engines/manifest.py
- python/cudnn/sdpa/init.py
- python/cudnn/sdpa/bwd/engine.py
- llms.txt
- python/cudnn/sdpa/fwd/engines.py
|
Round 2 addressed in |
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 `@test/python/fe_api/sdpa/test_sdpa_bwd_sm80.py`:
- Around line 67-71: Add the explicit pytest.mark.L0 decorator to
test_sdpa_bwd_sm80_smoke, alongside its existing RNG fixture decorator, so this
representative smoke test is selectable at the intended L0 level.
🪄 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: 4b720b74-1910-47f3-acec-8cafe103433c
📒 Files selected for processing (3)
python/cudnn/sdpa/bwd/api_sm80.pytest/python/fe_api/sdpa/test_sdpa_bwd_sm80.pytest/python/fe_api/sdpa/test_sdpa_fwd_sm80.py
🚧 Files skipped from review as they are similar to previous changes (1)
- test/python/fe_api/sdpa/test_sdpa_fwd_sm80.py
253f37c to
9bd899c
Compare
|
@cudnn-ci-bot run frost,oss,python_tests |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-493-9bd899c |
|
@cudnn-ci-bot run frost,oss,python_tests |
|
Rebased onto current develop (past #494/#498/#502/#505/#506/#507 — the family/slot manifest, full-causal SM120 bwd, and SWA merges simplified this PR's diff: the arch-range and causal-shape gates it carried are gone) and re-triggered CI. Triage of the previous pipeline (61460182):
|
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-493-8e3b62f |
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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/graph_analyzer.py`:
- Around line 809-811: Update the facts.seq_q_t handling in the graph-analysis
flow to resolve seq_len_q with _need(), requiring the corresponding variant-pack
buffer whenever the graph binds it. Preserve the existing conditional behavior
when facts.seq_q_t is None and keep the resolved value assigned to
ops.seq_len_q.
- Line 848: Remove the unused shape unpacking assignment `b, h, s_, d = t.shape`
from the surrounding graph-analysis logic, leaving the rest of the code
unchanged.
In `@test/python/fe_api/sdpa/test_sdpa_bwd_sm80.py`:
- Line 52: Update the shape unpacking in the SDPA backward test to replace the
unused b and d_qk bindings with underscores, while retaining h_q and s_q for
subsequent use and resolving Ruff RUF059.
In `@test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py`:
- Around line 225-226: Replace the local phys lambda with a locally defined
helper function before the sdpa_fwd_wrapper_sm80 call, preserving its existing
permute-and-contiguous behavior and invocation for q, k, and v tensors.
🪄 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: f22f774d-d9ef-4bc0-8f73-2cb8cba0ddf8
📒 Files selected for processing (11)
python/cudnn/engines/manifest.pypython/cudnn/sdpa/__init__.pypython/cudnn/sdpa/bwd/__init__.pypython/cudnn/sdpa/bwd/engines.pypython/cudnn/sdpa/fwd/__init__.pypython/cudnn/sdpa/fwd/engines.pypython/cudnn/sdpa/graph_analyzer.pytest/python/fe_api/sdpa/test_sdpa_bwd_sm80.pytest/python/fe_api/sdpa/test_sdpa_fwd_sm80.pytest/python/sdpa/frost/test_sdpa_graph_analyzer.pytest/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py
🚧 Files skipped from review as they are similar to previous changes (4)
- test/python/sdpa/frost/test_sdpa_graph_analyzer.py
- test/python/fe_api/sdpa/test_sdpa_fwd_sm80.py
- python/cudnn/sdpa/bwd/engines.py
- python/cudnn/sdpa/fwd/engines.py
|
Determinism failures root-caused and fixed in Root cause: the mhas harness's cudnn handle carries raw stream handle Fix: Verified in the CI container: the failing IDs pass, and Known residual (filed separately, not CI-relevant): on NGC torch, the first execute after re-pointing a handle to a brand-new stream still no-ops once on an already-compiled kernel — suspected DSL/tvm-ffi launch-state caching (could not be isolated further; these kernels cannot compile without tvm-ffi). @cudnn-ci-bot run frost,oss,python_tests |
| seq_kv_t: Any = None | ||
| seq_q_t: Any = None | ||
| # backward-only refs | ||
| # Feature operands (bias / block-mask / score-stat outputs). |
There was a problem hiding this comment.
given that nothing supports this today, lets just remove from this too.
There was a problem hiding this comment.
I'd push back on this one: these aren't dead fields. They're real pygraph SDPA ports — python/pygraph/sdpa.cpp binds bias, block_mask, score_max (→ set_logit_max), and score_sum_exp — and the SM80 rows in this PR claim and serve all four from graphs (Capabilities.bias/block_mask/score_max/score_sum_exp, bound through SdpaBinding in both lowerings). bias/dBias is covered at graph level in test_sdpa_graph_analyzer.py and the integration suite. Dropping the facts would make an engine that never writes score_max pass the probe and leave the caller's buffer as garbage — the exact failure mode the _record_from_node comment warns about. If you'd rather the SM80 rows not claim the score-stat outputs yet, I can flip those two Capabilities bits off instead — but the analyzer facts need to exist for the probe to reject them.
There was a problem hiding this comment.
Resolving this in favour of keeping the facts, so the PR can land — and flagging the cleanup rather than blocking on it.
They are read on the execute path, not just declared: fwd/engines.py:585 copies out["score_max"] / out["score_sum_exp"] back through facts.score_max_t / facts.score_sum_exp_t, and block_mask_t is passed at fwd/engines.py:547 and bwd/engines.py:477. So removing them would drop the write-back, which is the failure mode @egilliam-nv described.
The asymmetry is what decides it for me: the request here is to delete, so merging as-is keeps the working path and the only cost of being wrong is a few unused fields. Deleting and being wrong costs a silently unwritten caller buffer. Cleanup is cheap later and reversible; the other direction is not.
@vedaanta if you still want the surface narrowed, the fallback @egilliam-nv offered — turn the SM80 Capabilities bits off and keep the analyzer facts — gets you "nothing claims it today" while the probe can still reject an engine that does not write those outputs. Happy to do that as a follow-up.
|
Addressed the review in Re-verified on A100: import-boundaries/dispatch/parity (67), full fe_api SM80 suites (105), frost SM80 integration + stream-respect + analyzer (86) — all green. |
|
@cudnn-ci-bot run frost,oss,python_tests |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-493-5f95a8a |
CuTe-DSL SDPA forward and backward for SM80 as FROST engines, plus standalone cudnn.sdpa APIs. * sdpa_fwd_prefill_sm80 joins the FrostSdpaFwdEngines family (the manifest row's sm_lo drops to 80); sdpa_bwd_sm80 introduces the first backward opset as a new FrostSdpaBwdEngines family on the reserved FROST_SDPA_BWD_ID_BASE block, with a frost_sdpa_bwd manifest row anchored on SDPA_BWD. * The shared graph_analyzer learns sdpa_backward() graphs: backward facts, K/V transposed-input-view canonicalization, and a forward-direction gate in the shared mismatch(). * Kernels (fwd generic + d256; bwd generic + d64 fast path) build on the shared frost/tile_dsl library; torch-native host code, per-shape self-caching, stream-aware and CUDA-graph-capturable. * Standalone SdpafwdSm80 / SdpabwdSm80 APIBase adapters + wrappers, including a packed-THD path the engines do not expose yet. * Docs: FE OSS API pages + Attention.md sections. Verified on A100: reference suites (103), engine/analyzer/stream suites (123), and full test_mhas_v2 with FROST auto-selection: 1268 passed / 0 failed, 41.2% of graphs served on FROST. The kernels originate from earlier internal work by Roman Anders. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* drop the stray fe_cuda_13.3.cfg (local build config, added by accident) * THD forward: resolve the default softmax scale from the user's head dim BEFORE flavor padding (a d=96 call silently used 1/sqrt(128)) * THD forward/backward: reject dense-only features (bias / RoPE / block_mask / seq lens / scale_output / scheduler) instead of silently computing without them; drop the never-used backward max_s_q kwarg * dense mask resolution: window_size=(-1, r) without is_causal is now rejected instead of silently selecting a 0-token SWA window * bprop kernel: assert K's head dim matches Q's, and reject RoPE at d_qk > 128 (the sDQ SMEM staging exceeds the A100 budget beyond that) * bwd reference in the fe_api suite anchors the causal diagonal top-left, matching the mask the wrapper actually requests * docs: requirements sections no longer cite the retired ctm package; provenance notes restored in both kernel packages; THD feature gating documented * stale comments: flavor envelope (dsv3->qwen), d64 'not yet routed', gptoss draft artifact, prefill header diagram / swizzle note / mainloop wait counts, d256 wait count
The previous commit unstaged it, then a blanket git add -A swept the untracked local file back in before committing. Now excluded locally.
* bwd dense mask resolution: reject window_size=(-1, r) without is_causal instead of silently selecting a 0-token SWA window (the fwd adapter and both THD paths already resolve it this way) * bwd THD: resolve the default softmax scale from the user's head dim BEFORE flavor padding (same silent-wrong-gradients bug as the fwd twin) * fe_api sweeps (48 cases each) move from L0 to L2 per the repo guideline; a single representative smoke case stays at L0 for each direction * underscore the unused unpacks Ruff flags in the fwd reference
…dependency of anything in this PR
…dates cutlass.experimental The package imports are lazy (PEP 562) since the engine-family cleanup, so a missing or old nvidia-cutlass-dsl no longer fails at wrapper import — it erupts at kernel-load time mid-test. The oss:rel CI leg (older DSL) showed 9 such errors; probe cutlass.experimental in the module skip so those environments skip cleanly.
* resolve_feature_operands: a graph that BINDS seq_len_q must get its buffer through _need() like every other feature operand — silently omitting it would execute a different query-length contract (missed BR-under-padding base and padded-row LSE trim) * Ruff RUF059 (two unused unpacks) and E731 (lambda-to-def)
The rebase onto the engine-family cleanup appended our helper block including copies of resolve_variant_pack (identical, harmless) and tensor_desc_from_ir (STALE: still referencing the deleted _DTYPE_FROM_CUDNN map). The stale copy shadowed upstream's fixed one and took down every SM100/SM120 lowering with a NameError — invisible on A100, where those paths skip; pipeline 61601513's Blackwell frost leg caught it (716 failures).
…minism zeros) Root cause of the frost_tests:sdpa[Ampere] determinism failures (104/104 is_determin backward configs, second execution returning all-zero grads): the harness's cudnn handle carries raw stream 0, and _stream_ctx wrapped it in torch.cuda.ExternalStream(0). On the CI image's NGC torch build (2.12.0a0), every kernel launch inside that context after the compile run silently no-ops. Reproduced and verified in the CI container itself (gitlab/cudnn_frontend:cudnn_13.3.0 + the pipeline's build artifact): before — run 0 correct, runs 1+ all-zero; after — bitwise-identical runs, test_sdpa_random_bwd_L0 176 passed / 0 failed under -n 4. _stream_ctx now maps a raw handle equal to torch's current/default stream onto that torch stream object and reserves ExternalStream for genuine foreign streams — the same guard fwd/api_dsl._torch_stream_context and gemm/cutedsl/grouped/backend_utils.py already carry (these adapters were the only unguarded spot in the tree). Known residual, unrelated to CI: on NGC torch the FIRST execute after re-pointing a handle to a brand-new stream still no-ops once on a cached kernel (suspected DSL/tvm-ffi launch-state caching; to be reported upstream — the kernels cannot compile without tvm-ffi, so it could not be isolated further).
…drop the docs pages Per maintainer feedback on the PR: - SdpafwdSm80/SdpabwdSm80 + wrappers move from api_sm80.py into each opset's api.py, following the SM100 classes there (one api.py per opset, no per-arch files). Lazy exports and the engine lowerings repoint; no signature changes. - The per-flavor kernel configs leave kernels/ for parent-level config_sm80.py (one per direction), matching config_sm100/config_sm120. - The FE OSS docs pages and their overview/Attention/llms.txt entries are dropped for now; the engines remain the PR's product. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ced its own copy) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Rebased onto latest Previous pipeline 62053314 on |
|
@cudnn-ci-bot run frost,oss,python_tests |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-493-ddfbacc |
YangXu1990uiuc
left a comment
YangXu1990uiuc
left a comment
There was a problem hiding this comment.
Approving to unblock — this has been open since 08-05 and the one thing still outstanding is safe in the merge direction.
Scope of what I checked, so the approval is not read as broader than it is:
- The one unresolved thread. Verified in the code rather than taking either side's word:
score_max_t/score_sum_exp_tare read atfwd/engines.py:585to copy kernel outputs back into the caller's buffers, andblock_mask_tatfwd/engines.py:547andbwd/engines.py:477. They are live, not dead fields. And the request was to delete them, so merging as-is keeps the working path; the cost of being wrong that way is a few unused fields, against a silently unwritten output buffer the other way. Left a note on the thread with the narrower fallback if @vedaanta still wants the surface reduced. - CI is not this PR's problem. Pipeline 62367995 on
ddfbacc9chas 10 failing jobs and I diffed every one against the develop nightly (62325526): all 10 fail there too, zero unique to this PR. Three arejob_execution_timeout(oss:rubin,py_test:*[Ampere]), the rest chronicscript_failureincludingpy_test:dev:[Blackwell], which is broken on develop independently of this work. - 34 of the 35 review threads were already resolved, and the branch was rebased through #547 two days ago.
What I did NOT do: review the 10.5k lines of new SM80 kernel and adapter code on their merits. That rests on @vedaanta's earlier passes and on the suite.
Follow-up worth tracking separately, not blocking: narrowing the SM80 Capabilities bits if the score-stat outputs should not be claimed yet.
YangXu1990uiuc
left a comment
YangXu1990uiuc
left a comment
There was a problem hiding this comment.
approving to unblock, and avoid more merge conflicts, clean up can be deferred to next change
4 participants
Before submitting
pre-commit runand committed any formatting changes.cat-*, one or moremod-*, and oneorig-*(see label list).Affected area
FE OSS kernels or CuTeDSL
Summary
SM80 (A100) SDPA forward and backward as FROST engines, plus standalone
cudnn.sdpaAPIs. This extends the FROST engine work merged in #476 to Ampere — previously the manifest had no engine below SM100.python/cudnn/sdpa/{fwd,bwd}/kernels/): CuTe-DSL SM80 prefill (generic + d=256) and backward (generic + d=64 fast path, ~2x for plain dense MHA), building on the sharedcudnn/frost/tile_dsl/library. The only kernel dependency isnvidia-cutlass-dsl(+apache-tvm-ffi). The kernels originate from earlier internal work by Roman Anders and are maintained in-tree from here on (provenance notes inbwd/kernels/__init__.py).sdpa_fwd_prefill_sm80joins theFrostSdpaFwdEnginesfamily (engine-id offset 8, after Add SM120 per-tensor FP8 (e4m3) SDPA-forward engine #509's SM120 fp8 row; the fwd manifest row'ssm_lodrops to 80), andsdpa_bwd_sm80joins theFrostSdpaBwdEnginesfamily Add SM120 FROST SDPA backward engine (sdpa_bwd_sm120) #486 introduced (offset 1; thefrost_sdpa_bwdmanifest row widens from SM120-only tosm_lo=80). The shared backwardCapabilities/mismatchgeneralize for the second row: head-dim envelopes withd_qk >= d_v(rectangular 192/128),dense_flexlayouts, and decode / top-left-rectangular-causal /S_q > S_kvgates — all defaulting to the SM120 row's current behavior. Both engines are on the unified dispatch model (BaseEngineadapters,graph.plans/select_plan,ExecutionContextstreams) and CUDA-graph-capturable.graph_analyzer.pygains the feature-operand facts (bias / block-mask / score-stat refs) and the shared adapter helpers on top of Add SM120 FROST SDPA backward engine (sdpa_bwd_sm120) #486'ssdpa_backward()support.SdpafwdSm80/SdpabwdSm80+ wrappers in each opset'sapi.py(following the SM100 classes there; per-flavor kernel configs inconfig_sm80.py, matchingconfig_sm100/config_sm120), including a packed-THD path the engines do not expose yet. Per review, no FE OSS docs pages for now — the engines are the PR's product.Feature surface (engines)
fp16/bf16 (both tested), GQA/MQA, head dims ≤ (256, 256) — any multiple of 1, via kernel-flavor envelopes incl. 192/128;
d_qk >= d_von backward — causal / SWA / bottom-right (incl. BR+SWA) / padded (seq_len_kv+seq_len_q) / causal right-band widening, bias(+dBias), ALiBi, sinks(+dSink), block_mask, score_max/score_sum_exp, deterministic dQ,dense_flexlayouts, stream-aware + CUDA-graph-capturable.Not supported
SM80-specific: decode-shaped graphs (
s_q == 1); engine-side THD/ragged (the standalone wrappers serve THD); head dims beyond the (256, 256) envelope; FP8/MXFP8 (out of scope on SM80, as in the cuDNN backend).FROST-wide (no engine row supports these; graphs fall back to the backend): dropout /
rng_dump, paged attention,score_mod, tensor-valuedattn_scale,cu_seq_len_q/kv(declined at the analyzer level), RoPE-fused multi-node graphs (the sdpa analyzer is single-node by design).Why
FROST engine coverage currently starts at SM100; A100 fleets get no FROST serving at all (the frost coverage stats report 0% on Ampere). This adds the Ampere forward row and the second backward row (alongside #486's SM120) — with a substantially wider feature envelope: GQA, deterministic dQ, dBias/dSink, bias/ALiBi/block-mask, SWA/bottom-right/padded masks, and rectangular head dims.
Related issues
Related to #476 (FROST engines), #486 (SM120 backward family this PR's backward row joins), #484 (execute contract), #488/#509 (fwd engine-id offsets).
API and compatibility impact
New opt-in engines only (
CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1); no default-path behavior change. New public symbols:SdpafwdSm80/SdpabwdSm80+sdpa_fwd_wrapper_sm80/sdpa_bwd_wrapper_sm80(experimental FE OSS APIs). Requires an SM80 device andnvidia-cutlass-dslfor the kernels to be offered; without them the engines decline and planning proceeds on the backend as before.Execute-contract status (AGENTS.md Rule 1)
Compliant: no
.to()/ copyingreshapeon execute arguments anywhere in the lowerings (outputs bind bycopy_, which casts in place; stats binds as a strictview— the probe admits only the contiguous(B, H_q, S_q, 1)layout, mirroring #486); seq-lens/sinks dtypes are facts-gated (int32/fp32) so the feature-operand resolution is pure views. Known deviation: these lowerings are torch-native and allocate at execute fordense_flexnormalization (gather copy to BSHD-physical when the caller's layout is not already BSHD), GQA head expansion on the forward path, and the adapters' internal scratch (head-dim padding, kernel-layout O/LSE staging) — a cached-allocator hit per execute rather than caller-workspace carving. CUDA-graph capture is verified green despite this (the stream-respect suites capture and replay both directions). Converting the adapters to thescratch_workspace_bytes()carving contract is planned as a follow-up; the engines are opt-in until then.Testing
On A100 (CUDA 13,
nvidia-cutlass-dsl==4.7.0a0), all viapytestfromtest/python:fe_api/sdpa/test_sdpa_fwd_sm80.py+test_sdpa_bwd_sm80.py: reference suites across all four kernel flavors, fp16 + bf16, masks, GQA, deterministic-dQ repeatability, d64 fast-path agreement, padded-row LSE trim (both kernels, tile-aligned + unalignedS_q) — 103 passed.sdpa/frost/: engine integration,SDPA_BWDanalyzer unit tests, stream-respect + CUDA-graph capture (fwd and bwd), plus the existing router/frost integration suites with the SM80 rows present — 123 passed.test_mhas_v2.pywith FROST auto-selection enabled: 1268 passed / 0 failed; 41.2% of graphs served on FROST (sdpa_fwd_prefill_sm80=605,sdpa_bwd_sm80=192— every backward graph the suite generates), the rest on the backend (fwd=1137). This suite caught two kernel-level bugs during development, both fixed in this PR:dense_flexdelivery (lowerings normalize to BSHD-physical buffers) and a padded-row LSE trim missing on theis_even_mnfast store path.🤖 Generated with Claude Code
Summary by CodeRabbit