Support mixed-form sequence lengths in SDPA forward (cuDNN 9.25+) - #430
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:
📝 WalkthroughWalkthroughSDPA validation now supports independently selected explicit or cumulative sequence-length forms for Q and KV. Test configuration, tensor allocation, graph construction, version gating, documentation, samples, and parameterized mixed-form coverage are updated accordingly. ChangesMixed sequence-length support
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant ExecConfig
participant allocate_tensors
participant create_forward_graph
participant SDPA_attributes
ExecConfig->>allocate_tensors: select Q and KV sequence-length forms
allocate_tensors->>create_forward_graph: provide explicit or cumulative tensors
create_forward_graph->>SDPA_attributes: validate mixed-form inputs
SDPA_attributes-->>create_forward_graph: accept or reject configuration
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
egilliam-nv
added
cat-enhancements
mod-frontend
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-430-da67de4 |
da67de4 to
14e9175
Compare
egilliam-nv
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
test/python/test_mhas_v2.py (1)
192-192: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winFormat these calls with Black.
These three changed calls exceed the configured 160-character limit.
test/python/test_mhas_v2.py#L192-L192: wrap theRandomHiddenDimSizearguments.test/python/test_mhas_v2.py#L481-L481: wrap theRandomHiddenDimSizearguments.test/python/test_mhas_v2.py#L654-L654: wrap theRandomHiddenDimSizearguments.As per coding guidelines, “Format Python code and notebooks with Black using a line length of 160.”
🤖 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_mhas_v2.py` at line 192, Format the RandomHiddenDimSize calls with Black’s configured 160-character line length, wrapping their arguments at test/python/test_mhas_v2.py lines 192-192, 481-481, and 654-654 without changing behavior.Source: Coding guidelines
🧹 Nitpick comments (1)
test/python/test_mhas_v2.py (1)
1049-1079: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExercise the UNIFIED path explicitly.
The feature under test is gated in the UNIFIED surface, but this config requests
AUTO; the test can pass without covering that path. SetimplementationtoUNIFIED, or verify that AUTO is guaranteed to resolve to it here.Proposed fix
- implementation=cudnn.attention_implementation.AUTO, + implementation=cudnn.attention_implementation.UNIFIED,🤖 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_mhas_v2.py` around lines 1049 - 1079, Update the implementation setting in the SDPATestConfig and its nested ExecConfig within the test setup to explicitly use cudnn.attention_implementation.UNIFIED instead of AUTO, ensuring this test exercises the gated UNIFIED path.
🤖 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.
Outside diff comments:
In `@test/python/test_mhas_v2.py`:
- Line 192: Format the RandomHiddenDimSize calls with Black’s configured
160-character line length, wrapping their arguments at
test/python/test_mhas_v2.py lines 192-192, 481-481, and 654-654 without changing
behavior.
---
Nitpick comments:
In `@test/python/test_mhas_v2.py`:
- Around line 1049-1079: Update the implementation setting in the SDPATestConfig
and its nested ExecConfig within the test setup to explicitly use
cudnn.attention_implementation.UNIFIED instead of AUTO, ensuring this test
exercises the gated UNIFIED path.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 47ca0cff-b6df-42cc-b8db-1a49a661e6d5
📒 Files selected for processing (5)
include/cudnn_frontend/node/diagonal_band_mask.hinclude/cudnn_frontend/node/sdpa_support_surface.htest/python/sdpa/fp16.pytest/python/sdpa/random_config.pytest/python/test_mhas_v2.py
🚧 Files skipped from review as they are similar to previous changes (3)
- include/cudnn_frontend/node/diagonal_band_mask.h
- include/cudnn_frontend/node/sdpa_support_surface.h
- test/python/sdpa/random_config.py
14e9175 to
49d108f
Compare
|
Thanks @coderabbitai.
|
|
Tip For best results, initiate chat on the files or code changes.
🐇 ✅ |
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-430-49d108f |
|
Good catch on the stale sample doc — fixed in 62424a9.
I also swept the same now-stale wording out of the parallel parameter docs while here: |
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-430-62424a9 |
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 `@docs/operations/Attention.md`:
- Line 327: Update the duplicated SDPA sequence-length summaries consistently:
in docs/operations/Attention.md lines 327-327, allow the padding-mask summary to
accept either per-batch or cumulative lengths; in docs/operations/Attention.md
lines 914-914, state that cuDNN 9.24+ supports cumulative forms while cuDNN
9.25+ is required only for mixed forms; and in python/pygraph/sdpa.cpp lines
1307-1308, document opposite-side length pairing, the requirement for
use_padding_mask=True, and the mixed-form version requirement.
🪄 Autofix (Beta)
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: 8c6b0cbd-7ff0-49d2-8693-45ed1f764d19
📒 Files selected for processing (3)
docs/operations/Attention.mdpython/pygraph/sdpa.cppsamples/cpp/sdpa/fp16_fwd_with_cu_seq_len.cpp
|
Thanks — addressed in c282492.
One deliberate divergence from the suggestion on the FP8 summary: I did not state 9.24+ for cumulative forms there. The FP8 unified engine itself requires cuDNN 9.25+, so cu_seq_len on the FP8 path needs 9.25+ in any form; the 9.24-base / 9.25-mixed split applies only to fp16/bf16. The previous doc commit had incorrectly carried 9.24 into the FP8 docs — that's now corrected to 9.25+. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/pygraph/sdpa.cpp`:
- Around line 1141-1142: Wrap the cu_seq_len_q and cu_seq_len_kv documentation
in the relevant API comments, including the corresponding lines around their
later occurrence, so no C++ source line exceeds 120 columns. Preserve the
rendered wording and formatting semantics while splitting each description
across adjacent documentation lines.
🪄 Autofix (Beta)
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: c30765a8-230c-420c-9960-9da48549cec7
📒 Files selected for processing (2)
docs/operations/Attention.mdpython/pygraph/sdpa.cpp
🚧 Files skipped from review as they are similar to previous changes (1)
- docs/operations/Attention.md
| cu_seq_len_q (Optional[cudnn_tensor]): Cumulative sequence length of the query, shape (b+1, 1, 1, 1) or 1-D (b+1,) (promoted automatically), int32 or int64. Mutually exclusive with seq_len_q (a side cannot use both forms); must be paired with a KV-side length tensor (seq_len_kv or cu_seq_len_kv) and requires use_padding_mask=True. The two sides may use different forms (e.g. cu_seq_len_q with seq_len_kv), which requires cuDNN 9.25.0 or newer. Supplying cu_seq_len_q requires cuDNN 9.24.0 or newer and the UNIFIED implementation. | ||
| cu_seq_len_kv (Optional[cudnn_tensor]): Cumulative sequence length of the key, shape (b+1, 1, 1, 1) or 1-D (b+1,) (promoted automatically), int32 or int64. Mutually exclusive with seq_len_kv (a side cannot use both forms); must be paired with a Q-side length tensor (seq_len_q or cu_seq_len_q) and requires use_padding_mask=True. The two sides may use different forms (e.g. seq_len_q with cu_seq_len_kv), which requires cuDNN 9.25.0 or newer. Supplying cu_seq_len_kv requires cuDNN 9.24.0 or newer and the UNIFIED implementation. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Wrap the new documentation lines to the repository width.
The added C++ source lines substantially exceed the required 120-column limit. Wrap the text across adjacent documentation lines while preserving the rendered wording.
As per coding guidelines, C++ files must use Google-based formatting and a 120-column width.
Also applies to: 1307-1308
🤖 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/pygraph/sdpa.cpp` around lines 1141 - 1142, Wrap the cu_seq_len_q and
cu_seq_len_kv documentation in the relevant API comments, including the
corresponding lines around their later occurrence, so no C++ source line exceeds
120 columns. Preserve the rendered wording and formatting semantics while
splitting each description across adjacent documentation lines.
Source: Coding guidelines
Allow the Q and KV sides of SDPA forward to independently choose their sequence length representation: per-batch (seq_len_*) or cumulative (cu_seq_len_*). Previously the two forms were mutually exclusive across the whole operation, forcing paged-attention integrations -- whose natural state is a cumulative Q prefix sum alongside per-batch KV lengths -- to materialize a KV-side prefix sum purely to satisfy the pairing rule. - sdpa_support_surface: validation is now per-side (exactly one form per side, sides provided together); the paged requirement accepts either form per side; mixed forms are gated on cuDNN 9.25+ in the UNIFIED surface. Hoist a has_input helper for the presence checks. - diagonal_band_mask: relax the node's metadata seqlen check to per-side exclusivity, matching the SDPA rules. - tests: per-side form control in the SDPA test harness (cu_seq_len_sides) and a deterministic mixed-form test with non-uniform lengths, including bottom-right causal cases that guard the backend's band-mask alignment derivation. Validated on H100 against a cuDNN 9.27 dev build: mixed forms in both directions match the both-per-batch reference for non-paged, ragged, and paged (ragged Q + page tables, cu_seq_len_q + seq_len_kv) graphs; the deterministic tests skip below the version gate. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…forms The prior wording described the pre-relaxation constraints -- "cu_seq_len_q and cu_seq_len_kv must both be set or both unset" and "cu_seq_len_* are mutually exclusive with seq_len_q / seq_len_kv" -- which this change makes stale and misleading. Update them to the per-side rules: each side independently picks one form (per-batch or cumulative), the two sides must both be present but may use different forms, and mixed forms require cuDNN 9.25+ (plain cu_seq_len on both sides remains 9.24+). Touches the fp16_fwd_with_cu_seq_len sample header, the C++/Python parameter docs in docs/operations/Attention.md, and the sdpa pybind docstrings. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Follow-up to review: several duplicated sequence-length summaries still described the pre-relaxation contract, and the previous doc commit set the wrong version floor on the FP8 path. - Attention.md: the regular-sdpa use_padding_mask summary still said "Must also provide seq_len_q and seq_len_kv"; now allows either per-batch or cumulative lengths. - python/pygraph/sdpa.cpp: enrich the regular sdpa cu_seq_len_q/kv docstrings (opposite-side pairing, use_padding_mask=True, mixed-form 9.25 requirement) to match the detailed docs. - Correct the FP8 cu_seq_len docs (Attention.md and the sdpa_fp8 docstrings): the FP8 path requires cuDNN 9.25+ for cumulative sequence lengths in any form (the FP8 unified engine itself is 9.25+), so the 9.24-base / 9.25-mixed split applies only to fp16/bf16. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Review follow-up: the enriched cu_seq_len_q/kv docstring lines in the sdpa and sdpa_fp8 pybind blocks were ~550 chars, well beyond the other parameter lines. Condense the wording (dropping redundant phrasings) while keeping the constraints: same-side exclusivity, opposite-side pairing, use_padding_mask=True, and the version floors (regular 9.24+, mixed 9.25+; fp8 9.25+). These live inside a single R"pbdoc(...)pbdoc" raw string where the block documents one parameter per line, so a source line break would insert a literal newline into the rendered help text; the lines are shortened rather than wrapped. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
c282492 to
a1772cb
Compare
|
Addressed in a1772cb — but shortened rather than wrapped, for a structural reason. These Instead I condensed the wording so the lines are no longer outliers (~550 → ~370 chars, in line with existing long entries like |
4 participants
Before submitting
pre-commit runand committed any formatting changes.Affected area
Summary
Allow the Q and KV sides of SDPA forward to independently choose their sequence-length representation: per-batch (
seq_len_*) or cumulative (cu_seq_len_*). Previously the two forms were mutually exclusive across the whole operation.sdpa_support_surface: validation is now per-side (exactly one form per side; both sides provided together); the paged requirement accepts either form per side; mixed forms are gated on cuDNN 9.25+ in the UNIFIED surface. Ahas_inputhelper is hoisted for the presence checks.diagonal_band_mask: the node's metadata seqlen check is relaxed to per-side exclusivity, matching the SDPA rules.cu_seq_len_sides) plus a deterministic mixed-form test (test_sdpa_mixed_seq_len_forms_L0) with non-uniform per-batch lengths, including bottom-right-causal cases that guard the backend's band-mask alignment derivation.Why
Paged-attention integrations naturally hold a cumulative Q prefix sum (
qo_indptr) alongside per-batch KV lengths. The old same-form-both-sides rule forced them to materialize a KV-side prefix sum purely to satisfy the pairing constraint. Permitting mixed forms removes that unnecessary conversion.Mixed forms require the per-side sequence-length support added to the cuDNN backend in 9.25, so they are gated on
min(compiled, runtime) >= 92500in the UNIFIED support surface and rejected with a clear error below that.Related issues
Backend support merged separately (cuDNN backend MRs, internal). Related to #366 (expose
cu_seq_len_q/kvon thesdpa_fp8Python binding).API and compatibility impact
No public API additions or removals. Purely relaxes an existing validation constraint on the
cu_seq_len_*/seq_len_*SDPA inputs; all previously-valid graphs remain valid. New capability (mixed forms) is version-gated to cuDNN 9.25+ and rejected with a clear error below that.Testing
On H100 (sm90) against a cuDNN 9.27 dev build with the backend support:
pytest test/python/test_mhas_v2.py -k mixed_seq_len_forms— 4/4 pass (both mixed directions x {top-left, bottom-right-causal}).pytest test/python/test_sdpa_edge_cases.py— 10/10 pass (no harness regression).pytest test/python/test_mhas_v2.py::test_sdpa_random_fwd_L0(slice) — 203 passed, 9 skipped.cu_seq_len_q+seq_len_kvand bottom-right causal, all match the both-per-batch reference.Below cuDNN 9.25 the mixed-form tests skip via the version gate.
Summary by CodeRabbit
cu_seq_len_*rules, mixed-form support, and version/implementation requirements.