Fix stream-ordering and input-validation holes in the SM100 DSA backward interface - #429
Open
zkyue wants to merge 5 commits into
Open
Fix stream-ordering and input-validation holes in the SM100 DSA backward interface#429zkyue wants to merge 5 commits into
zkyue wants to merge 5 commits into
Conversation
flash_attn_bwd_sm100 allocates dq/dkv/d_sink and the two workspaces (and makes contiguity copies) with plain torch calls, which enqueue on the ambient torch stream, while the kernel launches on the caller-provided current_stream. When the caller passes a non-default stream, the semantically required zero-initialization of dkv/d_sink and the workspaces is unordered with the kernel: a busy ambient stream lets the zero-fills land after the kernel and wipe the accumulated gradients (or, in the other interleaving, the kernel accumulates into uninitialized memory). Resolve the stream first and scope the normalization/allocation section with torch_stream_context(current_stream), the same pattern the other DSA interfaces (score_recompute, indexer_forward, indexer_backward) already use. The default-stream path is unchanged. Add a deterministic regression test that keeps the ambient stream busy with torch.cuda._sleep while launching on a side stream: on the unpatched interface the returned dkv comes back all-zero. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
The dtype checks in SparseAttentionBackward.check_support and flash_attn_bwd_sm100 accept both fp16 and bf16, but the SM100 kernel (FlashAttentionDSABackwardSm100) hardcodes BF16 as its element type and never receives the input dtype. fp16 inputs on SM100 pass the checks, compile, run without any error, and return silently wrong gradients: on the same reference harness where bf16 passes, ~96% of the fp16 dq elements fall outside 5e-2 tolerances against the fp16 autograd reference. Restrict the dtype gate to bf16 when dispatching to SM100 (the SM90 kernels are dtype-parameterized and keep fp16), update the DSA docs and the DSA backward benchmark (which offered --dtype float16 uncondition- ally) to match, and skip the fp16 benchmark combination on non-SM90. Plumbing the dtype through the SM100 kernel would restore fp16 there and is left as a follow-up. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
flash_attn_bwd_sm100 derives every kernel dimension from q and passes the companion tensors through with no cross-tensor shape validation. Since the compiled kernel treats all dimensions as dynamic values, a mis-shaped companion tensor does not fail: a transposed dout or lse runs without any error and returns silently corrupted gradients (measured rel-L2 vs the correct result: ~1.1 and ~45 respectively). - Assert the shape contract of kv/out/dout/lse/attn_sink/topk_idxs/ topk_length against q in the interface, in the same style as the existing dq/dkv out-parameter asserts, and require all inputs on q's device (the launch-stream context is bound to that device). - Enforce the same contract in SparseAttentionBackward.check_support, which is the advertised metadata-only support gate (it previously accepted any companion shapes and omitted out/dout/topk_length dtype checks). - Extend the contiguity normalization, which covers q/kv/out/dout/lse, to attn_sink/topk_idxs/topk_length: non-contiguous aux tensors currently escape down to the CuTe DSL layer and fail there with low-level stride errors (a signature mismatch against the shared compile-cache entry on the warm path, a leading-stride assert on the cold path). - Require caller-provided dq/dkv to be contiguous: the compile cache is keyed without output strides, so a strided out-parameter would be written through the wrong layout (it cannot be silently copied without breaking out-parameter identity). Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughSM100 sparse-attention backward now validates cross-tensor contracts, supports FP16 and BF16 element dtypes, performs setup on the resolved CUDA stream, and adds SM100-gated tests for stream ordering, numerical behavior, tensor normalization, and validation. ChangesDSA backward contract and execution updates
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant flash_attn_bwd_sm100
participant torch_stream_context
participant CUDA_kernel
Caller->>flash_attn_bwd_sm100: submit backward tensors
flash_attn_bwd_sm100->>torch_stream_context: scope work to resolved CUDA stream
torch_stream_context->>flash_attn_bwd_sm100: convert tensors and zero outputs
flash_attn_bwd_sm100->>CUDA_kernel: launch SM100 backward operation
CUDA_kernel-->>Caller: return dq, dkv, and d_sink
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
2 tasks
Collaborator
|
Thanks @zkyue for the contribution. |
Anerudhan
added
orig-external
Contributor
|
Thanks @zkyue for tracking these issues down and adding the regression coverage. I reviewed the three commits and ran focused checks on B200. The stream-ordering fix looks correct to me. For FP16 on SM100, I would prefer enabling it rather than rejecting it. I prepared a small commit based directly on the current head of this PR: Both cache layers already include the dtype, so the change only threads the mapped input dtype into FlashAttentionDSABackwardSm100 and replaces the hardcoded BF16 element type. It also restores the FP16 documentation/benchmark behavior and replaces the rejection test with an FP16 numerical regression. I tested D=512 on B200 both with and without topk_length; the dq/dkv/d_sink relative L2 errors against the FP32 autograd reference were approximately 2.5e-4–3.0e-4. The existing BF16 and stream/contract regression tests also pass. Please feel free to cherry-pick bdbb731 if this direction looks good. |
Thread the interface dtype into FlashAttentionDSABackwardSm100 instead of hardcoding BF16. Both interface cache layers already include dtype, so no cache-key changes are needed. Restore FP16 API, documentation, and benchmark support, and replace the rejection coverage with an SM100 numerical regression against the FP32 autograd reference. The incorrect-FP16 behavior and reproduction were identified by @zkyue in NVIDIA#429. Signed-off-by: Jiayu Sun <jiayus@nvidia.com> (cherry picked from commit bdbb731)
Contributor
Author
|
Thanks @jiayus-nvidia — I verified bdbb731 independently on a B200 before taking it, and everything checks out:
Cherry-picked onto the branch with your authorship preserved (now |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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 (2)
python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py (2)
62-75: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winMake the support gate reject mixed-device inputs.
These checks validate dtype but not device placement, so
check_support()can returnTruefor CPU descriptors or tensors split across CUDA devices. The SM100 runtime rejects the same call at Line 70, making the advertised support gate inconsistent with execution. Validate all descriptor devices and CUDA placement here using the APIBase/device helper.🤖 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/deepseek_sparse_attention/sparse_attention_backward/api.py` around lines 62 - 75, Update the support validation around the dtype checks in the API class so check_support() rejects non-CUDA descriptors and inputs distributed across different CUDA devices. Reuse the existing APIBase/device helper to validate every descriptor’s device, including the optional topk_length_desc, before reporting support; keep the current dtype validation unchanged.
80-85: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winReject unsupported head dimensions in both validation layers.
The kernel’s layout and KV-loading logic only support
head_dim=512andhead_dim=576; the current fallback accepts other values and can produce incorrect or uninitialized output.
python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py#L80-L85: reject anyhead_dimnot in(512, 576)before derivingexpected_o_shape.python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py#L75-L81: add the equivalent runtime assertion for direct callers.🤖 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/deepseek_sparse_attention/sparse_attention_backward/api.py` around lines 80 - 85, Reject unsupported head dimensions in both validation layers: in python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py lines 80-85, validate that head_dim is 512 or 576 before deriving expected_o_shape; in python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py lines 75-81, add the equivalent runtime assertion for direct callers. Preserve the existing handling for supported dimensions.
🤖 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 `@python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py`:
- Around line 62-75: Update the support validation around the dtype checks in
the API class so check_support() rejects non-CUDA descriptors and inputs
distributed across different CUDA devices. Reuse the existing APIBase/device
helper to validate every descriptor’s device, including the optional
topk_length_desc, before reporting support; keep the current dtype validation
unchanged.
- Around line 80-85: Reject unsupported head dimensions in both validation
layers: in
python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py lines
80-85, validate that head_dim is 512 or 576 before deriving expected_o_shape; in
python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py
lines 75-81, add the equivalent runtime assertion for direct callers. Preserve
the existing handling for supported dimensions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ae5ea95f-6831-4732-9538-770c82bc3383
📒 Files selected for processing (4)
python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.pypython/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.pypython/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.pytest/python/fe_api/dsa/test_DSA_sparse_attention_backward.py
🚧 Files skipped from review as they are similar to previous changes (1)
- test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py
SparseAttentionBackward.check_support validates dtype and the
cross-tensor shape contract, but two gaps (both outside the FP16 diff)
let it accept inputs the SM100 runtime then rejects or crashes on.
- Device placement: check_support never inspected any descriptor's
device, so an all-CPU descriptor set or a cross-CUDA-device split
passed the gate (verified: CPU inputs returned True) even though
flash_attn_bwd_sm100 asserts that every input is a CUDA tensor on
Q's device. Validate that Q is on CUDA and that every descriptor
(including the optional topk_length) shares Q's device, using the
existing _value_error_if helper; the dtype checks are unchanged.
- head_dim: the SM100 kernel is tiled only for head_dim in {512, 576}
(the 576 MLA case splits QK=576 / V=512); any other head_dim takes
the non-512 KV-load path and indexes shared memory out of bounds.
check_support returned True for head_dim=128 (verified). Gate
head_dim in check_support (ValueError) and mirror it with a runtime
assert in flash_attn_bwd_sm100, before any compile/launch.
Negative coverage added for both: an all-CPU input, a cross-device
input (Q on CUDA, KV on CPU), and head_dim=128 now raise at the
support gate, and head_dim=128 raises at the runtime interface.
Supported configurations (head_dim 512 in BF16/FP16, head_dim 576 MLA)
are unchanged.
Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
Contributor
Author
|
Also addressed the two CodeRabbit findings from the latest scan (both verified real, both in the validation layers this PR owns): |
2 tasks
Contributor
|
Thanks, LGTM. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Before submitting
pre-commit runand committed any formatting changes.Affected area
FE OSS kernels or CuTeDSL
Summary
Three correctness/hardening fixes for the Python interface of the SM100 DSA
backward (
sparse_attention_backward/_interface_sm100.py,api.py), onecommit each. All three failure modes were confirmed empirically on
develop @3a9ed3f (B200); each commit carries regression tests that fail on
the parent commit and pass with the fix.
Stream-order the allocations with the launch stream (185728d).
flash_attn_bwd_sm100allocates dq/dkv/d_sink and the two workspaces(and makes contiguity copies) with plain torch calls, which enqueue on
the ambient torch stream, while the kernel launches on the
caller-provided
current_stream. With a non-defaultcurrent_streamthe semantically required zero-fills of dkv/d_sink and the workspaces
are unordered with the kernel. Observed on develop: with the ambient
stream busy, the zero-fills land after the kernel and the returned dkv
comes back all-zero (3/3 trials; the other interleaving lets the
kernel accumulate into uninitialized memory). Fix: resolve the stream
first and scope the normalization/allocation section with
torch_stream_context(current_stream)— the same pattern thescore_recompute / indexer_forward / indexer_backward interfaces already
use.
Reject fp16 on SM100 (05c0965).
The dtype checks (
api.pycheck_support, interface assert) acceptfp16, but
FlashAttentionDSABackwardSm100hardcodes BF16 as its elementtype and never receives the input dtype. Observed on develop: fp16
inputs compile, run without any error, and return silently wrong
gradients — ~96% of dq elements outside 5e-2 tolerances against the
fp16 autograd reference, on the same harness where bf16 passes. Fix:
gate the SM100 dispatch to bf16 (the SM90 kernels are
dtype-parameterized and keep fp16), and update the two in-repo surfaces
that advertised fp16 unconditionally (
docs/fe-oss-apis/dsa.md, thebenchmark/dsabackward benchmark — which now skips thefp16-on-non-SM90 combination). Plumbing the dtype through the SM100
kernel would restore fp16 there and is left as a follow-up.
Validate the input contract (2536e6d).
Every kernel dimension is derived from
qand the compiled kerneltreats dimensions as dynamic values, so mis-shaped companion tensors do
not fail. Observed on develop: a transposed
doutorlseruns withoutany error and returns silently corrupted gradients (rel-L2 vs the
correct result: ~1.1 and ~45 respectively). Fix:
topk_length against
qin the interface (style of the existingdq/dkv out-parameter asserts), and require all inputs on
q's device;SparseAttentionBackward.check_support,the advertised metadata-only support gate (it previously accepted any
companion shapes and omitted out/dout/topk_length dtype checks);
only) to attn_sink/topk_idxs/topk_length, which today escape to the
CuTe DSL layer and fail there with low-level stride errors (signature
mismatch against the shared compile-cache entry on the warm path, a
leading-stride assert on the cold path);
keyed without output strides, so a strided out-parameter would be
written through the wrong layout (silently copying would break
out-parameter identity).
Related issues
None filed; found during an interface audit and confirmed empirically on
develop.
API and compatibility impact
(S=512, S_kv=2048, H=64, D=576, topk=128; both with and without
topk_length), dq and d_sink are bitwise-identical before/after this
PR in both configurations. dkv is an FP32 atomic accumulation followed
by BF16 storage and is not bitwise-stable even against itself
(base-vs-base rerun: rel-L2 9.2e-6 / 80 of 1.18M elements differ without
topk_length, 8.6e-6 / 27 with); base-vs-fixed sits inside that jitter
(7.5e-6 / 87 and 1.5e-6 / 21 respectively).
ValueErrorfromcheck_support/AssertionErrorfrom the interface instead of silently returning wronggradients. SM90 fp16 acceptance is unchanged (code-inspected; we had no
SM90 machine to run it, maintainers' CI covers that arch).
interface,
ValueErrorfromcheck_support) instead of running withcorrupted results; non-contiguous attn_sink/topk_idxs/topk_length are
normalized instead of failing with DSL-level stride errors;
non-contiguous caller-provided dq/dkv are rejected.
_interface_sm90.py) shares theallocation-before-stream-resolve pattern and the narrower contiguity
normalization; it is left untouched here (no SM90 hardware available to
validate a change) and could receive the same treatment as a follow-up.
Testing
All on B200 (SM100), CUDA 13.3 toolchain, nvidia-cutlass-dsl 4.6.1, torch 2.13.
New regression tests:
All five fail on the parent commit (the stream test with dkv returned
all-zero; the fp16 and contract tests because nothing raises; the
noncontiguous test with a DSL stride error) and pass with the fixes.
The stream test parks the ambient stream on
torch.cuda._sleepand ismarked
gpu_exclusive(a concurrently busy GPU could mask theordering violation).
Full DSA suite
pytest test/python/fe_api/dsa -q:before:
4 failed, 26 passed, 4 skipped— after: the same 4 failed,31 passed, 4 skipped. The pre-existing failures(2×
test_DSA_indexer_top_k[705],2×
test_DSA_sparse_attention_backward_wrapper[...576...]) failidentically on the unpatched base; the wrapper[576] pair looks related
to the attn-sink issue addressed by DSA: fix SM100 sink normalization for 576/512 backward #421.
Default-path bitwise check as described under compatibility impact.
Summary by CodeRabbit