sdpa: add torch custom ops cudnn::sdpa_fwd / cudnn::sdpa_bwd - #517
sdpa: add torch custom ops cudnn::sdpa_fwd / cudnn::sdpa_bwd#517vedaanta wants to merge 1 commit into
Conversation
vedaanta
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:
📝 WalkthroughWalkthroughAdds cuDNN PyTorch SDPA forward and THD backward operators. The implementation supports extended masking, sinks, padding, packed sequences, LSE output, graph caching, fake kernels, CUDA execution, and varlen autograd. Tests and documentation cover the new contracts. ChangesExtended SDPA execution
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant sdpa
participant sdpa_fwd
participant cuDNNGraph
Caller->>sdpa: provide q, k, v, scale, and options
sdpa->>sdpa_fwd: pass tensors and sequence metadata
sdpa_fwd->>cuDNNGraph: build or retrieve cached graph
cuDNNGraph->>sdpa_fwd: return output and optional LSE
sdpa_fwd->>Caller: return SDPA result
sequenceDiagram
participant Autograd
participant sdpa_fwd
participant sdpa_bwd
participant cuDNNGraph
Autograd->>sdpa_fwd: run varlen forward
sdpa_fwd->>Autograd: save output and packed LSE
Autograd->>sdpa_bwd: pass gradients and padded LSE
sdpa_bwd->>cuDNNGraph: execute THD backward graph
cuDNNGraph->>Autograd: return packed DQ, DK, and DV
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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: 16
🧹 Nitpick comments (4)
python/cudnn/sdpa/fwd/torch_op.py (4)
43-44: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winBound the graph cache.
_graph_cachegrows without limit. The key includesB,S_q,S_kv, and all strides. A serving workload with variable batch or sequence shapes creates a new built cuDNN graph plus its execution plans for every distinct shape, and nothing evicts them. This retains device-side plan state for the process lifetime.Consider an LRU bound (for example
functools.lru_cacheon a builder helper, or anOrderedDictwith a max size) before this graduates from prototype.🤖 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/torch_op.py` around lines 43 - 44, Bound the global _graph_cache used by the cuDNN graph builder so it evicts least-recently-used entries after a fixed maximum size. Update the cache access and insertion logic around _graph_cache to refresh recency on hits and remove the oldest graph, including its execution-plan state, when the limit is exceeded.
451-466: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe non-THD branch of
_build_bwd_graphis unreachable.
_sdpa_bwd_ex_implraisesNotImplementedErrorfor the dense path at line 542 and always passesis_thd=Trueat line 616. Sois_thdis a constant here. Theif is_thd:guard, theuse_padding_mask=is_thdat line 484, and the conditionalmax_total_seq_len_*at lines 487-488 are dead branches.Keeping the parameter is reasonable if dense backward lands soon. If not, drop it and simplify. Add a short comment either way so a reader does not assume the dense path is exercised.
🤖 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/torch_op.py` around lines 451 - 466, In _build_bwd_graph, document that the current backward implementation only reaches the THD path because _sdpa_bwd_ex_impl rejects dense inputs and invokes it with is_thd=True. If dense backward is not being added now, remove the unreachable is_thd branching, including the guard and dependent use_padding_mask and max_total_seq_len conditionals; otherwise retain the parameter but add a short comment identifying the current restriction.
350-361: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winGuard the int32 ragged offsets against overflow.
cu_seqlens_qis int32.cu_seqlens_q * q.stride(0)stays int32, so the product wraps silently for large packed buffers. With a kv-interleaved view the token stride doubles, which halves the safe token budget. Example:T = 131072,H = 64,D = 128, kv-packed token stride2 * H * D = 16384gives2.1e9, pastINT32_MAX. A wrapped offset produces a negative base pointer offset and reads outside the allocation.cuDNN ragged offsets are int32 by contract, so add an explicit check rather than widening.
🛡️ Proposed guard
if is_thd: + max_elem = max(T_q * q.stride(0), T_kv * k.stride(0), T_kv * v.stride(0), T_q * H_q * D_v) + if max_elem > 2**31 - 1: + raise ValueError(f"packed tensors exceed the int32 ragged-offset limit ({max_elem} elements); split the batch") # cuDNN ragged offsets are ELEMENT offsets per tensor, so each scalesThis requires keeping
T_qandT_kvbound at lines 256-257 instead of discarding them.🤖 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/torch_op.py` around lines 350 - 361, Retain the bound T_q and T_kv values in the setup around the existing sequence-shape handling, then validate that every computed ragged element offset fits within the signed int32 range before populating the variant. Guard the products used by RAGGED_Q, RAGGED_KV, RAGGED_V, RAGGED_O, and RAGGED_STATS, raising the established argument/validation error when any packed token count and stride would exceed INT32_MAX; keep cuDNN offsets int32 rather than widening them.
487-488: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winPass the actual token counts to
max_total_seq_len_qandmax_total_seq_len_kv.These attributes only size workspace. They do not define
dQ,dK, ordVextents. The cuDNN documentation does not require 64-token rounding. Usetotal_qandtotal_kvdirectly unless a version-specific engine requires this alignment.🤖 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/torch_op.py` around lines 487 - 488, Update the SDPA forward operation configuration to pass total_q directly as max_total_seq_len_q and total_kv directly as max_total_seq_len_kv when is_thd, removing the _round64 alignment for these workspace-sizing attributes while preserving the existing non-THD None behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudnn/sdpa/fwd/torch_op.py`:
- Around line 95-102: Update _normalize_thd in
python/cudnn/sdpa/fwd/torch_op.py:95-102 to use
clone(memory_format=torch.contiguous_format) instead of contiguous(), ensuring
misaligned tensors always receive fresh storage. At
python/cudnn/sdpa/fwd/torch_op.py:561-563, extend the lse normalization guard
with lse.data_ptr() % 16 and use the same clone allocation when contiguity or
alignment fails.
- Around line 256-265: Validate k.shape on both forward paths before building
its cuDNN descriptor: in python/cudnn/sdpa/fwd/torch_op.py lines 256-265,
require k.shape to equal (v.shape[0], H_kv, D_qk); in lines 267-276, require (B,
H_kv, S_kv, D_qk) and validate matching batch sizes across q, k, and v. Raise
ValueError on mismatches, and prefix unused T_q/T_kv bindings with underscores
unless they are used for the ragged-offset overflow guard.
- Around line 391-399: Update the fake-kernel allocations in
python/cudnn/sdpa/fwd/torch_op.py:391-399 to mirror the real o_stride selection
used around lines 272-275, allocating o with torch.empty_strided for both THD
and non-THD layouts, and make empty stats placeholders use float32 on q.device.
Also update python/cudnn/sdpa/fwd/torch_op.py:660-678 to replace
torch.empty_like(k) and torch.empty_like(v) with contiguous torch.empty
allocations using q.dtype, matching the freshly allocated gradients returned
around lines 628-629.
- Around line 686-728: Complete the frontend-only API packaging for sdpa_ex: add
its APIBase subclass and wrapper around the existing sdpa_ex implementation,
register a lazy export in python/cudnn/__init__.py, add documentation under
docs/fe-oss-apis/, and move or add pytest coverage under test/python/fe_api/.
Ensure the public wrapper preserves sdpa_ex’s current arguments and return
behavior.
- Line 62: Update the O assignment in the SDPA forward implementation to add an
inline Ruff E741 suppression, keeping the existing name and value unchanged.
- Around line 170-172: Validate the `causal_bottom_right` combination before
constructing `alignment`: reject calls where `is_causal` is false and
`window_left` is negative, since both `lb` and `rb` are unset and the flag would
be inert. Raise the existing argument-validation error type with a clear
message; preserve current behavior for all other combinations.
- Around line 475-493: Update _sdpa_bwd_ex_impl to explicitly reject backward
execution when attention sinks are present, since g.sdpa_backward cannot consume
them. Add the sinks argument to the backward schema and _sdpa_bwd_ex_fake so the
guard is reachable, and raise a clear error rather than computing incorrect dq,
dk, or dv.
- Around line 21-22: Update the module docstring in torch_op.py to state that
backward support is implemented and registered through cudnn::sdpa_bwd_ex, while
documenting that the backward contract is restricted to THD inputs. Remove the
forward-only wording and preserve the existing forward behavior description.
- Around line 337-343: Update both real and fake kernel paths around the THD and
non-THD `stats` allocations so the `return_lse=False` placeholder is created on
`q.device` with the expected stats dtype, matching the device and dtype of the
`return_lse=True` output. Apply the same change to the fake kernel definitions
to keep compiled output metadata consistent with the CUDA implementation.
- Around line 561-563: Update the lse handling to call contiguous() on the
original lse tensor before reshaping it to (B, H_q, S_q, 1), preserving the
existing defensive normalization while ensuring the base-pointer alignment
requirement is applied before reshape.
- Around line 686-703: Update sdpa_ex to reject calls when query, key, or value
requires gradients, since cudnn::sdpa_fwd_ex has no registered autograd formula.
Add the guard at the start of sdpa_ex before invoking the forward operation,
while preserving inference behavior for non-grad inputs; do not wire sdpa_bwd_ex
in this change.
- Around line 627-629: Update the SDPA forward/backward dtype validation around
_build_graph and the gradient allocations so q, k, and v must share the same
supported dtype before deriving io_dtype; raise a clear validation error for
mismatches or unsupported dtypes instead of allowing a mapping KeyError.
Allocate dq, dk, and dv using the validated common dtype, preserving the
existing device and shape behavior.
- Around line 70-74: Update _get_handle and the surrounding execution path to
use a thread-local cuDNN handle per device, or otherwise serialize all handle
rebinding and g.execute calls so concurrent threads cannot share a handle.
Synchronize first-use _graph_cache construction as well, preserving correct
stream association for every execution.
In `@test/python/test_cudnn_sdpa_ex_ops.py`:
- Line 36: In test/python/test_cudnn_sdpa_ex_ops.py, update the shape unpacking
at lines 36 and 168 to use `_` for the unused D and T bindings respectively;
leave the other destructured values unchanged.
- Around line 154-160: Extend the reference path used by _ref in the causal
sdpa_fwd_ex test to produce expected fp32 LSE values, then compare the returned
packed lse tensor against that reference using the existing dtype-appropriate
tolerance conventions. Retain the current shape, dtype, and finiteness checks
while adding value validation for every LSE entry.
- Around line 23-24: Extend the module-level gate before the torch_op import to
validate the cuDNN version, GPU compute capability, and BF16 support using
cudnn.backend_version() and torch.cuda.get_device_capability(). Skip the
extended SDPA tests when these requirements are unsupported, while retaining the
existing CUDA availability check.
---
Nitpick comments:
In `@python/cudnn/sdpa/fwd/torch_op.py`:
- Around line 43-44: Bound the global _graph_cache used by the cuDNN graph
builder so it evicts least-recently-used entries after a fixed maximum size.
Update the cache access and insertion logic around _graph_cache to refresh
recency on hits and remove the oldest graph, including its execution-plan state,
when the limit is exceeded.
- Around line 451-466: In _build_bwd_graph, document that the current backward
implementation only reaches the THD path because _sdpa_bwd_ex_impl rejects dense
inputs and invokes it with is_thd=True. If dense backward is not being added
now, remove the unreachable is_thd branching, including the guard and dependent
use_padding_mask and max_total_seq_len conditionals; otherwise retain the
parameter but add a short comment identifying the current restriction.
- Around line 350-361: Retain the bound T_q and T_kv values in the setup around
the existing sequence-shape handling, then validate that every computed ragged
element offset fits within the signed int32 range before populating the variant.
Guard the products used by RAGGED_Q, RAGGED_KV, RAGGED_V, RAGGED_O, and
RAGGED_STATS, raising the established argument/validation error when any packed
token count and stride would exceed INT32_MAX; keep cuDNN offsets int32 rather
than widening them.
- Around line 487-488: Update the SDPA forward operation configuration to pass
total_q directly as max_total_seq_len_q and total_kv directly as
max_total_seq_len_kv when is_thd, removing the _round64 alignment for these
workspace-sizing attributes while preserving the existing non-THD None behavior.
🪄 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: 2605e9b1-7ebf-406d-a6ab-0f0fab19f8ca
📒 Files selected for processing (2)
python/cudnn/sdpa/fwd/torch_op.pytest/python/test_cudnn_sdpa_ex_ops.py
| def sdpa_ex( | ||
| query: torch.Tensor, | ||
| key: torch.Tensor, | ||
| value: torch.Tensor, | ||
| *, | ||
| scale: Optional[float] = None, | ||
| is_causal: bool = False, | ||
| causal_bottom_right: bool = False, | ||
| window_left: int = -1, | ||
| sinks: Optional[torch.Tensor] = None, | ||
| seq_len_q: Optional[torch.Tensor] = None, | ||
| seq_len_kv: Optional[torch.Tensor] = None, | ||
| cu_seqlens_q: Optional[torch.Tensor] = None, | ||
| cu_seqlens_kv: Optional[torch.Tensor] = None, | ||
| max_seqlen_q: int = 0, | ||
| max_seqlen_kv: int = 0, | ||
| return_lse: bool = False, | ||
| ): | ||
| """cuDNN SDPA forward with the extended feature surface (see module docstring). | ||
|
|
||
| Returns ``o`` or ``(o, lse)`` when ``return_lse=True``. | ||
| """ | ||
| import math | ||
|
|
||
| attn_scale = scale if scale is not None else 1.0 / math.sqrt(query.shape[-1]) | ||
| o, lse = torch.ops.cudnn.sdpa_fwd_ex( | ||
| query, | ||
| key, | ||
| value, | ||
| attn_scale, | ||
| is_causal=is_causal, | ||
| causal_bottom_right=causal_bottom_right, | ||
| window_left=window_left, | ||
| sinks=sinks, | ||
| seq_len_q=seq_len_q, | ||
| seq_len_kv=seq_len_kv, | ||
| cu_seqlens_q=cu_seqlens_q, | ||
| cu_seqlens_kv=cu_seqlens_kv, | ||
| max_seqlen_q=max_seqlen_q, | ||
| max_seqlen_kv=max_seqlen_kv, | ||
| return_lse=return_lse, | ||
| ) | ||
| return (o, lse) if return_lse else o |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
Missing artifacts required for a new frontend-only Python API.
The coding guidelines require every new frontend-only Python API under python/cudnn/ to ship four artifacts. sdpa_ex currently has none of them:
- No
APIBasesubclass and wrapper. - No lazy export in
python/cudnn/__init__.py. The module docstring at lines 24-25 confirms this. - No documentation under
docs/fe-oss-apis/. - Tests live at
test/python/test_cudnn_sdpa_ex_ops.py, not undertest/python/fe_api/.
The PR description lists documentation and the lazy-table export as follow-ups. Track the APIBase wrapper and the test relocation as well, so the prototype does not graduate without them.
As per coding guidelines: "Every new frontend-only Python API must include an APIBase subclass and wrapper, a lazy export in python/cudnn/__init__.py, documentation under docs/fe-oss-apis/, and pytest coverage under test/python/fe_api/."
🤖 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/torch_op.py` around lines 686 - 728, Complete the
frontend-only API packaging for sdpa_ex: add its APIBase subclass and wrapper
around the existing sdpa_ex implementation, register a lazy export in
python/cudnn/__init__.py, add documentation under docs/fe-oss-apis/, and move or
add pytest coverage under test/python/fe_api/. Ensure the public wrapper
preserves sdpa_ex’s current arguments and return behavior.
Source: Coding guidelines
| (the cuDNN diagonal_band_left_bound convention). sinks: (H,) extra softmax | ||
| logit per query head, contributing no value.""" | ||
| q, k, v = q.float(), k.float(), v.float() | ||
| B, Hq, Sq, D = q.shape |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Remove unused destructured values.
Ruff reports RUF059 for both bindings.
test/python/test_cudnn_sdpa_ex_ops.py#L36-L36: replaceDwith_.test/python/test_cudnn_sdpa_ex_ops.py#L168-L168: replaceTwith_.
🧰 Tools
🪛 Ruff (0.16.1)
[warning] 36-36: Unpacked variable D is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
📍 Affects 1 file
test/python/test_cudnn_sdpa_ex_ops.py#L36-L36(this comment)test/python/test_cudnn_sdpa_ex_ops.py#L168-L168
🤖 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/test_cudnn_sdpa_ex_ops.py` at line 36, In
test/python/test_cudnn_sdpa_ex_ops.py, update the shape unpacking at lines 36
and 168 to use `_` for the unused D and T bindings respectively; leave the other
destructured values unchanged.
Source: Linters/SAST tools
| o, lse = torch.ops.cudnn.sdpa_fwd_ex( | ||
| q, k, v, D**-0.5, is_causal=True, | ||
| cu_seqlens_q=cu, cu_seqlens_kv=cu, max_seqlen_q=mx, max_seqlen_kv=mx, return_lse=True, | ||
| ) # fmt: skip | ||
| ref, _, _, _ = self._ref(q, k, v, cu, is_causal=True) | ||
| assert (o.float() - ref).abs().max().item() < TOL | ||
| assert lse.shape == (T, H, 1) and lse.dtype == torch.float32 and torch.isfinite(lse).all() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Compare returned LSE values with the fp32 reference.
Lines 159-160 validate only metadata and finiteness. Incorrect finite LSE values will pass. Extend _ref or add an LSE reference, then compare the packed lse values with dtype-appropriate tolerances.
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/test_cudnn_sdpa_ex_ops.py` around lines 154 - 160, Extend the
reference path used by _ref in the causal sdpa_fwd_ex test to produce expected
fp32 LSE values, then compare the returned packed lse tensor against that
reference using the existing dtype-appropriate tolerance conventions. Retain the
current shape, dtype, and finiteness checks while adding value validation for
every LSE entry.
Source: Coding guidelines
41beced to
b9e8330
Compare
vedaanta
changed the title
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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/experimental/ops/sdpa.py`:
- Around line 599-602: Update the downstream calls in
test_cudnn_sdpa_torch_ops.py to use torch.ops.cudnn.sdpa_bwd_legacy, matching
the registered operator declaration in the sdpa operation schema. Retain the
sdpa_bwd name only if an explicitly schema-compatible alias is required.
In `@test/python/test_cudnn_sdpa_torch_ops.py`:
- Around line 67-68: Add an L0 test under TestSdpaFwdExDense that supplies
unequal per-batch seq_len_q and seq_len_kv values to the dense SDPA operation,
masks invalid KV positions in the reference implementation, and compares outputs
only for valid Q positions using the suite’s existing reference-module pattern
and dtype-appropriate tolerances.
- Around line 76-78: Extend the reference helper _ref to return packed
log-sum-exp values, then update the dense and THD SDPA tests around sdpa_fwd to
request LSE output instead of disabling it. Compare each implementation’s LSE
against the reference using dtype-appropriate tolerances, while retaining the
existing output comparisons and finite-value checks.
🪄 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: baf0a09f-d8d9-42e7-b51f-88cdb6e6a2e7
📒 Files selected for processing (3)
python/cudnn/experimental/ops/sdpa.pypython/cudnn/sdpa/fwd/torch_op.pytest/python/test_cudnn_sdpa_torch_ops.py
| # Renamed from cudnn::sdpa_bwd: the canonical name now belongs to the | ||
| # consolidated op family in cudnn.sdpa.fwd.torch_op (this experimental | ||
| # module is slated to fold into it). | ||
| "sdpa_bwd_legacy(Tensor dO, Tensor q, Tensor k, Tensor v, Tensor o, Tensor stats, " |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Update the stale downstream operator call.
The declaration now registers cudnn::sdpa_bwd_legacy, but test/python/test_cudnn_sdpa_torch_ops.py still calls torch.ops.cudnn.sdpa_bwd. Update that test to use torch.ops.cudnn.sdpa_bwd_legacy, or preserve sdpa_bwd only if a schema-compatible alias is required.
Proposed test fix
- dq, dk, dv = torch.ops.cudnn.sdpa_bwd(
+ dq, dk, dv = torch.ops.cudnn.sdpa_bwd_legacy(🤖 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/experimental/ops/sdpa.py` around lines 599 - 602, Update the
downstream calls in test_cudnn_sdpa_torch_ops.py to use
torch.ops.cudnn.sdpa_bwd_legacy, matching the registered operator declaration in
the sdpa operation schema. Retain the sdpa_bwd name only if an explicitly
schema-compatible alias is required.
b9e8330 to
98331fe
Compare
|
All CodeRabbit findings triaged and addressed in Fixed — correctness:
Fixed — hygiene/tests:
Skipped, with reasons:
|
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/fe-oss-apis/sdpa-torch-ops.md`:
- Around line 21-35: Add a lazy public export for the SDPA torch wrapper in the
cudnn package initializer, ensuring torch and cutlass are imported only when the
exported API is accessed. Update the documentation example to import and invoke
that public entry point instead of importing cudnn.sdpa.fwd.torch_op directly,
while preserving the existing operator usage.
In `@python/cudnn/sdpa/fwd/torch_op.py`:
- Around line 662-682: Update the backward graph cache key in the sdpa_bwd key
construction to include the rounded packed token totals for T_q and T_kv, using
the same 64-token bucket calculation as _build_bwd_graph. Preserve the existing
key fields and ensure both totals participate in cache identity before the plan
is reused.
🪄 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: 570da5ff-3749-4643-9dc8-0545fcf56e6f
📒 Files selected for processing (3)
docs/fe-oss-apis/sdpa-torch-ops.mdpython/cudnn/sdpa/fwd/torch_op.pytest/python/test_cudnn_sdpa_torch_ops.py
|
Addendum — two more contract fixes landed in the same commit after re-running the upstream PyTorch suites against these ops (they were previously masked by the experimental dense module, which the provider no longer uses):
Upstream verification on the final tip: |
98331fe to
1877d67
Compare
|
|
||
|
|
||
| @torch.library.register_fake("cudnn::sdpa_bwd") | ||
| @torch.library.register_fake("cudnn::sdpa_bwd_legacy") |
There was a problem hiding this comment.
Can you fix the name about legacy ?
There was a problem hiding this comment.
Done in be0cf75 — the implementation identifiers now match the registered name: _sdpa_bwd_legacy_impl / _sdpa_bwd_legacy_fake. (The whole experimental module is slated for removal once dense backward lands in cudnn::sdpa_bwd.)
| if handles is None: | ||
| handles = _tls.handles = {} | ||
| if device not in handles: | ||
| handles[device] = cudnn.create_handle() |
There was a problem hiding this comment.
May be we should guard it with the device contenxt. It silently creates on current GPU
There was a problem hiding this comment.
Done in be0cf75 — handle creation is pinned with with torch.cuda.device(device): so a tensor on cuda:1 never gets a handle silently created against the caller's current device.
| # cuDNN ragged offsets are ELEMENT offsets per tensor, so each scales | ||
| # its token prefix sums by that tensor's OWN token stride. Small | ||
| # on-stream int ops — CUDA-graph-capture safe. | ||
| variant[int(_UIDs.RAGGED_Q)] = _int32_col(cu_seqlens_q * q.stride(0)) |
There was a problem hiding this comment.
We should confirm cu_seqlen_q is on device ? And _int64 as pointer ?
There was a problem hiding this comment.
Both done in be0cf75: _check_cu_seqlens validates that cu_seqlens_q/kv live on q's device before they are bound (a CPU tensor would hand cuDNN a host address), and the offsets are int64 pointers now (INT64 graph declarations + _int64_col binding).
| S_q, S_kv = max_seqlen_q, max_seqlen_kv | ||
| # cuDNN ragged offsets are int32 ELEMENT offsets: the largest offset | ||
| # (total tokens x token stride) must fit. | ||
| for name, t, total in (("q", q, T_q), ("k", k, T_kv), ("v", v, T_kv)): |
There was a problem hiding this comment.
We fixed it for int64? ?
There was a problem hiding this comment.
Done in be0cf75 — ragged offsets are now int64 end to end: the graph declares all 12 ragged-offset tensors as INT64 and the variant binds _int64_col columns (matching the test_mhas harness convention). The int32 overflow guards are gone — int64 element offsets cannot realistically overflow. Verified on both engine routes (FROST + backend), 19/19.
Family-local torch contract for the features torch.nn.functional.scaled_dot_product_attention cannot express: attention sinks, sliding window, bottom-right causal, padded batches, and THD/varlen packing (FA-style (T,H,D) + cu_seqlens). The ops build pygraph sdpa/sdpa_backward nodes; the Router picks the serving plan (FROST OSS kernels or backend engines) per config. Contract highlights: - register_fake meta kernels mirror the real kernels' output strides; torch.library.opcheck passes on both paths, including dynamic-shape AOT dispatch (torch.compile contract), and is locked in by a test. - sdpa_fwd is differentiable on the varlen path via register_autograd; the glue converts packed TH1 stats to the padded LSE layout device-side (no host reads, capture/tracing-safe). Dense and sink backward raise NotImplementedError until their engine contracts land. - Thread-safe: thread-local cuDNN handles (a handle must not be used from two threads), serialized graph builds, bounded (FIFO) graph cache. - Validation: one io dtype per call, k/o/grad_out shape checks, int32 ragged-offset overflow guards, inert-flag rejection (causal_bottom_right without an active band), clone() not contiguous() for base-pointer realignment (contiguous() cannot fix a misaligned base). cudnn::sdpa_fwd / cudnn::sdpa_bwd are the canonical names; the experimental dense module's backward is renamed cudnn::sdpa_bwd_legacy so both modules coexist in one process until it is removed. Tests (14, L0): sinks/window/bottom-right/padded dense with LSE value checks against an fp32 reference; THD fwd/bwd incl. GQA, kv-interleaved views, end-to-end autograd; opcheck. Docs: docs/fe-oss-apis/sdpa-torch-ops.md. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1877d67 to
be0cf75
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/cudnn/__init__.py (1)
328-337: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winRegister
jaxin_LAZY_OPTIONAL_IMPORTS.The dedicated
jaxbranch bypasses_load_optional_symbol, which adds_OPTIONAL_DEPENDENCY_INSTALL_HINTto import failures. Add"jax": (".jax", None)to_LAZY_OPTIONAL_IMPORTSand remove the dedicated branch.🤖 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/__init__.py` around lines 328 - 337, Update `_LAZY_OPTIONAL_IMPORTS` to include `"jax": (".jax", None)`, then remove the dedicated `if name == "jax"` branch so `__getattr__` routes JAX through `_load_optional_symbol` and preserves the existing lazy import behavior and installation hint handling.Sources: Coding guidelines, Learnings
🧹 Nitpick comments (3)
python/cudnn/sdpa/fwd/torch_op.py (2)
648-658: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winBackward accepts an inert
causal_bottom_right.
_sdpa_fwd_implrejectscausal_bottom_right=Truewhen neitheris_causalnorwindow_left >= 0is set (line 333)._sdpa_bwd_implaccepts it.alignmentthen re-anchors no band, so the flag has no effect and the two ops disagree on the same argument set. Repeat the forward check here so a mismatched backward call fails instead of computing a different mask than intended.🤖 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/torch_op.py` around lines 648 - 658, The _sdpa_bwd_impl validation must reject causal_bottom_right=True when neither is_causal nor window_left >= 0 is enabled, matching the existing _sdpa_fwd_impl check. Add this guard alongside the other backward argument validations so unsupported alignment does not proceed with an unanchored mask.
364-377: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winThe dense path skips the alignment and dense-innermost guard.
_normalize_thdenforcesstride(-1) == 1and a 16B-aligned base pointer for THD inputs. The dense branch passesq,k, andvthrough unchecked. A sliced dense view with an odd element storage offset then reaches the kernels with a misaligned address, which is the exact failure_normalize_thddocuments at lines 163-167.Apply an alignment-only normalization here, or validate and raise.
♻️ Proposed guard
B, H_q, S_q, D_qk = q.shape _, H_v, S_kv, D_v = v.shape H_k = k.shape[1] + for name, t in (("q", q), ("k", k), ("v", v)): + if t.data_ptr() % 16: + raise ValueError(f"{name} base pointer must be 16B-aligned for the cuDNN descriptors; got offset {t.storage_offset()}")🤖 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/torch_op.py` around lines 364 - 377, Update the dense-path handling around q_stride, k_stride, and v_stride to enforce dense-innermost layout and 16-byte base-pointer alignment for q, k, and v, matching the guarantees provided by _normalize_thd. Normalize eligible views or raise a clear validation error before kernel dispatch when alignment or stride requirements are not met.python/cudnn/__init__.py (1)
251-251: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExpose only explicit GEMM symbols through
cudnn.Remove the
attr_name=Noneentries forgrouped_gemmanddiscrete_grouped_gemm. Returning the package module exposes internal scheduler and utility submodules throughcudnn;__all__does not restrict direct attribute access. Register only the public classes and wrappers.🤖 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/__init__.py` at line 251, Update the `cudnn` registration entries for `grouped_gemm` and `discrete_grouped_gemm` to specify only their explicit public GEMM classes or wrappers, removing any `attr_name=None` values. Preserve the existing module paths while preventing direct exposure of internal scheduler and utility submodules through `cudnn`.Source: Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudnn/sdpa/fwd/torch_op.py`:
- Around line 774-780: Update the backward ragged-offset construction around the
_UIDs.RAGGED_Q through _UIDs.RAGGED_DV entries by converting cu_seqlens_q and
cu_seqlens_kv to torch.int64 once before these multiplications. Reuse the
widened prefix-sum tensors for all seven entries, preserving the existing
_int64_col conversions and stride/shape calculations.
---
Outside diff comments:
In `@python/cudnn/__init__.py`:
- Around line 328-337: Update `_LAZY_OPTIONAL_IMPORTS` to include `"jax":
(".jax", None)`, then remove the dedicated `if name == "jax"` branch so
`__getattr__` routes JAX through `_load_optional_symbol` and preserves the
existing lazy import behavior and installation hint handling.
---
Nitpick comments:
In `@python/cudnn/__init__.py`:
- Line 251: Update the `cudnn` registration entries for `grouped_gemm` and
`discrete_grouped_gemm` to specify only their explicit public GEMM classes or
wrappers, removing any `attr_name=None` values. Preserve the existing module
paths while preventing direct exposure of internal scheduler and utility
submodules through `cudnn`.
In `@python/cudnn/sdpa/fwd/torch_op.py`:
- Around line 648-658: The _sdpa_bwd_impl validation must reject
causal_bottom_right=True when neither is_causal nor window_left >= 0 is enabled,
matching the existing _sdpa_fwd_impl check. Add this guard alongside the other
backward argument validations so unsupported alignment does not proceed with an
unanchored mask.
- Around line 364-377: Update the dense-path handling around q_stride, k_stride,
and v_stride to enforce dense-innermost layout and 16-byte base-pointer
alignment for q, k, and v, matching the guarantees provided by _normalize_thd.
Normalize eligible views or raise a clear validation error before kernel
dispatch when alignment or stride requirements are not met.
🪄 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: dbe6014a-d59d-4b4d-b977-c37e5415e439
📒 Files selected for processing (4)
docs/fe-oss-apis/sdpa-torch-ops.mdpython/cudnn/__init__.pypython/cudnn/experimental/ops/sdpa.pypython/cudnn/sdpa/fwd/torch_op.py
🚧 Files skipped from review as they are similar to previous changes (2)
- docs/fe-oss-apis/sdpa-torch-ops.md
- python/cudnn/experimental/ops/sdpa.py
| int(_UIDs.RAGGED_Q): _int64_col(cu_seqlens_q * q.stride(0)), | ||
| int(_UIDs.RAGGED_KV): _int64_col(cu_seqlens_kv * k.stride(0)), | ||
| int(_UIDs.RAGGED_V): _int64_col(cu_seqlens_kv * v.stride(0)), | ||
| int(_UIDs.RAGGED_O): _int64_col(cu_seqlens_q * o.stride(0)), | ||
| int(_UIDs.RAGGED_DQ): _int64_col(cu_seqlens_q * (H_q * D_qk)), | ||
| int(_UIDs.RAGGED_DK): _int64_col(cu_seqlens_kv * (H_k * D_qk)), | ||
| int(_UIDs.RAGGED_DV): _int64_col(cu_seqlens_kv * (H_v * D_v)), |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Same int32 ragged-offset overflow on the backward path.
All seven ragged-offset entries multiply an int32 cu_seqlens_* tensor by a stride before _int64_col widens the result. Cast the two prefix-sum tensors with .to(torch.int64) once, then multiply.
[skip_comment]
⛔ Skipped due to learnings
Learnt from: YangXu1990uiuc
Repo: NVIDIA/cudnn-frontend PR: 509
File: python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py:1218-1224
Timestamp: 2026-08-09T19:32:05.011Z
Learning: In `python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py`, `cutlass.Vector` vectorizes `prims.mul_packed_f32x2` when `fmul2` receives the eight-element `cutlass.Array` slice `o_regs[o_off:8]`. Do not assume that the `Tuple[Float32, Float32]` annotation limits this operation to one pair.
Learnt from: YangXu1990uiuc
Repo: NVIDIA/cudnn-frontend PR: 509
File: python/cudnn/sdpa/fwd/api_dsl.py:2018-2024
Timestamp: 2026-08-09T19:32:34.987Z
Learning: In `python/cudnn/sdpa/fwd/api_dsl.py`, FP8 Amax buffers are written in place by kernel atomics. The shared `SdpaFwdDsl._amax_slot()` helper must use `tensor.view(-1)[:1]`, not `reshape()`, so non-contiguous caller buffers raise `ValueError` instead of creating a detached copy. The SM100 and SM120 per-tensor FP8 paths use this helper.
Learnt from: CR
Repo: NVIDIA/cudnn-frontend PR: 0
File: python/cudnn/AGENTS.md:0-0
Timestamp: 2026-07-31T21:15:22.165Z
Learning: Applies to python/cudnn/**/*.py : For FP4x2 packing, use `_tensor_shape` and `_tensor_stride`; when `interpret_uint8_as_fp4x2` is enabled, the innermost dimension is doubled.
Learnt from: YangXu1990uiuc
Repo: NVIDIA/cudnn-frontend PR: 509
File: python/cudnn/sdpa/fwd/api_dsl.py:1985-1991
Timestamp: 2026-08-09T19:32:36.684Z
Learning: In `python/cudnn/sdpa/fwd/engines.py`, `lower_dsl_prefill` clears `seq_q_lens_present` for FP8/MXFP8 before adapter-level checks. Therefore, FP8 handling for dense per-batch Q lengths must validate at execute time: allow `seq_len_q == S_q`, and reject shorter Q lengths to prevent writes of O and finite LSE values past the valid query range.
🤖 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/torch_op.py` around lines 774 - 780, Update the
backward ragged-offset construction around the _UIDs.RAGGED_Q through
_UIDs.RAGGED_DV entries by converting cu_seqlens_q and cu_seqlens_kv to
torch.int64 once before these multiplications. Reuse the widened prefix-sum
tensors for all seven entries, preserving the existing _int64_col conversions
and stride/shape calculations.
2 participants
What
PyTorch custom ops —
torch.ops.cudnn.sdpa_fwd/sdpa_bwd— exposing the cuDNN SDPA feature surface thattorch.nn.functional.scaled_dot_product_attention's aten contract cannot express:window_left, cuDNN visible-tokens convention)seq_len_q/seq_len_kv)(T, H, D)+cu_seqlens; non-contiguous K/V views of a fused KV projection are declared with their true strides)The ops build pygraph
sdpa/sdpa_backwardnodes; the engine Router picks the serving plan (FROST OSS kernels or cuDNN-backend engines) per configuration. One new file + tests + docs; no existing behavior changes (the experimental dense module's backward is renamedcudnn::sdpa_bwd_legacyso both modules can coexist in one process until it is removed).Contract
torch.compileregister_fakemeta kernels mirror the real output strides;torch.library.opcheckpasses on dense and THD paths including dynamic-shape AOT dispatch, locked in by a testsdpa_fwdis differentiable on the varlen path (register_autograd); the glue converts packed TH1 stats to the padded LSE layout device-side (no host reads — tracing/capture safe). Dense + sink backward raiseNotImplementedErroruntil their engine contracts landh_k != h_vGQA head splits validated (independent K/V head counts); int32 ragged-offset overflow guards; inertcausal_bottom_rightrejected; base-pointer realignment usesclone()(contiguous()cannot fix a misaligned base)(B, H, max_seqlen_q, 1)fp32 — backend restriction (bprop THD rejects ragged LSE on SM8X/SM12X)Tests (19, all L0 —
test/python/test_cudnn_sdpa_torch_ops.py)Dense: sinks (±causal) with LSE value checks,
h_k != h_vGQA, query-layout adoption (4 permutations) vs an fp32 reference, sliding windows, bottom-right cross-seqlen, sinks+window, padded per-batch lengths. THD: fwd with LSE value checks, fwd+bwd (MHA + GQA), end-to-end autograd, kv-interleaved(T,2,H,D)views. Contract: opcheck (dense + THD/dynamic). All green on SM100 (cc 10.0) + previously SM120, both engine routes.Docs:
docs/fe-oss-apis/sdpa-torch-ops.md.Consumers
torch.nn.attention.varlenandF.sdpaintegration builds on these ops in #554 (thecudnn.torchprovider); TE/vLLM/Megatron-style frameworks can call them directly.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Tests