DSA backward SM100: reject topk_length entries < 1 (a zero entry hangs the kernel) - #433
Closed
zkyue wants to merge 3 commits into
Closed
DSA backward SM100: reject topk_length entries < 1 (a zero entry hangs the kernel)#433zkyue wants to merge 3 commits into
zkyue wants to merge 3 commits into
Conversation
A topk_length entry of 0 gives the SM100 kernel zero KV tiles for that token (tile_count = ceil_div(0, block_tile) = 0 in dsa_bwd_sm100), a case the kernel has no defined behavior for. Measured on B200 (head_dim 512, head_dim_v 512, h 64, topk 512, s_q 128, s_kv 4096): a call whose topk_length contains zero entries (3 of 128 rows) never terminates -- the launch hangs at 100% SM until killed -- while the identical call with all lengths >= 1 completes normally (~40 s including first-call JIT compilation). Code inspection suggests a zero-tile pipeline deadlock, and that a zero-tile token reaching the dq epilogue would store accumulator TMEM that no MMA instruction has written. The mathematically correct dq for a query attending to nothing is exactly 0. Make the failure loud at the interface layer instead of hanging the device: flash_attn_bwd_sm100 raises ValueError when any topk_length entry is < 1 (zero entries are the measured hazard; negatives are rejected by the same check). The check needs a device-to-host sync, so it runs outside CUDA graph capture only -- under capture the caller must uphold the precondition, which is documented on the interface and on sparse_attention_backward_wrapper. Add a test covering 0 and a negative entry. No kernel change: whether topk_length == 0 should instead be supported input (skipping the pipeline handshake and zero-filling dq for zero-tile tokens) is left as an upstream semantic decision; this is minimal interface validation and is trivially removable in that world. 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 (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe SM100 sparse-attention backward path validates ChangesSM100
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
🚥 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py (1)
307-317: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winReconsider the
L0classification.On a cold cache, the wrapper compiles the SM100 object before
execute()reaches this validation (see Lines [176]-[192] inapi.py). This test can therefore pay the full CuTe compilation cost despite never launching the kernel. Move it to the appropriate higher level, or verify that compilation is already warmed in this test suite.🤖 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/fe_api/dsa/test_DSA_sparse_attention_backward.py` around lines 307 - 317, Reclassify test_DSA_sparse_attention_backward_topk_length_zero_raises from L0 because wrapper compilation occurs before validation and can incur a cold-cache CuTe compile. Move the test to the appropriate higher-level suite, or ensure the suite explicitly warms the SM100 compilation before running this validation-only test.
🤖 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/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py`:
- Around line 87-102: Resolve the launch stream from current_stream before the
CUDA graph-capture and topk_length validation in
SparseAttentionBackward.execute. Use that resolved stream for capture detection
and the required topk_length synchronization, or explicitly reject non-current
streams; ensure explicit-stream capture cannot perform an illegal sync and
empty-row validation remains enforced on the launch stream.
In `@test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py`:
- Around line 324-326: Update the test’s SM100 guard to reuse the repository’s
existing DSA support check and backend-version gate, including
cudnn.backend_version() alongside torch.cuda.get_device_capability(). Skip
unsupported architecture, dtype, or backend combinations before the DSA wrapper
setup, while preserving execution for supported configurations.
---
Nitpick comments:
In `@test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py`:
- Around line 307-317: Reclassify
test_DSA_sparse_attention_backward_topk_length_zero_raises from L0 because
wrapper compilation occurs before validation and can incur a cold-cache CuTe
compile. Move the test to the appropriate higher-level suite, or ensure the
suite explicitly warms the SM100 compilation before running this validation-only
test.
🪄 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: 2d1ade90-6ab9-4108-9e5b-d12dedd29480
📒 Files selected for processing (3)
python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.pypython/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.pytest/python/fe_api/dsa/test_DSA_sparse_attention_backward.py
Comment on lines
+324
to
+326
| major, _ = torch.cuda.get_device_capability() | ||
| if major != 10: | ||
| pytest.skip("the topk_length >= 1 guard is implemented on the SM100 interface") |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Add the repository’s backend/support gate.
The changed test checks only major == 10 and whether cudnn.DSA imports. Reuse the existing DSA support check and backend-version gating so an unsupported SM100/backend combination skips instead of failing during wrapper setup.
As per coding guidelines, test/python/**/*.py tests must gate unsupported architecture, dtype, and backend-version combinations using support checks, cudnn.backend_version(), and torch.cuda.get_device_capability().
🤖 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/fe_api/dsa/test_DSA_sparse_attention_backward.py` around lines
324 - 326, Update the test’s SM100 guard to reuse the repository’s existing DSA
support check and backend-version gate, including cudnn.backend_version()
alongside torch.cuda.get_device_capability(). Skip unsupported architecture,
dtype, or backend combinations before the DSA wrapper setup, while preserving
execution for supported configurations.
Source: Coding guidelines
…tream, test gating _interface_sm100: resolve the launch stream before the topk_length guard and evaluate both the capture check and the device-to-host sync inside torch_stream_context(current_stream). Previously the guard consulted torch.cuda.is_current_stream_capturing() on whatever torch stream was current: with an explicit current_stream, a capture on that stream was not detected (validation ran mid-capture, syncing a stream it must not), and outside capture the sync was not ordered against the caller producer work on the launch stream. Add a test that exercises both sides on a real explicit side stream: outside capture a zero row is still rejected; during capture on the explicit stream -- while the torch-current stream is a different, non-capturing stream -- the guard detects the capture on the resolved stream and skips validation per the documented capture semantics (caller upholds the precondition), the launch is recorded into the graph, and the graph is never replayed. Verified red on the previous guard placement (ValueError raised mid-capture) and green with this one. Test gating: replace the bare major == 10 check with the suite's DSA capability gate (major * 10 + minor < 100, matching test_DSA_sparse_attention_backward_staged_store and dsa_init's min_compute_capability=100). The cudnn[cutedsl] import skip remains the backend gate: this CuTe-DSL-only path does not depend on cudnn.backend_version(). L0 classification of the zero-raises test: kept, with the reasoning documented in the test docstring -- SparseAttentionBackward.compile() defers compilation to execute(), and the guard raises before the interface consults its compile cache, so the validation-only test pays no CuTe compilation cost even on a cold cache. test/python fe_api/dsa: 4 failed, 28 passed, 4 skipped; the failure set is identical to the pre-change baseline. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py`:
- Around line 360-378: Change the pytest marker on
test_DSA_sparse_attention_backward_topk_length_guard_explicit_stream from L0 to
L1, since its pre-capture call performs compile-cache warmup and real kernel
execution. Leave the test logic and neighboring markers unchanged.
🪄 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: 14d5f5cc-70e6-4e49-b9d5-401263c417c4
📒 Files selected for processing (3)
python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.pypython/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.pytest/python/fe_api/dsa/test_DSA_sparse_attention_backward.py
🚧 Files skipped from review as they are similar to previous changes (2)
- python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py
- python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py
…fore capturing) Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
Collaborator
|
Thanks @zkyue for the contribution. |
Anerudhan
added
orig-external
Contributor
Contributor
Author
|
Closing in favor of #439, which supports empty rows inside the kernel — no host-side synchronization and consistent behavior under CUDA graph capture. Thanks for the quick turnaround! |
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
-l 160) and committed any changes — files unchanged. (clang-format n/a: Python only.)Affected area
FE OSS kernels or CuTeDSL
Summary
sparse_attention_backwardaccepts a per-querytopk_length, but the SM100kernel has no defined behavior for a zero entry. In
dsa_bwd_sm100.pytheper-token tile count is
with no zero guard. In practice a
topk_lengthcontaining zero entrieshangs the launch: on the configuration we measured, the call never
returns — 100% SM utilization until the process is SIGKILLed — while the
identical call with every length >= 1 completes normally. Code inspection
suggests a zero-tile pipeline deadlock (the warp-specialized stages have no
tiles to hand through for that token), and that a zero-tile token reaching
the dq epilogue would store accumulator TMEM that no MMA has written; the
mathematically correct dq for a query attending to nothing is exactly 0.
This PR adds minimal interface validation, with no kernel changes:
_interface_sm100.flash_attn_bwd_sm100raisesValueErrorwhen anytopk_lengthentry is< 1. Zero entries are the measured hazard;negative values are rejected by the same check. The check needs a
device-to-host sync, so it runs outside CUDA graph capture only —
under capture the caller must uphold the precondition, which the
docstrings state explicitly;
sparse_attention_backward_wrapper;0and a negative value both raise).We deliberately did not implement kernel-side support (skipping the
pipeline handshake and zero-filling dq for zero-tile tokens): that is a
semantic decision about whether empty rows are supported input, and it
touches the same kernel region #421 is currently changing. If you would
rather support
topk_length == 0in the kernel, we are happy to follow upwith that variant instead — this guard simply becomes removable in that
world.
Why
Repro on unmodified develop @3a9ed3f (B200, CUDA 13.3, nvidia-cutlass-dsl
4.6.1, torch 2.13): s_q=128, s_kv=4096, h=64, head_dim=head_dim_v=512,
topk=512;
topk_length[{3, 64, 127}] = 0(3 of 128 rows), all other rows512:
While hung: GPU at 100% utilization, host thread parked in
torch.cuda.synchronize()(py-spy), no forward progress. Reproduces from afresh process. With this PR the same call raises
ValueErrorup front.The current test generator floors lengths at 1
(
torch.randint(1, topk_k + 1, ...)), so nothing in CI exercises 0 — thecase is reachable only from user input, which is exactly the case that
should fail loudly rather than hang the device.
Related issues
None filed; found while auditing the DSA OSS kernel contracts. Adjacent but
independent of #421 (which changes the same kernel's sink normalization and
extends the same test file — this PR touches only the Python interface layer
plus one appended test; trivial to rebase in either order).
API and compatibility impact
topk_lengthentry< 1now raiseValueErroroutsideCUDA graph capture (previously: zero entries hung the device on the
configuration we measured). Under capture the check cannot run and the
documented precondition applies. Calls satisfying the precondition are
unaffected in behavior.
(topk_length < 1).any()reduction + D2H syncper call on the SM100 path (skipped under capture). If that is unwanted on
hot paths, we can gate it behind an env var or debug flag instead —
maintainer preference.
guard is scoped to the SM100 interface where the failure is demonstrated.
Testing
Environment: B200 (SM100), CUDA 13.3, nvidia-cutlass-dsl 4.6.1, torch 2.13.
test/python/fe_api/dsasuite on this branch: 4 failed, 27 passed,4 skipped — the failure set is identical to unpatched develop @3a9ed3f
(2×
test_DSA_indexer_top_k[705], 2×test_DSA_sparse_attention_backward_wrapper[576], pre-existing and inDSA: fix SM100 sink normalization for 576/512 backward #421's scope); the +1 passed is the new guard test.
test_DSA_sparse_attention_backward_topk_length_zero_raises(the guard raises before any kernel compile/launch, so it runs in ~1 s).
The hang itself was demonstrated with the separate repro configuration
above, not by this test.
Summary by CodeRabbit
Bug Fixes
topk_lengthentries less than 1 (including0and negatives), raising a clearValueError.topk_length == 0.Tests
topk_lengthrejection.