Tell DLPack we are only reading, so probe() survives a graph capture - #506
Conversation
buffers.probe() turns a caller's buffer into (pointer, shape, strides, dtype,
device) for a kernel launch. It reads metadata and does nothing else, so it
never needed synchronisation -- but its DLPack fallback called __dlpack__()
with the default stream argument, and torch reads that as "prepare this tensor
for a consumer on another stream" and calls record_stream. That is illegal
during a CUDA graph capture.
Passing stream=-1 -- DLPack's "the caller handles synchronisation, do no
bookkeeping" -- says what probe() actually wants. The TypeError fallback covers
a producer whose __dlpack__ predates the argument.
Only bfloat16 reached this. probe() prefers __cuda_array_interface__, whose
typestr is numpy-flavoured: float16 is "<f2" and resolves, bfloat16 has no
numpy spelling so torch reports raw "<V2" and it falls through to DLPack. So
the failure needed bf16 buffers AND a capture AND probe() on the execute path,
which is why it looked like a GDN/KDA problem: 286 tests across
test_gdn_bprop_kernel, test_gdn_prefill_kernel and test_kda_prefill_kernel,
all reporting only "operation failed due to a previous error during capture",
an error that names neither this file nor this function. SDPA and GEMM use
fp16/fp8 and take the CAI path; bf16 SDPA under capture would have hit it too.
test_buffer_probe.py pins it directly rather than leaving it to the linear
attention suites to notice again. Reverting the fix turns
test_probe_is_capture_safe[bfloat16] red, which was checked.
fp8 is left out of the test: probe() has no mapping for it at all
("unsupported buffer dtype (code=10, bits=8)"), which is a separate gap.
On sm100: linear_attention 353 passed / 1769 skipped (was 286 failed / 67
passed on this same tree); with sdpa/frost, gemm/frost and the new test,
4859 passed / 3892 skipped.
|
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)
📝 WalkthroughWalkthrough
ChangesDLPack probe behavior
Estimated code review effort: 2 (Simple) | ~10 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
YangXu1990uiuc
added
cat-bug
![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
🧹 Nitpick comments (1)
python/cudnn/frost/buffers.py (1)
242-245: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winDistinguish a legacy signature mismatch from a producer failure.
The
except TypeErrorcovers the complete__dlpack__call. If a stream-aware producer raisesTypeErrorinside its implementation, Line [245] callsdl()again. This can hide the original error and repeat producer or stream side effects.Restrict the retry to a confirmed unsupported-keyword error. Add tests for a legacy signature and an internal
TypeError. Verify whether legacy producers are intentionally outside the capture-safety guarantee because the no-argument call can still perform default stream bookkeeping.🤖 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/frost/buffers.py` around lines 242 - 245, Update the __dlpack__ invocation around dl so the no-argument retry occurs only when the TypeError confirms the stream keyword is unsupported, while propagating internal producer TypeErrors without calling dl twice. Add coverage for both a legacy signature and an implementation-raised TypeError, and preserve or explicitly validate the intended capture-safety behavior for legacy producers.
🤖 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/test_buffer_probe.py`:
- Around line 58-62: Update the side-stream warmup in the test setup around
torch.cuda.Stream so it waits for the current stream before writing: call
side.wait_stream(torch.cuda.current_stream()) before t.add_(1), while preserving
the existing synchronization afterward.
- Around line 26-37: Update the BF16 parameter gating around _DTYPES and the
requires_cuda test marker to skip torch.bfloat16 when the active CUDA device
lacks native BF16 support, including compute capability major version below 8.
Keep float16 and float32 enabled whenever CUDA is available, and do not add a
cuDNN version check.
---
Nitpick comments:
In `@python/cudnn/frost/buffers.py`:
- Around line 242-245: Update the __dlpack__ invocation around dl so the
no-argument retry occurs only when the TypeError confirms the stream keyword is
unsupported, while propagating internal producer TypeErrors without calling dl
twice. Add coverage for both a legacy signature and an implementation-raised
TypeError, and preserve or explicitly validate the intended capture-safety
behavior for legacy producers.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a45e288b-59b3-4c98-aec0-aa3504055d79
📒 Files selected for processing (2)
python/cudnn/frost/buffers.pytest/python/test_buffer_probe.py
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-506-675601f |
bfloat16 needs SM80+, but the whole module was gated on CUDA presence alone, so on an older card the bf16 cases would have failed rather than skipped. Gated per dtype, so float16/float32 still run there -- they take the CAI path, which is the control for the DLPack one. The side-stream warmup did not order the zeros before the write: torch.cuda.synchronize() waits for both to finish but imposes no order between them. side.wait_stream(current_stream()) does. Both from review. Reverting the fix still turns test_probe_is_capture_safe [bfloat16] red, which was re-checked.
|
Both valid, both fixed in BF16 gating. Right — the module was gated on CUDA presence alone, so on a pre-SM80 card the bf16 cases would have failed rather than skipped. Now per-dtype: _BF16_OK = torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8
_DTYPES = [torch.float16, torch.float32,
pytest.param(torch.bfloat16, marks=pytest.mark.skipif(not _BF16_OK, ...))]Per-dtype rather than module-level on purpose: float16 and float32 resolve through Stream ordering. Also right, and a real hole rather than a style point: Re-checked that the test still does its job: reverting the |
|
@cudnn-ci-bot frost |
|
cuDNN CI bot commands
Only allowlisted maintainers can use |
3 participants
Before submitting
pre-commit runand committed any formatting changes.cat-bug,mod-frost,mod-cutedsl,orig-nv-eng.Affected area
FE OSS kernels or CuTeDSL
Summary
buffers.probe()turns a caller's buffer into(pointer, shape, strides, dtype, device)for a kernel launch. It reads metadata and does nothing else — so it never needed synchronisation. Its DLPack fallback called__dlpack__()with the default stream argument, which torch reads as prepare this tensor for a consumer on another stream and answers withrecord_stream. That is illegal inside a CUDA graph capture.Why it looked like a GDN/KDA bug
probe()prefers__cuda_array_interface__, whosetypestris numpy-flavoured.float16is"<f2"and resolves there; bfloat16 has no numpy spelling, so torch reports raw"<V2", the lookup misses, and it falls through to DLPack.So the failure needed three things at once — bf16 buffers, a capture, and
probe()on the execute path — and only the linear-attention suites have all three:every one reporting
operation failed due to a previous error during capture— an error naming neither this file nor this function. SDPA and GEMM use fp16/fp8 and take the CAI path; a bf16 SDPA capture would have hit the same thing.Minimal reproduction, no cudnn involved:
Why
probe()reads metadata and launches nothing, so it has no business asking a producer to synchronise. Passing the default stream argument made torch do bookkeeping on its behalf — correct for a cross-stream handoff, illegal inside a capture, and unnecessary either way.Related issues
Found while working on #502 (python engine dispatch cleanup); independent of it and based on current
develop.API and compatibility impact
None.
buffers.probe()is internal and its return value is unchanged.stream=-1is the DLPack-specified way to request no synchronisation; theTypeErrorfallback keeps producers whose__dlpack__predates the argument working.Testing
test/python/test_buffer_probe.pypins it onprobe()directly rather than leaving the linear-attention suites to notice again — they do notice, but a hundred lines into unrelated setup and pointing at torch's stream bookkeeping. Reverting the fix turnstest_probe_is_capture_safe[bfloat16]red; that was checked.On sm100:
Not addressed here
probe()has no mapping for fp8 at all —TypeError: unsupported buffer dtype (code=10, bits=8, lanes=1)— so fp8 is absent from the test. Separate gap, and I have no fp8-under-capture case to verify a fix against.Worth knowing for anyone weighing the two protocols: CAI is not strictly cheaper-and-better. It cannot express bf16 or fp8 (fp8 raises outright), it is CUDA-only (CPU tensors and JAX arrays expose only DLPack), it carries no ownership or lifetime semantics, and torch reports version 2, which has no
streamfield at all. Its capture-safety here comes precisely from knowing nothing about streams. The existing CAI-first-then-DLPack order is right; only the DLPack call was under-specified./cc @Anerudhan
Summary by CodeRabbit
Bug Fixes
Tests