Fix ragged SDPA backward workspace under-allocation for non-token-major stats layouts - #462
Conversation
max_total_seq_len_* switches softmax_sum / dQ_accum / dK_fullhead / dV_fullhead to a THD-ragged layout and sized them as max_total_seq_len * stride[2] * sizeof(float), which assumes the per-token footprint is stride[2]. That holds only for token-major tensors. For head-major stats ([h, t], so stride[2] == 1 -- FlashAttention's and PyTorch varlen's softmax_lse layout) the buffer is under-allocated by h_q and softmax_sum overruns into the backend engine's own workspace, which the frontend places immediately after the FE region. The first word there is the persistent scheduler's tile_id_counter; once it holds a float bit pattern the unsigned tile_id < num_tiles test fails in every CTA, so all CTAs retire without storing and dQ/dK/dV come back exactly zero with no error and a correct forward. Size from every axis instead. For token-major inputs the new expression is bit-identical to the old one, so allocations are unchanged for every layout that works today. Verified on sm90 and sm100 against cuDNN 9.26. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The ragged SDPA tests only ever generated token-major stats ([t, h, 1], sequence stride h_q), so nothing exercised a packed Stats tensor whose per-token footprint is not stride[2]. That blind spot is what let the max_total_seq_len workspace under-allocation fixed in the previous commit survive from 9.22 to 9.26. Adds ragged_stats_layout to ExecConfig, randomized in the ragged fwd/bwd tests, selecting between token-major and head-major ([h, t], sequence stride 1 -- the FlashAttention / PyTorch varlen softmax_lse layout). The stats allocation and the stats ragged offset now both derive from cfg.stride_stats instead of hardcoding the token-major assumption. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
📝 WalkthroughWalkthroughRagged SDPA configuration now supports token-major and head-major statistics layouts. Test allocations preserve configured strides and packed capacities. Flash attention workspace sizing now accounts for complete strided tensor footprints. ChangesRagged SDPA sizing and layouts
Estimated code review effort: 3 (Moderate) | ~20 minutes 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: 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 `@include/cudnn_frontend/node/scaled_dot_product_flash_attention.h`:
- Around line 1094-1105: The ragged_workspace_size helper currently sizes only
max_total_seq_len and can under-allocate when offsets contain inter-batch gaps.
Update the backward workspace sizing flow around ragged_workspace_size to use
the physical ragged span through the final batch boundary, or explicitly reject
unsupported partially packed offsets with documentation; preserve correct
stride-based sizing for all axes. Add a regression test covering inter-batch
padding and verifying the workspace does not overlap subsequent allocations.
In `@test/python/test_mhas_v2.py`:
- Line 368: The ragged_stats_layout randomization in the forward test
configuration (around line 368) does not provide meaningful coverage because the
forward test runs with is_infer=True (default), which means TensorUid.stats and
related output tensors are not allocated or consumed, so the stats_ragged_offset
input tensor is never actually used. Remove the RandomChoice ragged_stats_layout
configuration from the forward test and relocate it to the backward test
configuration (around line 469), where is_infer=False ensures both layout
options are properly exercised. Alternatively, enable with_score_max or
with_score_sum_exp in the forward test to activate the stats outputs and allow
the randomization to provide real coverage.
🪄 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: 8d86a621-cd9a-41a0-bc93-12213b4ea74d
📒 Files selected for processing (4)
include/cudnn_frontend/node/scaled_dot_product_flash_attention.htest/python/sdpa/fp16.pytest/python/sdpa/random_config.pytest/python/test_mhas_v2.py
| // Byte span of a workspace tensor packed (ragged) over the token axis: the sequence axis is | ||
| // bounded by max_total_seq_len instead of dim[2], every other axis contributes its own extent. | ||
| // Sizing from stride[2] alone under-allocates any layout whose per-token footprint is not | ||
| // stride[2] -- e.g. head-major [h, t] stats, where stride[2] == 1. | ||
| static int64_t | ||
| ragged_workspace_size(int64_t max_total_seq_len, | ||
| std::vector<int64_t> const& dim, | ||
| std::vector<int64_t> const& stride, | ||
| int64_t elem_size) { | ||
| return ((max_total_seq_len - 1) * stride[2] + (dim[1] - 1) * stride[1] + (dim[3] - 1) * stride[3] + 1) * | ||
| elem_size; | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="$(fd --type f 'scaled_dot_product_flash_attention.h' | head -n 1)"
test -n "$file"
ast-grep outline "$file" --items all --match 'CompositeSDPABackwardNode'
rg -n -C 8 'ragged_workspace_size|set_ragged_offset|max_total_seq_len' "$file"
rg -n -C 5 'partially.?packed|ragged.?offset|max_total_seq_len|packed_token_capacity' .Repository: NVIDIA/cudnn-frontend
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="include/cudnn_frontend/node/scaled_dot_product_flash_attention.h"
printf '%s\n' '--- helper and call sites ---'
sed -n '1088,1110p;1718,1734p;1924,1943p;2036,2056p;2080,2098p' "$file"
printf '%s\n' '--- ragged offset definitions and workspace construction ---'
rg -n -C 6 'get_ragged_offset|set_ragged_offset|ragged_offset_multiplier|workspace_size' include/cudnn_frontend | head -n 500
printf '%s\n' '--- offset construction in tests and examples ---'
rg -n -C 8 'cum_q|cum_kv|ragged_offset|set_ragged_offset' test python samples include | \
rg -v '(__pycache__|\.pyc)' | head -n 500Repository: NVIDIA/cudnn-frontend
Length of output: 50377
🌐 Web query:
cuDNN frontend ragged offset partially packed tensor max_total_seq_len semantics
💡 Result:
In the NVIDIA cuDNN frontend library, when using attention operations (such as SDPA) with packed or partially packed tensor layouts, specific semantics apply to ragged offsets and maximum sequence length parameters [1][2]. Ragged Offset Semantics: When employing a packed layout (such as THD) or a partially packed layout—where variable-length sequences are stored contiguously without padding—you must provide a ragged offset tensor to define the boundaries of individual sequences within the batch [2][3]. 1. Tensor Shape: The ragged offset tensor must have a shape of (B+1, 1, 1, 1) [1][2]. 2. Functionality: It is an integer tensor containing the cumulative token offsets (in terms of number of elements, not bytes) [1][2]. The value at index i represents the starting position of batch i, and the last value (at index B) specifies the offset of the past-the-end element of the entire ragged structure [1][2]. 3. Partially Packed Layouts: This layout allows tokens within each batch to be contiguous, even if the batches are not globally packed (e.g., gaps may exist between batches) [2][3]. The ragged offset tensor is still required to map these sequences correctly [1][2]. Max Total Seq Len Semantics: The parameter often used in conjunction with these layouts is related to workspace allocation and memory planning [1][2]. 1. Purpose: When using packed layouts with backward propagation (bprop), it is recommended to use an API (e.g., set_max_total_seq_len_q) to specify the maximum number of tokens [1][2]. 2. Default Behavior: If this parameter is not explicitly set, the library typically assumes the maximum sequence length is B x S (where B is the batch size and S is the maximum sequence length per batch), which may result in a larger-than-necessary workspace allocation [1][2]. These semantics ensure that the cuDNN kernels correctly identify the valid data segments within the memory buffers, facilitating efficient execution of attention operations on variable-length sequence data [2][3].
Citations:
- 1: https://docs.nvidia.com/deeplearning/cudnn/v1.12.1/operations/Attention.html
- 2: https://docs.nvidia.com/deeplearning/cudnn/frontend/latest/operations/Attention.html
- 3: https://docs.nvidia.com/deeplearning/cudnn/frontend/v1.23.0/operations/Attention.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="include/cudnn_frontend/node/scaled_dot_product_flash_attention.h"
printf '%s\n' '--- helper and call sites ---'
sed -n '1088,1110p;1718,1734p;1924,1943p;2036,2056p;2080,2098p' "$file"
printf '%s\n' '--- ragged offset definitions and workspace construction ---'
rg -n -C 6 'get_ragged_offset|set_ragged_offset|ragged_offset_multiplier|workspace_size' include/cudnn_frontend | head -n 500
printf '%s\n' '--- offset construction in tests and examples ---'
rg -n -C 8 'cum_q|cum_kv|ragged_offset|set_ragged_offset' test python samples include/ | \
rg -v '(__pycache__|\.pyc)' | head -n 500Repository: NVIDIA/cudnn-frontend
Length of output: 50377
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- exact ragged-offset and workspace tensor handling ---'
rg -n -C 12 'collect_tensors_in_workspace_node|get_fe_workspace_size_node|softmax_sum_size|dV_fullhead_size|dK_fullhead_size|dQ_accum_size' \
include/cudnn_frontend/node/scaled_dot_product_flash_attention.h
printf '%s\n' '--- documentation and tests for partially packed layouts ---'
rg -n -i -C 10 'partially[ -]?packed|partially.?pack|ragged.?offset|cum(ulative)?[ _-]?seq|offset.*gap|gap.*offset' \
docs test python samples include | head -n 800
printf '%s\n' '--- current branch metadata ---'
git remote -v
git rev-parse --short HEADRepository: NVIDIA/cudnn-frontend
Length of output: 50379
🌐 Web query:
site:docs.nvidia.com/deeplearning/cudnn/frontend/latest/operations/Attention.html "partially packed" "ragged offset"
💡 Result:
The term partially packed in the context of ragged offsets refers to a memory layout optimization for handling variable-length sequences in deep learning frameworks, specifically within NVIDIA cuDNN [1]. In this layout, sequences within a batch are stored contiguously but are not required to be globally packed across the entire memory buffer [1]. This contrasts with standard packed (or THD) layouts, where all valid tokens are concatenated into a single contiguous block ($T = \sum(\text{seq_len})$) to eliminate padding [1]. Key aspects of the partially packed layout include: 1. Ragged Offsets: The layout uses a ragged offset tensor—an array of size (B+1)—to define the start and end positions of each sequence within the buffer [1]. By tracking these offsets, the system can identify and process variable-length data without needing global contiguity [1]. 2. Contiguity Requirements: While individual batches can be partially packed, tokens within each batch must remain contiguous to be valid [1]. Interleaved or non-contiguous tokens within a single batch are not supported under this representation [1]. 3. Memory Efficiency: This approach provides a balance between the efficiency of packed layouts (which avoid padding) and the flexibility needed for specific memory management patterns, such as those used in attention mechanisms or paged attention systems where frameworks may prefer to track necessary block indices via ragged offsets [1]. Essentially, the ragged offset tensor acts as a map, allowing the compute engine to navigate the memory buffer and correctly identify the bounds of each variable-length sequence even when the sequences are not stored in one continuous global sequence [1].
Citations:
🌐 Web query:
site:docs.nvidia.com/deeplearning/cudnn "CUDNN_ATTR_TENSOR_RAGGED_OFFSET" "partially packed"
💡 Result:
In the NVIDIA cuDNN library, CUDNN_ATTR_TENSOR_RAGGED_OFFSET_DESC is an attribute used in the cuDNN backend API to define the descriptor for a ragged offset tensor [1][2]. This attribute is explicitly associated with ragged tensors, which are tensors containing nested variable-length lists, commonly utilized in operations like attention [3][4]. There is no concept of a "partially packed" ragged offset; instead, "partially-packed" is a distinct technical term in cuDNN referring to a tensor's memory layout [5][6]. 1. CUDNN_ATTR_TENSOR_RAGGED_OFFSET_DESC: This attribute is set on a tensor descriptor to provide an integer tensor (typically of size B+1) that specifies the offsets or starting positions for variable-length segments within a ragged tensor [3][7]. It is required when handling packed layouts (such as THD) where multiple batches are concatenated together, allowing the library to identify where each batch begins and ends [3][4]. 2. Partially-Packed Tensors: This terminology refers to the packing status of a tensor's dimensions [5][8]. A tensor with a given dimension format (e.g., WXYZ) is "XYZ-packed" if specific striding conditions are met for the dimensions referenced in the suffix [5][9]. This definition is independent of ragged tensor functionality [5][6]. In summary, CUDNN_ATTR_TENSOR_RAGGED_OFFSET_DESC is used to manage the boundaries of ragged data, while "partially-packed" describes specific constraints on how a tensor's strides map its dimensions in memory [5][3].
Citations:
- 1: https://docs.nvidia.com/deeplearning/cudnn/backend/v9.10.2/api/cudnn-graph-library.html
- 2: https://docs.nvidia.com/deeplearning/cudnn/backend/v9.23.2/api/cudnn-graph-library.html
- 3: https://docs.nvidia.com/deeplearning/cudnn/v1.14.1/operations/Attention.html
- 4: https://docs.nvidia.com/deeplearning/cudnn/v1.11.0/operations/Attention.html
- 5: https://docs.nvidia.com/deeplearning/cudnn/latest/developer/core-concepts.html
- 6: https://docs.nvidia.com/deeplearning/cudnn/archives/cudnn-897/developer-guide/index.html
- 7: https://docs.nvidia.com/deeplearning/cudnn/v1.9.0/operations/Attention.html
- 8: https://docs.nvidia.com/deeplearning/cudnn/backend/v9.5.1/developer/core-concepts.html
- 9: https://docs.nvidia.com/deeplearning/cudnn/frontend/v1.19.1/developer/core-concepts.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- exact ragged-offset and workspace tensor handling ---'
rg -n -C 12 'collect_tensors_in_workspace_node|get_fe_workspace_size_node|softmax_sum_size|dV_fullhead_size|dK_fullhead_size|dQ_accum_size' \
include/cudnn_frontend/node/scaled_dot_product_flash_attention.h
printf '%s\n' '--- documentation and tests for partially packed layouts ---'
rg -n -i -C 10 'partially[ -]?packed|partially.?pack|ragged.?offset|cum(ulative)?[ _-]?seq|offset.*gap|gap.*offset' \
docs test python samples include/ | head -n 800
printf '%s\n' '--- current branch metadata ---'
git remote -v
git rev-parse --short HEADRepository: NVIDIA/cudnn-frontend
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- max_total_seq_len documentation and validation ---'
rg -n -i -C 12 'max_total_seq_len(_q|_kv)?|maximum total sequence|total sequence length' \
docs include python test | head -n 1000
printf '%s\n' '--- ragged workspace output layouts and stride assumptions ---'
rg -n -C 10 'set_dim\(\{b, h_q|set_stride\(.*outputs\[output_names::d[QKV]\]|dQ_accum|dK_fullhead|dV_fullhead' \
include/cudnn_frontend/node/scaled_dot_product_flash_attention.h | head -n 700
printf '%s\n' '--- exact documentation block ---'
sed -n '61,110p' docs/operations/Attention.mdRepository: NVIDIA/cudnn-frontend
Length of output: 50379
🌐 Web query:
site:docs.nvidia.com/deeplearning/cudnn/frontend/latest/operations/Attention.html "max_total_seq_len"
💡 Result:
In the context of machine learning and high-performance computing, particularly within NVIDIA's cuDNN framework, max_total_seq_len (often split into max_total_seq_len_q and max_total_seq_len_kv) is a parameter used to define the maximum total number of tokens across a batch when processing ragged tensors [1]. Ragged tensors occur when sequences within a single batch have varying lengths. Because memory needs to be pre-allocated for operations like Scaled Dot-Product Attention (SDPA), this parameter informs the system of the maximum aggregate sequence length (the sum of the lengths of all sequences in the batch) [1]. Providing this value allows the framework to correctly allocate the workspace memory required for backward passes and other operations, preventing overflow or inefficient memory usage [1]. If not explicitly provided, it often defaults to the batch size multiplied by the maximum sequence length per item (B x S) [1].
Citations:
🌐 Web query:
site:docs.nvidia.com/deeplearning/cudnn "max_total_seq_len_q" packed layout
💡 Result:
In the context of NVIDIA cuDNN and its attention operations, max_total_seq_len_q is an optional configuration parameter used when working with ragged (packed) tensors [1][2][3]. When processing variable-length sequences packed together into a single continuous memory block (often referred to as the THD layout, where T = sum(sequence_lengths)) to avoid padding, this parameter specifies the maximum total number of tokens for all queries in a batch [1][3]. Key details regarding max_total_seq_len_q: Purpose: It is primarily used for efficient workspace memory allocation [1][2]. By providing this value, the library can allocate the necessary memory buffer for intermediate computations without needing to over-allocate based on a worst-case scenario [1][4]. Default Behavior: If max_total_seq_len_q is not explicitly provided, it defaults to B * S_q (where B is the batch size and S_q is the maximum sequence length per query), which may result in unnecessary memory consumption [1][3]. Usage: It should be set to the maximum sum of sequence lengths expected for the queries in your input data [1][2]. A corresponding parameter, max_total_seq_len_kv, is also typically used for keys and values [1][3]. In the THD (Total-length, Head-count, Dimension) layout, sequences are concatenated, and ragged offsets (or cumulative sequence length tensors) are used to track the boundaries of individual sequences within that packed structure [1][5][6]. This approach is widely used in high-performance Transformer implementations to optimize memory bandwidth and compute utilization by eliminating the overhead of padding tokens [1][2].
Citations:
- 1: https://docs.nvidia.com/deeplearning/cudnn/frontend/latest/operations/Attention.html
- 2: https://docs.nvidia.com/deeplearning/cudnn/v1.25.0/operations/Attention.html
- 3: https://docs.nvidia.com/deeplearning/cudnn/latest/operations/Attention.html
- 4: https://docs.nvidia.com/deeplearning/cudnn/v1.19.1/operations/Attention.html
- 5: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/nsa.html
- 6: https://docs.nvidia.com/deeplearning/cudnn/frontend/latest/fe-oss-apis/nsa.html
Size workspace for the physical ragged span.
max_total_seq_len is the sum of valid tokens, but ragged offsets can include inter-batch gaps. For offsets [0, 4*H*D, 7*H*D] and five valid tokens, this helper sizes only the five-token span even though the final batch ends at ten tokens. The resulting write can overlap the next frontend workspace allocation.
Use a physical-capacity bound that includes gaps, or reject and document partially packed offsets for this backward path. Add a regression test with inter-batch padding.
🤖 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 `@include/cudnn_frontend/node/scaled_dot_product_flash_attention.h` around
lines 1094 - 1105, The ragged_workspace_size helper currently sizes only
max_total_seq_len and can under-allocate when offsets contain inter-batch gaps.
Update the backward workspace sizing flow around ragged_workspace_size to use
the physical ragged span through the final batch boundary, or explicitly reject
unsupported partially packed offsets with documentation; preserve correct
stride-based sizing for all axes. Add a regression test covering inter-batch
padding and verifying the workspace does not overlap subsequent allocations.
| diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}), | ||
| is_ragged_or_padded_or_full=RandomChoice({"ragged" : 1, "padded" : 0, "full" : 0}), | ||
| with_sink_token=RandomChoice({True : 1, False : 3}), | ||
| ragged_stats_layout=RandomChoice({"token_major" : 1, "head_major" : 1}), |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -u
rg -n -C 8 'TensorUid\.stats|with_score_max|with_score_sum_exp|stride_stats|exec_sdpa' \
test/python/sdpa test/python/test_mhas_v2.py || trueRepository: NVIDIA/cudnn-frontend
Length of output: 50377
🏁 Script executed:
# Check if stride_stats is used in graph construction or only when allocating stats
rg -A 5 -B 5 'stride_stats' test/python/sdpa/fp16.py | head -100Repository: NVIDIA/cudnn-frontend
Length of output: 6398
🏁 Script executed:
# Trace through exec_sdpa to see if it reads stride_stats for forward inference
ast-grep outline test/python/sdpa/fp16.py --pattern 'def exec_sdpa'Repository: NVIDIA/cudnn-frontend
Length of output: 336
🏁 Script executed:
# Check if there's graph construction that references stats stride even when stats=None
rg 'set_stride.*stats|stride.*stats' test/python/sdpa/ -A 2 -B 2 | head -80Repository: NVIDIA/cudnn-frontend
Length of output: 7625
🏁 Script executed:
# Check what happens with stats_ragged_offset in forward path when stats is None
sed -n '220,260p' test/python/sdpa/fp16.pyRepository: NVIDIA/cudnn-frontend
Length of output: 2382
🏁 Script executed:
# Verify that backward test enables stats allocation
sed -n '449,478p' test/python/test_mhas_v2.pyRepository: NVIDIA/cudnn-frontend
Length of output: 1902
🏁 Script executed:
# Check if graph building code conditionally uses stats_ragged_offset
sed -n '230,250p' test/python/sdpa/fp16.pyRepository: NVIDIA/cudnn-frontend
Length of output: 1361
🏁 Script executed:
# Find the actual graph building code that uses stats_ragged_offset
sed -n '250,310p' test/python/sdpa/fp16.pyRepository: NVIDIA/cudnn-frontend
Length of output: 5006
🏁 Script executed:
# Check the complete graph building logic for stats
sed -n '295,330p' test/python/sdpa/fp16.pyRepository: NVIDIA/cudnn-frontend
Length of output: 2855
🏁 Script executed:
# Find where the stats tensor is created and where stats_ragged_offset is used
sed -n '325,360p' test/python/sdpa/fp16.pyRepository: NVIDIA/cudnn-frontend
Length of output: 1741
🏁 Script executed:
# Check if stats_ragged_offset is used only when stats exists
rg -A 3 'stats.set_ragged_offset' test/python/sdpa/fp16.pyRepository: NVIDIA/cudnn-frontend
Length of output: 445
🏁 Script executed:
# Get broader context around stats.set_ragged_offset usage in forward graph
sed -n '360,385p' test/python/sdpa/fp16.pyRepository: NVIDIA/cudnn-frontend
Length of output: 1187
Move the forward test's ragged-stats-layout randomization to the backward test or enable score outputs in forward.
The forward test sets is_infer = True (default) and does not enable with_score_max or with_score_sum_exp. At lines 149–151 of test/python/sdpa/fp16.py, TensorUid.stats is allocated only when cfg.is_train. The stats_ragged_offset tensor is still created as a graph input (line 301) when ragged, but none of the output tensors (stats, score_max, score_sum_exp) consume it in forward inference mode. The randomization at line 368 does not provide real coverage of the stride configuration. The backward test at line 469 already exercises both layouts because it sets is_infer = False.
🤖 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 368, The ragged_stats_layout
randomization in the forward test configuration (around line 368) does not
provide meaningful coverage because the forward test runs with is_infer=True
(default), which means TensorUid.stats and related output tensors are not
allocated or consumed, so the stats_ragged_offset input tensor is never actually
used. Remove the RandomChoice ragged_stats_layout configuration from the forward
test and relocate it to the backward test configuration (around line 469), where
is_infer=False ensures both layout options are properly exercised.
Alternatively, enable with_score_max or with_score_sum_exp in the forward test
to activate the stats outputs and allow the randomization to provide real
coverage.
There was a problem hiding this comment.
Pull request overview
Fixes a correctness-critical workspace under-allocation in ragged (THD) sdpa_backward when max_total_seq_len_* is set and the Stats tensor uses a non-token-major (e.g., head-major) layout—preventing silent all-zero gradients due to downstream workspace corruption in the backend engine workspace region.
Changes:
- Update ragged workspace sizing in
CompositeSDPABackwardNodeto account for full tensor extents/strides rather than onlystride[2]. - Expand Python ragged SDPA randomized tests to cover both token-major and head-major packed
Statslayouts, including correct packed capacity and ragged offset handling. - Deduplicate packed token-capacity computation into a shared helper (
packed_token_capacity) and use it consistently.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
include/cudnn_frontend/node/scaled_dot_product_flash_attention.h |
Fixes ragged backward intermediate workspace sizing via a stride/dim-aware span calculation when max_total_seq_len_* is provided. |
test/python/sdpa/random_config.py |
Adds packed_token_capacity() and introduces ragged_stats_layout to randomize token-major vs head-major packed stats strides for ragged configs. |
test/python/sdpa/fp16.py |
Uses packed_token_capacity(); allocates ragged Stats/score_* tensors and their ragged offsets based on cfg.stride_stats (supports head-major). |
test/python/test_mhas_v2.py |
Extends ragged fwd/bwd random tests to include ragged_stats_layout randomization. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
@cudnn-ci-bot run |
Only allowlisted maintainers can use |
|
@cudnn-ci-bot run backend,frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-462-f0d2dc4 |
|
Thanks @vedaanta. Triaging the two CodeRabbit comments — I checked both against the code, and I'd like to land this as-is and carry them separately. On the test coverage comment ( On the partially-packed workspace comment (
The underlying issue is that The suggested remedy of computing a physical-capacity bound isn't available to the frontend at build time: the gaps live in the ragged-offset device tensor and aren't known until execute. The only lever here is the host-side |
|
note to self: claude::88367f80-c82c-4c8f-93a2-e66ed5a89e07 — "审计PyTorch中CuDNN Attention相关问题" |
Ragged Q/K/V/O were always generated with canonical packed strides (token stride == h*d): the gaps stride machinery was dense-only, and PR NVIDIA#462 varies only the STATS layout. So no sweep could ever produce a tensor like a K/V view of a kv-interleaved [T, 2, H, D] buffer (token stride 2*h*d) — the layout torch.nn.attention.varlen users get by slicing a fused KV projection. New with_ragged_token_gap knob in both ragged random sweeps: each of Q/K/V/O independently draws token stride h*d + gap, gap in {0, 8, 64, roundup8(h*d)} seeded from rng_geom_seed (deterministic through serialize/deserialize repro; explicit strides still win). Gaps stay multiples of 8 elements so ragged base addresses keep the packed layout's alignment class (the graph API requires 16-byte-aligned pointers). Harness: ragged buffers (incl. gradients) are allocated with the configured strides, and ragged offsets scale by each tensor's actual stride[2] — the same generalization NVIDIA#462 made for stats offsets. KNOWN FAILURES this exposes (intentionally not masked): the FROST THD forward engines (sdpa_fwd_prefill_sm120, sdpa_fwd_prefill_sm100_d128) claim non-packed-stride THD graphs and silently mis-address them (100% of O wrong; the stride ORDER is still BSHD, so the order-only layout gate passes). Backend engines serve every gapped combination correctly, fwd and bwd. The failing configs under CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 are the repro set for fixing the THD lowerings to honor declared strides. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Ragged Q/K/V/O were always generated with canonical packed strides (token stride == h*d): the gaps stride machinery was dense-only, and PR NVIDIA#462 varies only the STATS layout. So no sweep could ever produce a tensor like a K/V view of a kv-interleaved [T, 2, H, D] buffer (token stride 2*h*d) — the layout torch.nn.attention.varlen users get by slicing a fused KV projection. New with_ragged_token_gap knob in both ragged random sweeps: each of Q/K/V/O independently draws token stride h*d + gap, gap in {0, 8, 64, roundup8(h*d)} seeded from rng_geom_seed (deterministic through serialize/deserialize repro; explicit strides still win). Gaps stay multiples of 8 elements so ragged base addresses keep the packed layout's alignment class (the graph API requires 16-byte-aligned pointers). Harness: ragged buffers (incl. gradients) are allocated with the configured strides, and ragged offsets scale by each tensor's actual stride[2] — the same generalization NVIDIA#462 made for stats offsets. KNOWN FAILURES this exposes (intentionally not masked): the FROST THD forward engines (sdpa_fwd_prefill_sm120, sdpa_fwd_prefill_sm100_d128) claim non-packed-stride THD graphs and silently mis-address them (100% of O wrong; the stride ORDER is still BSHD, so the order-only layout gate passes). Backend engines serve every gapped combination correctly, fwd and bwd. The failing configs under CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 are the repro set for fixing the THD lowerings to honor declared strides. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Ragged Q/K/V/O were always generated with canonical packed strides (token stride == h*d): the gaps stride machinery was dense-only, and PR NVIDIA#462 varies only the STATS layout. So no sweep could ever produce a tensor like a K/V view of a kv-interleaved [T, 2, H, D] buffer (token stride 2*h*d) — the layout torch.nn.attention.varlen users get by slicing a fused KV projection. New with_ragged_token_gap knob in both ragged random sweeps: each of Q/K/V/O independently draws token stride h*d + gap, gap in {0, 8, 64, roundup8(h*d)} seeded from rng_geom_seed (deterministic through serialize/deserialize repro; explicit strides still win). Gaps stay multiples of 8 elements so ragged base addresses keep the packed layout's alignment class (the graph API requires 16-byte-aligned pointers). Harness: ragged buffers (incl. gradients) are allocated with the configured strides, and ragged offsets scale by each tensor's actual stride[2] — the same generalization NVIDIA#462 made for stats offsets. KNOWN FAILURES this exposes (intentionally not masked): the FROST THD forward engines (sdpa_fwd_prefill_sm120, sdpa_fwd_prefill_sm100_d128) claim non-packed-stride THD graphs and silently mis-address them (100% of O wrong; the stride ORDER is still BSHD, so the order-only layout gate passes). Backend engines serve every gapped combination correctly, fwd and bwd. The failing configs under CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 are the repro set for fixing the THD lowerings to honor declared strides. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Ragged Q/K/V/O were always generated with canonical packed strides (token stride == h*d): the gaps stride machinery was dense-only, and PR NVIDIA#462 varies only the STATS layout. So no sweep could ever produce a tensor like a K/V view of a kv-interleaved [T, 2, H, D] buffer (token stride 2*h*d) — the layout torch.nn.attention.varlen users get by slicing a fused KV projection. New with_ragged_token_gap knob in both ragged random sweeps: each of Q/K/V/O independently draws token stride h*d + gap, gap in {0, 8, 64, roundup8(h*d)} seeded from rng_geom_seed (deterministic through serialize/deserialize repro; explicit strides still win). Gaps stay multiples of 8 elements so ragged base addresses keep the packed layout's alignment class (the graph API requires 16-byte-aligned pointers). Harness: ragged buffers (incl. gradients) are allocated with the configured strides, and ragged offsets scale by each tensor's actual stride[2] — the same generalization NVIDIA#462 made for stats offsets. KNOWN FAILURES this exposes (intentionally not masked): the FROST THD forward engines (sdpa_fwd_prefill_sm120, sdpa_fwd_prefill_sm100_d128) claim non-packed-stride THD graphs and silently mis-address them (100% of O wrong; the stride ORDER is still BSHD, so the order-only layout gate passes). Backend engines serve every gapped combination correctly, fwd and bwd. The failing configs under CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 are the repro set for fixing the THD lowerings to honor declared strides. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…516) * test(sdpa): fuzz per-tensor ragged token-stride gaps in test_mhas_v2 Ragged Q/K/V/O were always generated with canonical packed strides (token stride == h*d): the gaps stride machinery was dense-only, and PR #462 varies only the STATS layout. So no sweep could ever produce a tensor like a K/V view of a kv-interleaved [T, 2, H, D] buffer (token stride 2*h*d) — the layout torch.nn.attention.varlen users get by slicing a fused KV projection. New with_ragged_token_gap knob in both ragged random sweeps: each of Q/K/V/O independently draws token stride h*d + gap, gap in {0, 8, 64, roundup8(h*d)} seeded from rng_geom_seed (deterministic through serialize/deserialize repro; explicit strides still win). Gaps stay multiples of 8 elements so ragged base addresses keep the packed layout's alignment class (the graph API requires 16-byte-aligned pointers). Harness: ragged buffers (incl. gradients) are allocated with the configured strides, and ragged offsets scale by each tensor's actual stride[2] — the same generalization #462 made for stats offsets. KNOWN FAILURES this exposes (intentionally not masked): the FROST THD forward engines (sdpa_fwd_prefill_sm120, sdpa_fwd_prefill_sm100_d128) claim non-packed-stride THD graphs and silently mis-address them (100% of O wrong; the stride ORDER is still BSHD, so the order-only layout gate passes). Backend engines serve every gapped combination correctly, fwd and bwd. The failing configs under CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 are the repro set for fixing the THD lowerings to honor declared strides. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * test(sdpa): simplify ragged token gaps to whole tokens; always-on in the fp16 ragged sweeps Per review: drop the True/False dice — the fp16 ragged sweeps enable the knob unconditionally, and the packed case comes from the gap draw itself. Each of Q/K/V/O independently draws gap = n*h*d tokens, n in 0..3 (n=0 packed, n=1 a kv-interleaved [T, 2, H, D] view, n=2 a [T, 3, H, D] QKV-interleave). Whole-token gaps keep every ragged base address in the packed layout's alignment class by construction, replacing the previous multiples-of-8-elements rule. The config field stays default-False: fixed configs and the fp8/mxfp8 harnesses still assume packed allocations, and the cu_ragged form derives offsets internally. Verified: backend engines pass 5/5 seeded repros, a bf16 training config, and a 128-test slice of the fwd ragged L0 sweep; FROST THD fwd engines keep failing gapped draws (the intended standing repro set). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * test(sdpa): draw all four ragged token gaps up front; drop unused binding The gap RNG was consumed lazily, one draw per stride left unspecified — so explicitly pinning e.g. stride_q shifted the gaps K/V/O derive from the same rng_geom_seed. Draw all four values up front in fixed Q/K/V/O order and apply each only where the stride is missing: per-tensor layouts are now a function of the seed alone. The all-defaults path (the sweeps) consumes draws in the same order as before, so existing seeded repros reproduce identical strides. Adds a GPU-free regression test pinning stride_q and asserting the K/V/O strides match the all-defaults derivation (seed chosen so the old lazy behavior visibly shifts two of the three gaps). Also renames the unused batch binding in compute_packed_strides (RUF059). Addresses CodeRabbit review feedback on #516. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * test(sdpa): ragged token gaps on by default; auto-packed where not yet expressible Drop the per-sweep with_ragged_token_gap opt-in: the field now defaults True, so EVERY ragged config fuzzes its per-tensor token strides. fill_derived_fields falls back to packed automatically where a gap is not yet expressible or handled — cu / offset-multiplier forms bind offsets as cu (x multiplier) and cannot declare a token gap (#538), and the fp8/mxfp8 harnesses (1-byte data types) allocate assuming packed (#537) — so mixed sweeps (ragged + cu_ragged in one RandomChoice) gap exactly the draws that support it. Explicit strides are never touched (the gap only fills strides left None), so recorded repro dicts reproduce exactly. The regression test also locks in the new semantics: default-on for plain ragged, packed for the three fallback forms. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
What
A ragged (THD)
sdpa_backwardgraph that setsmax_total_seq_len_qsilently returns all-zero dQ/dK/dV when theStatstensor is head-major. The forwardOandstatsare correct and no error is raised, so this corrupts training without ever surfacing. Present unchanged in 9.22, 9.23, 9.24, 9.25 and 9.26, on both sm90 and sm100.Head-major
[h, total_q]is FlashAttention's and PyTorch varlen'ssoftmax_lselayout, so this is reachable from exactly the integrations that would reach formax_total_seq_lenin the first place.Root cause
max_total_seq_len_*is documented as a workspace-sizing hint. It is not — setting it makesCompositeSDPABackwardNodere-lay-out its non-virtual intermediates onto the corresponding user tensor's stride and ragged offset, and then size them from a single stride element:stride[2]is the sequence stride, so the formula assumes the per-token footprint isstride[2]. That is true only for token-major[t, h, 1]stats, wherestride[2] == h_q. For head-major[h, t]statsstride[2] == 1, and the buffer is under-allocated by exactlyh_q.softmax_sumthen writes(h_q - 1) * t_q * 4bytes past its slot. The frontend places the backend engine's own workspace immediately after the FE region (graph_interface.h:engine_workspace = workspace + fe_workspace_size), and the first word of that region is the persistent scheduler's tile counter. Every CTA claims work withtile_id = atomicAdd(tile_id_counter, 1)and retires unlesstile_id < num_tiles;tile_idisuint32_t, so any nonzero float bit pattern is astronomically out of range. Every CTA therefore retires without storing anything — no dQ accumulation, no dK/dV TMA store — and the gradients are simply never written. No fault, no error, no hang.Measured directly, comparing the exact
int32bits of the clobbered word against the truesoftmax_sumvalue at the same element from a clean run at the same geometry:A float bit pattern plus exactly 128 integer increments — one
atomicAddper CTA, every one of which then failed the tile bound and retired.The mask plays no part in the trigger:
TOP_LEFTand the completely unmasked case fail identically. It only decides whether the clobbering bytes happen to be zero — a fully-masked leading query row has row-sum 0, which is what the memset already put there. That is why a handful of geometries appear to work, and why one intermediate case produces|dQ| ≈ 6e36garbage rather than zeros.The fix (commit 1)
Size the workspace from every axis instead of from
stride[2]alone:applied to
softmax_sum,dQ_accum,dK_fullheadanddV_fullhead.For token-major inputs the new expression is bit-identical to the old one —
softmax_sumgives(maxT-1)·h_q + (h_q-1)·1 + 0 + 1 = maxT·h_q,dQ_accumgives(maxT-1)·h_q·d + (h_q-1)·d + (d-1) + 1 = maxT·h_q·d. Allocations are unchanged for every layout that works today; only non-token-major layouts move.Why this was never caught (commit 2)
test/python/sdpa/fp16.pyonly ever allocated token-major ragged stats, the one layout for which the old formula happens to be exactly right. There was no coverage of any other packedStatslayout, so the suite was structurally blind to this class of defect.Commit 2 adds a
ragged_stats_layoutaxis toExecConfig, randomized in the ragged fwd/bwd tests, choosing between token-major and head-major. The stats allocation and the stats ragged offset now both derive fromcfg.stride_statsrather than hardcoding the token-major assumption.Against the unpatched frontend the new axis reproduces the defect immediately — 11 tests fail with
Backward workspace overwritten outside its boundaries(the suite's existing guard-band check) before the run aborts. The first failing config hasstride_stats = (58880, 7360, 1, 1), i.e. sequence stride 1.Verification
-k ragged(fwd + bwd), sm100Sanity check on the numbers rather than just pass/fail: with the fix, a
TOP_LEFT2048/512 case withmax_total_seq_lenset gives|dQ| = 5.53e5, exactly matching the same case withmax_total_seq_lenunset; aBOTTOM_RIGHT512/512 case gives1.976e5in both.Not addressed here
docs/operations/Attention.mdstill describesmax_total_seq_len_*as a workspace allocation hint. It changes intermediate tensor layout and can change results, and should say so.ragged_offset == cum_token × stride[2]. Offsets carrying inter-sequence padding would still under-allocate, and the frontend cannot detect this since the offsets are device data. This wants either a documented contract or formax_total_seq_lento move into the backend attention descriptor.Summary by CodeRabbit
Bug Fixes
Tests