Add SM120 FROST SDPA backward engine (sdpa_bwd_sm120) - #486
Conversation
|
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:
📝 WalkthroughWalkthroughChangesSM120/SM121 FROST SDPA backward support now includes graph analysis, capability routing, DSL execution, fused dQ/dK/dV kernels, manifest registration, shared tile-DSL helpers, and validation. SDPA backward support
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Graph
participant Engine
participant DSL
participant Kernel
Graph->>Engine: analyze backward tensors and features
Engine->>DSL: lower and compile SM120 plan
DSL->>Kernel: launch dot, main, and dQ conversion kernels
Kernel-->>DSL: write dQ, dK, and dV
DSL-->>Engine: return gradients and workspace state
Possibly related PRs
Suggested labels: 🚥 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: 7
🧹 Nitpick comments (6)
python/cudnn/engines/manifest.py (1)
180-191: 📐 Maintainability & Code Quality | 🔵 TrivialRow is consistent with the reserved ID block and the SM encoding.
id_hi = FROST_SDPA_BWD_ID_BASE + 100makes the row own 20_400..20_499, which matches theengine_ids.pycomment, and it does not overlap the forward block 20_300..20_399.sm_lo=120/sm_hi=121follows themajor * 10 + minorconvention used by thefrost_gemmrow.Line 187 records a TODO to widen the architecture bounds when an SM100 or SM80 spec lands. Do you want me to open an issue to track it?
🤖 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/engines/manifest.py` around lines 180 - 191, Retain the FrostSdpaBwdEngines row as implemented; no code change is required. The reserved ID range and SM encoding are consistent, and the TODO in the FrostSdpaBwd configuration may remain until SM100 or SM80 specifications are available.python/cudnn/sdpa/graph_analyzer.py (1)
211-221: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUpdate the docstring to cover backward nodes.
Line 220 now accepts
SDPA_BWD, but the docstring still describes only an SDPA-forward node. The predicate is shared by the forward and backward engine registries, so the stale text misstates the routing contract.♻️ Proposed docstring fix
- """The graph's sole SDPA-forward node, or None if the graph is anything else.""" + """The graph's sole SDPA node (forward, backward, MXFP8 or FP8), or None + if the graph is anything else."""🤖 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 211 - 221, Update the _single_sdpa_node docstring to describe the graph’s sole supported SDPA node, including both forward and backward variants accepted by the node_type predicate. Keep the routing behavior unchanged.test/python/sdpa/frost/test_sdpa_graph_analyzer.py (1)
825-827: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
bias=Truesets bothbiasanddBias, so this test does not isolate dBias.
_mk_bwd_graph(bias=True)adds abiasinput and adBiasoutput. Inmismatch()thehas_dbiasgate runs before thehas_biasgate, so the rejection could come from either capability. Split the helper flag intobiasanddbiasto cover the two gates separately.♻️ Proposed split of the bias flag
- bias: bool = False, + bias: bool = False, + dbias: bool = False, **bwd_kwargs, ): @@ - if bias: + if bias or dbias: bias_t = g.tensor(dim=(1, H, s_q, s_kv), stride=(H * s_q * s_kv, s_q * s_kv, s_kv, 1), data_type=DTYPE, name="bias") - dbias_t = g.tensor(dim=(1, H, s_q, s_kv), stride=(H * s_q * s_kv, s_q * s_kv, s_kv, 1), data_type=DTYPE, name="dBias") - bwd_kwargs.update(bias=bias_t, dBias=dbias_t) + bwd_kwargs.update(bias=bias_t) + if dbias: + dbias_t = g.tensor(dim=(1, H, s_q, s_kv), stride=(H * s_q * s_kv, s_q * s_kv, s_kv, 1), data_type=DTYPE, name="dBias") + bwd_kwargs.update(dBias=dbias_t)def test_bwd_probe_rejects_dbias(monkeypatch): monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) - assert not _bwd_eligible(_mk_bwd_graph(bias=True)) + assert not _bwd_eligible(_mk_bwd_graph(bias=True, dbias=True)) + assert not _bwd_eligible(_mk_bwd_graph(bias=True))🤖 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_graph_analyzer.py` around lines 825 - 827, Update _mk_bwd_graph to accept independent bias and dbias flags, wiring them to the bias input and dBias output respectively. Adjust test_bwd_probe_rejects_dbias to enable only dbias, and add or update a separate test to verify bias rejection with only bias enabled, so each mismatch() gate is isolated.test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py (1)
173-178: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove the unconditional skip to the test that triggers it.
_run_bwd_graphskips wheneverq_tileorkv_tileis set.test_sdpa_bwd_dsl_sm120_tile_knobsalways sets both, so both parameterizations always skip. The skip reason is hidden in a helper, and the helper builds the whole graph before it skips. Mark the test itself with@pytest.mark.skipand drop theq_tile/kv_tileplumbing from_run_bwd_graphand_run_caseuntil the knob plans exist.♻️ Proposed move of the skip to the test site
- if q_tile is not None or kv_tile is not None: - pytest.skip( - "graph.set_engine_knobs() was removed with the monkey-patch dispatch layer and has no replacement in " - "this MR: knobs now ride on the plan (engines.base.PlanConfig.knobs), one ranked entry per knob set, " - "picked with select_plan(). Re-enable once the SDPA bwd family proposes its tile domain as plans." - ) - graph.validate()`@pytest.mark.L0` +@pytest.mark.skip( + reason="graph.set_engine_knobs() was removed with the monkey-patch dispatch layer and has no replacement in " + "this MR: knobs now ride on the plan (engines.base.PlanConfig.knobs), one ranked entry per knob set, " + "picked with select_plan(). Re-enable once the SDPA bwd family proposes its tile domain as plans." +) `@pytest.mark.parametrize`(🤖 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_bwd_dsl_sm120.py` around lines 173 - 178, Move the unconditional skip to test_sdpa_bwd_dsl_sm120_tile_knobs using pytest.mark.skip with the existing reason, and remove the q_tile/kv_tile parameters and related skip logic from _run_bwd_graph and _run_case. Keep non-knob test paths unchanged while allowing the skipped test to be marked before graph construction.python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py (2)
304-312: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRaise a clear error for an unsupported
head_dim.Line 304 indexes
DEFAULT_TILESdirectly. An unsupportedhead_dimproduces a bareKeyError. Every neighboring constraint in this constructor reports a descriptiveValueError(Line 312 and Line 329), andvalidate_paramsexists as a backstop for direct template use. Make this failure consistent.♻️ Proposed refactor
+ if head_dim not in self.DEFAULT_TILES: + raise ValueError(f"head_dim must be one of {tuple(self.DEFAULT_TILES)}; got {head_dim}") self.q_tile, self.kv_tile = self.DEFAULT_TILES[head_dim]🤖 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_sm120.py` around lines 304 - 312, Update the constructor initialization around DEFAULT_TILES so an unsupported head_dim is detected before direct dictionary indexing and raises a descriptive ValueError instead of propagating KeyError. Preserve the existing tile overrides and validation behavior for supported head dimensions, and keep validate_params as the backstop for direct template use.
1176-1184: 🚀 Performance & Scalability | 🔵 TrivialConsider a bound and a metric for the compile cache.
@lru_cache(maxsize=None)never evicts. Each entry holds three compiled CUDA modules for one(compute_capability, b, qh, sq, skv, d)tuple. A serving process that sees many distinct batch sizes and sequence lengths grows host and device memory for the lifetime of the process. The PR notes 14 OOMs at 32 GB attributed to a pre-existing capacity issue; an unbounded kernel cache adds to that pressure.A bound here alone is not sufficient, because
_wrapper_api_cacheinpython/cudnn/sdpa/bwd/api_dsl.pyLine 369 is also unbounded and holds references to the compiled objects. Evicting one and not the other causes silent recompilation. Treat the cache size as one cross-layer decision, and export the entry count so operators can see the growth.🤖 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_sm120.py` around lines 1176 - 1184, Bound the compile cache used by compile and coordinate the same entry limit with _wrapper_api_cache so evicting a kernel entry also releases its corresponding wrapper references instead of triggering silent recompilation. Define the cache capacity as a shared configurable/exported metric, expose the current entry count for operators, and update both cache implementations to use that single cross-layer limit.
🤖 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_dsl.py`:
- Around line 105-110: Update _to_bshd and the output handling for dq_tensor,
dk_tensor, and dv_tensor so non-contiguous output views are rejected instead of
converted with contiguous(). Keep copying permitted for input tensors, but make
execute-time output validation fail clearly before the kernel writes, regardless
of cached descriptors or sample tensors.
- Around line 14-26: Add the SM120 backward API symbols SdpaBwdDslSm120 and
sdpa_bwd_wrapper_dsl_sm120 to the optional lazy-export mapping used by the cudnn
package, ensuring resolving those exports does not eagerly import api_dsl.py.
Document both exports under docs/fe-oss-apis/ and add pytest coverage under
test/python/fe_api/ verifying lazy resolution and the exported API names.
In `@python/cudnn/sdpa/bwd/engines.py`:
- Line 374: Sort the exports in the __all__ declaration alphabetically to
satisfy RUF022, preserving all existing names and their spelling.
In `@python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py`:
- Around line 1-2: Align the SPDX license header in the bprop_f16_sm120 module
with the confirmed package license, matching the sibling files if MIT applies.
If the FlashAttention-2 attribution requires BSD-3-Clause coverage, preserve the
applicable package license and add explicit BSD-3-Clause attribution rather than
substituting Apache-2.0.
In `@python/cudnn/sdpa/graph_analyzer.py`:
- Around line 311-320: Update the backward K/V canonicalization block to detect
native (B, H, D, S) views using K/V stride ordering rather than only dimension
comparisons. In the is_backward path, adjust k_dim/k_stride and v_dim/v_stride
when their strides identify the transposed layout, including square cases where
S_kv equals d_qk or d_v, while preserving canonical BSHD dimensions and strides
for bshd_layout_ok.
In `@test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py`:
- Around line 310-319: Make test_sdpa_bwd_dsl_sm120_stream_respect observably
verify stream usage instead of relying on _run_bwd_graph’s device-wide
synchronization. Add stream-ordering instrumentation, such as blocking stream
with a long-running kernel and checking the default stream remains complete, or
recording an event on stream and confirming the engine work is captured;
preserve the test’s no-default-stream-leak assertion.
- Around line 238-256: Update _run_case to assert that the selected-engine test
actually uses the FROST engine: when select=True, require plan_name to equal
ENGINE, alongside the existing workspace validation. Keep the numeric assertions
and behavior for non-selected runs unchanged.
---
Nitpick comments:
In `@python/cudnn/engines/manifest.py`:
- Around line 180-191: Retain the FrostSdpaBwdEngines row as implemented; no
code change is required. The reserved ID range and SM encoding are consistent,
and the TODO in the FrostSdpaBwd configuration may remain until SM100 or SM80
specifications are available.
In `@python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py`:
- Around line 304-312: Update the constructor initialization around
DEFAULT_TILES so an unsupported head_dim is detected before direct dictionary
indexing and raises a descriptive ValueError instead of propagating KeyError.
Preserve the existing tile overrides and validation behavior for supported head
dimensions, and keep validate_params as the backstop for direct template use.
- Around line 1176-1184: Bound the compile cache used by compile and coordinate
the same entry limit with _wrapper_api_cache so evicting a kernel entry also
releases its corresponding wrapper references instead of triggering silent
recompilation. Define the cache capacity as a shared configurable/exported
metric, expose the current entry count for operators, and update both cache
implementations to use that single cross-layer limit.
In `@python/cudnn/sdpa/graph_analyzer.py`:
- Around line 211-221: Update the _single_sdpa_node docstring to describe the
graph’s sole supported SDPA node, including both forward and backward variants
accepted by the node_type predicate. Keep the routing behavior unchanged.
In `@test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py`:
- Around line 173-178: Move the unconditional skip to
test_sdpa_bwd_dsl_sm120_tile_knobs using pytest.mark.skip with the existing
reason, and remove the q_tile/kv_tile parameters and related skip logic from
_run_bwd_graph and _run_case. Keep non-knob test paths unchanged while allowing
the skipped test to be marked before graph construction.
In `@test/python/sdpa/frost/test_sdpa_graph_analyzer.py`:
- Around line 825-827: Update _mk_bwd_graph to accept independent bias and dbias
flags, wiring them to the bias input and dBias output respectively. Adjust
test_bwd_probe_rejects_dbias to enable only dbias, and add or update a separate
test to verify bias rejection with only bias enabled, so each mismatch() gate is
isolated.
🪄 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: 0dbb453b-615a-42ee-85d7-230b33b50ef7
📒 Files selected for processing (16)
python/cudnn/engines/engine_ids.pypython/cudnn/engines/manifest.pypython/cudnn/frost/tile_dsl/mma.pypython/cudnn/frost/tile_dsl/swizzle.pypython/cudnn/sdpa/bwd/__init__.pypython/cudnn/sdpa/bwd/api_dsl.pypython/cudnn/sdpa/bwd/config_sm120.pypython/cudnn/sdpa/bwd/engine.pypython/cudnn/sdpa/bwd/engines.pypython/cudnn/sdpa/bwd/kernels/__init__.pypython/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.pypython/cudnn/sdpa/fwd/engines.pypython/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.pypython/cudnn/sdpa/graph_analyzer.pytest/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.pytest/python/sdpa/frost/test_sdpa_graph_analyzer.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 `@test/python/sdpa/frost/test_sdpa_graph_analyzer.py`:
- Around line 833-835: Update test_bwd_probe_rejects_dbias to isolate the dBias
condition instead of relying on _mk_bwd_graph(dbias=True), which also adds bias.
Construct a dBias-only graph if supported, or assert the rejection diagnostic
specifically identifies dBias, ensuring the test fails when only bias causes
rejection.
🪄 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: 862e94eb-60f2-4385-90c8-dca15afe9c25
📒 Files selected for processing (6)
python/cudnn/sdpa/bwd/api_dsl.pypython/cudnn/sdpa/bwd/engines.pypython/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.pypython/cudnn/sdpa/graph_analyzer.pytest/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.pytest/python/sdpa/frost/test_sdpa_graph_analyzer.py
🚧 Files skipped from review as they are similar to previous changes (4)
- python/cudnn/sdpa/bwd/api_dsl.py
- python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py
- python/cudnn/sdpa/graph_analyzer.py
- python/cudnn/sdpa/bwd/engines.py
b3a83ee to
48785e4
Compare
|
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. |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
python/cudnn/sdpa/bwd/api_dsl.py (1)
379-431: 🚀 Performance & Scalability | 🔵 Trivial | ⚖️ Poor tradeoff
_wrapper_api_cachegrows without a bound.The key contains the full shape, stride, dtype, and device signature of six tensors plus the causal flags and the scale. Each miss stores a
SdpaBwdDslSm120instance that owns compiled kernels. A workload with many distinct sequence lengths retains every specialization for the process lifetime. Consider a bounded cache with eviction, for examplefunctools.lru_cacheon a keyed factory or an explicitOrderedDictwith a size limit.🤖 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_dsl.py` around lines 379 - 431, Bound the specialization cache used by sdpa_bwd_wrapper_dsl_sm120 so it evicts older entries instead of retaining every SdpaBwdDslSm120 instance indefinitely. Preserve the existing full cache key and cache-miss behavior, using a fixed size limit via an LRU cache or equivalent bounded eviction mechanism around the factory.
🤖 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/bwd/api_dsl.py`:
- Around line 379-431: Bound the specialization cache used by
sdpa_bwd_wrapper_dsl_sm120 so it evicts older entries instead of retaining every
SdpaBwdDslSm120 instance indefinitely. Preserve the existing full cache key and
cache-miss behavior, using a fixed size limit via an LRU cache or equivalent
bounded eviction mechanism around the factory.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 024057ac-9482-4559-a92a-fe5fadfafd3d
📒 Files selected for processing (16)
python/cudnn/engines/engine_ids.pypython/cudnn/engines/manifest.pypython/cudnn/frost/tile_dsl/mma.pypython/cudnn/frost/tile_dsl/swizzle.pypython/cudnn/sdpa/bwd/__init__.pypython/cudnn/sdpa/bwd/api_dsl.pypython/cudnn/sdpa/bwd/config_sm120.pypython/cudnn/sdpa/bwd/engine.pypython/cudnn/sdpa/bwd/engines.pypython/cudnn/sdpa/bwd/kernels/__init__.pypython/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.pypython/cudnn/sdpa/fwd/engines.pypython/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.pypython/cudnn/sdpa/graph_analyzer.pytest/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.pytest/python/sdpa/frost/test_sdpa_graph_analyzer.py
🚧 Files skipped from review as they are similar to previous changes (11)
- python/cudnn/sdpa/bwd/kernels/init.py
- python/cudnn/sdpa/bwd/init.py
- python/cudnn/frost/tile_dsl/swizzle.py
- python/cudnn/engines/manifest.py
- python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py
- python/cudnn/frost/tile_dsl/mma.py
- python/cudnn/sdpa/bwd/engine.py
- python/cudnn/engines/engine_ids.py
- python/cudnn/sdpa/fwd/engines.py
- test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py
- python/cudnn/sdpa/graph_analyzer.py
|
Nice job! LGTM overall; @vedaanta may want to take a look as well. The next steps are to make the kernel perf-ready and add more features. By the way, you may run test_mhas_v2.py to see the coverage of the FROST sdpa_fwd kernel. |
|
@cudnn-ci-bot run |
|
cuDNN CI bot run targets
Targets are comma-separated: Aliases: Only allowlisted maintainers can use |
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-486-3ab04dc |
|
@cudnn-ci-bot run frost |
|
Pipeline not launched Reason: @Adnios is not allowlisted to run this bot. Ask an allowlisted maintainer to comment |
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-486-0f3bf39 |
4 participants
Before submitting
pre-commit runand committed any formatting changes.Affected area
Summary
Three-kernel chain (
dot_do_o→ main 5-GEMM KV-stationary kernel →convert_dq),PDL-linked, 12 warps (8 compute + 1 TMA producer), dQ accumulated through fp32
global atomics.
Why
Related issues
Related to #381
API and compatibility impact
Testing
test_sdpa_graph_analyzer.py: 63 passed (14 new GPU-free bwd probe/facts tests)test_sdpa_bwd_dsl_sm120.py(new): 14 passedtest_mhas_v2.py::test_sdpa_random_bwd_L0with the engine enabled:165 passed, 24.1% of graphs routed to FROST, zero numeric failures
(14 OOMs reproduce identically with FROST disabled — pre-existing 32 GB capacity issue)
Summary by CodeRabbit
New Features
Bug Fixes
Tests