Make the public gemm path work for the flavors it claims - #558
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:
📝 WalkthroughWalkthroughThis change adds asynchronous 32-bit device fills for FROST reduction outputs, passes runtime streams through GEMM launchers, treats zero workspace sizes as unknown, adds public execution tests, and updates runtime documentation. ChangesFROST execution
Estimated code review effort: 3 (Moderate) | ~20 minutes Mergeability Score: 🟡 Moderate · up to The PR enables reduction-output initialization through the public GEMM path, but validation of a later output can currently fail after an earlier caller buffer has already been modified, leaving partial state when execution aborts. This bounded correctness risk should be fixed or explicitly accepted before merge; a separate Python formatting follow-up remains outstanding. Sequence Diagram(s)sequenceDiagram
participant GEMMCompiler
participant fill_word_strided_async
participant CUDADriver
GEMMCompiler->>fill_word_strided_async: initialize reduction output on runtime stream
fill_word_strided_async->>CUDADriver: issue asynchronous 32-bit memset
CUDADriver-->>fill_word_strided_async: return CUDA status
fill_word_strided_async-->>GEMMCompiler: complete or raise RuntimeError
Possibly related PRs
Suggested labels: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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: 3
🤖 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/frost/workspace.py`:
- Around line 128-132: Update the Workspace handling around take() and
remaining() so a workspace with unknown capacity (workspace_bytes == 0) cannot
produce an invalid negative-sized DeviceView: either reject remaining()
explicitly or represent the remaining tail as unknown. Preserve valid behavior
for known capacities, and add coverage for a raw workspace address followed by
take() and remaining().
In `@python/cudnn/gemm/frost/compiler.py`:
- Around line 1775-1779: Update the reduction-output initialization in the
visible fill branch to honor the supplied stream: replace the direct torch fill
path with stream-aware fill_f32_async, or execute fill_ under a
torch.cuda.ExternalStream backed by ctx.stream. Preserve the existing fallback
for non-torch tensors and add a regression test covering differing current and
supplied CUDA streams.
In `@test/python/gemm/frost/test_public_execute_flavors.py`:
- Around line 62-65: Update the test helper or relevant tests around _run and
test_bare_address_operands to execute a graph requiring FROST workspace with
ws.data_ptr() rather than always passing the tensor. Add a passing bare-address
case that exercises the unknown-capacity path in Workspace.over and
WorkspaceCarve.carve, while retaining a known undersized workspace tensor case
that verifies bounds validation still raises.
🪄 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: 038ac86c-9595-45cf-92bc-a05a3a9f54d4
📒 Files selected for processing (9)
python/cudnn/engines/base.pypython/cudnn/frost/buffers.pypython/cudnn/frost/workspace.pypython/cudnn/gemm/frost/compiler.pypython/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.pypython/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.pypython/cudnn/linear_attention/frost/kernel/kda_prefill_f16.pypython/pygraph/variant_pack.cpptest/python/gemm/frost/test_public_execute_flavors.py
💤 Files with no reviewable changes (1)
- python/cudnn/engines/base.py
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-558-722770b |
|
Pushed
Landed: the driver where it is right, which is every contiguous buffer, and that also puts the seed on the stream the kernel runs on. A padded output keeps the torch path it already had. A padded output arriving as a slot — the public path, where this was broken and where nothing could seed it — is refused, with the measurement in a TODO pointing at the fill kernel that would close it. Verified: 4494 passed, 1 xfailed across gemm/, linear_attention/ and the variant-pack suites. The general point, since it cost three passes: replacing a borrowed operation means adopting a contract you did not write. |
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-558-897152c |
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 `@test/python/gemm/frost/test_public_execute_flavors.py`:
- Around line 190-210: Assign an L0–L4 test-level marker to each new test:
test/python/gemm/frost/test_public_execute_flavors.py:190-210
test_bare_address_workspace, :214-233 test_undersized_workspace_still_rejected,
and :236-254 test_unknown_size_workspace_refuses_to_measure_its_tail. Use the
project’s established Python test-level marker convention for all three tests.
- Line 231: Update the pytest.raises call in the relevant test to use a raw
regex string for the match pattern, preserving the existing regular expression
while satisfying Ruff RUF043.
🪄 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: f5d319b3-30dc-4443-8bf0-1fc59f1cee36
📒 Files selected for processing (4)
python/cudnn/frost/buffers.pypython/cudnn/frost/workspace.pypython/cudnn/gemm/frost/compiler.pytest/python/gemm/frost/test_public_execute_flavors.py
🚧 Files skipped from review as they are similar to previous changes (3)
- python/cudnn/frost/buffers.py
- python/cudnn/frost/workspace.py
- python/cudnn/gemm/frost/compiler.py
| ``tensormap_workspace``. Compiled + launched separately from the main | ||
| kernel and cached by input identity in the host bridge, so the builder | ||
| launches do not recur in steady-state replay. | ||
| ``tensormap_workspace``.Launched on every execute: the descriptors fold cu_seqlens contents into |
There was a problem hiding this comment.
Nit: missing space
| ``tensormap_workspace``.Launched on every execute: the descriptors fold cu_seqlens contents into | |
| ``tensormap_workspace``. Launched on every execute: the descriptors fold cu_seqlens contents into |
There was a problem hiding this comment.
Fixed — thank you. There was a third occurrence you did not hit, gdn2_prefill_f16.py, so all three are corrected rather than just the two flagged.
| o, h) into ``tensormap_workspace``. Compiled + launched separately from | ||
| the main kernel and cached by input identity in the host bridge, so the | ||
| builder launches do not recur in steady-state replay. Each descriptor | ||
| o, h) into ``tensormap_workspace``.Launched on every execute: the descriptors fold cu_seqlens contents into |
There was a problem hiding this comment.
Nit: missing space
| o, h) into ``tensormap_workspace``.Launched on every execute: the descriptors fold cu_seqlens contents into | |
| o, h) into ``tensormap_workspace``. Launched on every execute: the descriptors fold cu_seqlens contents into |
There was a problem hiding this comment.
Fixed — thank you. There was a third occurrence you did not hit, gdn2_prefill_f16.py, so all three are corrected rather than just the two flagged.
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 `@python/cudnn/gemm/frost/compiler.py`:
- Line 1781: Update the padded-output handling around the reduction output fill
to remove the tensor.fill_() fallback and use buffers.fill_word_async for every
variant-pack buffer type, preserving the reduction dtype/value and passing the
target stream so initialization is stream-ordered.
In `@test/python/gemm/frost/test_public_execute_flavors.py`:
- Around line 140-141: Add the appropriate L0–L4 test-level marker to
test_int32_reduction_seed_is_packed_as_int32, while retaining the existing `@_GPU`
hardware gate.
🪄 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: 140e34bb-d983-457d-8f43-79d56b7cfbe8
📒 Files selected for processing (3)
python/cudnn/frost/buffers.pypython/cudnn/gemm/frost/compiler.pytest/python/gemm/frost/test_public_execute_flavors.py
NVIDIA#547 normalized the variant pack, which changed what an engine is handed: buffers arrive as slots that describe memory rather than as tensors. Two places in the gemm engine were writing THROUGH the caller's buffer with torch methods, which worked only while the buffer happened to be a torch tensor: - a reduction output is seeded with its identity before the kernel runs, via tensor.fill_() - a norm2 reduction is finalized with tensor.sqrt_() The engine owns the first one now. Every seed a reduction uses (0, 1, +-inf) is a 32-bit pattern, so buffers.fill_f32_async drives cuMemsetD32Async and no kernel is needed. The second is unreachable: a norm2 reduction is refused while the graph is lowered, so no plan exists to execute -- recorded as a test rather than left as a live-looking path. A bare device address as the WORKSPACE measured 0 bytes, and Workspace.over read that as "empty" and refused any engine that needs scratch. It means "the pack could not measure it": a raw pointer carries no size and the backend takes one without checking, so refusing here made the same call depend on which plan ran. Zero now skips the size check, here and in the C carve's bounds check. The reason none of this was caught: no test drove a non-trivial gemm flavor through graph.execute(). The direct-call tests construct a fusion chain and invoke the compiled object with torch tensors, which skips operand binding, the pack, and the conversion execute() performs -- exactly the part that changed. test_public_execute_flavors.py covers plain matmul, epilogue fusion, the four reduction modes, norm2's refusal, and the bare-address operand form through the entry point a caller actually has. That last one is xfail: frost reads its extents by axis position, and a bare address describes the operand the way the GRAPH declares it (a matmul's B is [batch, K, N]) rather than the way a caller's buffer reports it. It was broken before this too -- the geometry-less Tensor made it an IndexError instead. The fix is the engine recording which axis is M/N/K at build. Also: two unused imports in engines/base.py that broke the lazy frost boundary, and three kernel docstrings claiming the descriptor builders do not recur.
Three passes over the same fix, each caught by measuring rather than by the
suite going green:
1. tensor.fill_() does not exist on a slot -> use cuMemsetD32Async.
2. That lost the STRIDE. torch's fill_ sets every element; a memset writes a
contiguous byte range, and the two agree only for a dense buffer. Six
strided reduction tests went NaN -- the elements past the first run were
never written.
3. One memset per contiguous run is correct and 572 us for a per-row scalar
output, against 3.5 for the single kernel torch launches. Correct is not
mergeable; nothing in the suite would have flagged it.
So: the driver where it is right, which is every contiguous buffer -- and that
also puts the seed on the stream the kernel will run on, where tensor.fill_()
queues on torch's current stream and is the same stream only by luck.
A padded output keeps the torch path it already had, and a padded output
arriving as a slot -- the public path, which is where this was broken and
where nothing could seed it -- is refused with the measurement in a TODO. The
fix is a fill kernel; this is the last place the engine writes through the
caller's buffer.
Also from review: remaining() refuses rather than returning a negative extent
when the workspace size is unknown, and three cases cover the unknown-capacity
path end to end (bare-address workspace, undersized-but-known, and the tail
refusal).
…stream A memset moves bits, not numbers, so the identity has to be packed as the dtype the kernel reads it back as. int32's identities are the ends of its range and are exactly where that bites: -2**31 packed as float is 0xcf000000, so an int32 MAX reduction returned -822083584 for every input below it. fill_f32_async becomes fill_word_async over a 32-bit pattern, with init_word turning a value into one. The four MoE launchers seeded on the null stream rather than the execute one -- they were the call sites that had no stream to pass before this path moved to the driver, and passing None was not the same thing afterwards. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
dee80c0 to
605e443
Compare
|
Force-pushed The one file this PR shares with those merges is
|
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)
python/cudnn/gemm/frost/compiler.py (1)
2334-2335: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winFormat these validation errors with Black.
Lines 2335 and 2346 exceed the 160-character limit.
Proposed formatting
- raise ValueError(f"TileConfig {config.name!r} cannot use M-major A: " f"per-MMA M={slice_m} is not divisible by " f"swizzle group {mn_group_elems}") + raise ValueError( + f"TileConfig {config.name!r} cannot use M-major A: " + f"per-MMA M={slice_m} is not divisible by " + f"swizzle group {mn_group_elems}" + ) @@ - raise ValueError( - f"TileConfig {config.name!r} cannot use N-major B: " f"per-MMA per-CTA SMEM N={slice_n} is not divisible by " f"swizzle group {mn_group_elems}" - ) + raise ValueError( + f"TileConfig {config.name!r} cannot use N-major B: " + f"per-MMA per-CTA SMEM N={slice_n} is not divisible by " + f"swizzle group {mn_group_elems}" + )As per coding guidelines, format Python code with Black and a maximum line length of 160 characters.
Also applies to: 2344-2346
🤖 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/gemm/frost/compiler.py` around lines 2334 - 2335, Reformat the long ValueError validation messages in the relevant TileConfig validation block, including the checks around slice_m and the lines 2344–2346, using Black-compatible wrapping while preserving the existing error text and behavior. Keep the maximum line length at 160 characters.Source: Coding guidelines
🤖 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/gemm/frost/compiler.py`:
- Around line 2334-2335: Reformat the long ValueError validation messages in the
relevant TileConfig validation block, including the checks around slice_m and
the lines 2344–2346, using Black-compatible wrapping while preserving the
existing error text and behavior. Keep the maximum line length at 160
characters.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: dca8b720-549c-454f-9b77-b1bde3450144
📒 Files selected for processing (1)
python/cudnn/gemm/frost/compiler.py
The contiguous case already went through cuMemsetD32Async on the execute-time stream. The padded one fell back to tensor.fill_(), which queues on torch's CURRENT stream -- the same stream only by luck -- and exists at all only while the caller happened to pass a torch tensor. That contradicted what this branch claims to do, so it is gone. strided_fill_plan collapses the layout to the runs a memset can cover and returns the 2D memsets that cover it exactly once. cuMemsetD2D32Async takes a pitch, so a per-row scalar tap is ONE call rather than one per row -- that reading, 572 us at one memset per row, is why the fallback was there. What remains is one call per point of whatever axis is left outside the 2D region, which for a rank-3 output is the batch and is usually 1. The plan is returned before anything is written, and is None for a layout that would write an element twice: a stride of 0 over a real extent, or an outer stride that does not clear the axis below it. Checking only the innermost pair (pitch >= width) is not enough -- shape (2, 2) stride (2, 2) has width 1 and lands both axes on element 2. Also in this commit: a missing space in three linear-attention kernel docstrings (review caught two of the three), and a raw-string pytest.raises pattern. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Pushed Fixed
Not changed, with reasoning
Test: Also rebased onto |
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)
python/cudnn/gemm/frost/compiler.py (1)
1813-1829: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winValidate all reduction outputs before writing any output.
Line 1829 plans and applies each strided fill during the loop. If an earlier output is valid and a later output has an unsupported or overlapping layout, the earlier caller buffer is modified before execution raises.
Build and validate all fill descriptors first. Apply the descriptors only after validation succeeds for every reduction output. Add a multi-reduction regression test with a valid first output and a rejected second output.
🤖 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/gemm/frost/compiler.py` around lines 1813 - 1829, The reduction-output loop must not mutate caller buffers before every layout is validated. In the output initialization flow around the reduction descriptor construction, first collect and validate all fill descriptors for every reduction output, rejecting unsupported or overlapping layouts; only after validation succeeds for the entire set should you invoke the contiguous or strided fill helpers. Add a regression test covering multiple reductions where the first output is valid and the second is rejected, verifying no output was modified.
🤖 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/gemm/frost/compiler.py`:
- Around line 1813-1829: The reduction-output loop must not mutate caller
buffers before every layout is validated. In the output initialization flow
around the reduction descriptor construction, first collect and validate all
fill descriptors for every reduction output, rejecting unsupported or
overlapping layouts; only after validation succeeds for the entire set should
you invoke the contiguous or strided fill helpers. Add a regression test
covering multiple reductions where the first output is valid and the second is
rejected, verifying no output was modified.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6207cf32-e1f0-4d93-9ec2-4822efc27c68
📒 Files selected for processing (6)
python/cudnn/frost/buffers.pypython/cudnn/gemm/frost/compiler.pypython/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.pypython/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.pypython/cudnn/linear_attention/frost/kernel/kda_prefill_f16.pytest/python/gemm/frost/test_public_execute_flavors.py
🚧 Files skipped from review as they are similar to previous changes (3)
- python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py
- python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py
- python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py
YangXu1990uiuc
added
cat-bug
4 participants
Follow-up to #547, which landed before this was ready.
#547 changed what an engine is handed: buffers arrive as slots that describe memory rather than as tensors. Two places in the gemm engine write THROUGH the caller's buffer with torch methods, which worked only while the buffer happened to be a torch tensor.
A reduction output is broken on
developright now. The epilogue seeds it withtensor.fill_()before the kernel runs; a slot has nofill_. The engine owns that operation now -- every seed a reduction uses (0, 1, ±inf) is a 32-bit pattern, socuMemsetD32Asynccovers them without a kernel.norm2'ssqrt_finalize is unreachable: the backend refuses that reduction while the graph is lowered, so no plan exists to execute. Recorded as a test rather than left looking live.A bare device address as the workspace measured 0 bytes and
Workspace.overread that as "empty", refusing any engine that needs scratch. Zero means "the pack could not measure it" -- a raw pointer carries no size, and the backend takes one without checking.Why the suite did not catch it
No test drove a non-trivial gemm flavor through
graph.execute(). The tests that cover reductions build a fusion chain and call the compiled object directly with torch tensors, which skips operand binding, the pack, and the conversionexecute()performs -- exactly the part #547 changed.test_public_execute_flavors.pycloses that: plain matmul, epilogue fusion, the four reduction modes, norm2's refusal, and the bare-address operand form, all through the public entry point.The bare-address case is
xfail(strict): frost reads its extents by axis position, and a bare address describes the operand the way the GRAPH declares it (a matmul's B is[batch, K, N]) rather than the way a caller's buffer reports it. Broken before #547 too, as an IndexError. The fix is the engine recording which axis is M/N/K at build, which belongs with the executor rewrite.Verified on SM100: 4557 passed, 1 xfailed across gemm/, linear_attention/ and the variant-pack suites.
note to self: claude::774e8e99-23ad-4a94-be0d-53ed5ee4def9 — "cuDNN FE variant-pack normalization" · cwd /home/scratch.yanxu_libs/cudnn_frontend · workspace /home/scratch.yanxu_gpu/fe_pr1
Summary by CodeRabbit
New Features
Tests
Documentation