frost(sdpa): bind THD host prep to the launch stream; plan-time-only THD compile keys (#552) - #543
Conversation
vedaanta
added
cat-bug
|
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:
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 (10)
🚧 Files skipped from review as they are similar to previous changes (9)
📝 WalkthroughWalkthroughSM100 and SM120 THD paths now compile from plan-time keys with runtime packed extents. Execution reuses compiled artifacts, reads launch values at runtime, and performs metadata and dummy work on the launch stream. Tests verify reuse across packed totals. ChangesDynamic THD compilation and execution
Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: ⚪ Minimal · up to The PR changes THD stream ordering, removes an unused fill launch, and makes compilation plan-time-only; no actionable merge-blocking risk remains beyond normal checks and review. Sequence Diagram(s)sequenceDiagram
participant Caller
participant LaunchStream
participant THDPacker
participant CompiledArtifact
Caller->>LaunchStream: resolve execution stream
Caller->>THDPacker: pack runtime THD metadata
THDPacker->>LaunchStream: allocate and upload metadata
Caller->>CompiledArtifact: rebind packed extents and strides
CompiledArtifact-->>Caller: return attention outputs
🚥 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: 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 `@python/cudnn/sdpa/fwd/api_dsl.py`:
- Line 2088: Move the FP8 THD `seq_q_lens` dummy initialization out of the
shared setup near `_thd_pack` and create `seq_q_dummy` only in the dense branch.
Ensure the THD path passes no preallocated query-length dummy so `_thd_pack`
allocates and initializes it inside `_torch_stream_context` on the launch
stream, while preserving the cached tensor return behavior at the corresponding
FP8 THD handling. Add a regression test covering different handle and PyTorch
current streams.
🪄 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: e88dbb84-f14e-4add-97bf-95fb52ae4dd2
📒 Files selected for processing (1)
python/cudnn/sdpa/fwd/api_dsl.py
b012c03 to
37d0ebc
Compare
vedaanta
changed the title
There was a problem hiding this comment.
🧹 Nitpick comments (4)
python/cudnn/sdpa/fwd/api_dsl.py (1)
1250-1250: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename the ambiguous loop variable
l.Ruff E741 flags
lon this line. Uselength, which matches the naming in the new regression tests.♻️ Proposed rename
- units = qh * sum((l + cga_tile_m - 1) // cga_tile_m for l in slq_host) + units = qh * sum((length + cga_tile_m - 1) // cga_tile_m for length in slq_host)🤖 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/sdpa/fwd/api_dsl.py` at line 1250, Rename the generator variable `l` in the `units` calculation to `length`, and update its references within the expression so behavior remains unchanged and Ruff E741 is resolved.Source: Linters/SAST tools
python/cudnn/AGENTS.md (1)
91-95: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRecord the bounded all-KV-zero exception under Rule 3.
Rule 3 states the execute path's cached call "must be a guaranteed hit". The all-KV-zero clamp in
api_dsl.py(lines 1291-1300 and 2441-2449) swaps in packed batch-1 K/V views and mints a second cache entry on its first hit. The entry count is bounded at one extra artifact, so the per-step pathology Rule 3 targets does not apply. Record the exception so a later audit does not read the clamp as a rule violation.📝 Proposed addition
- **Compile at plan time, re-bind at execute.** With a plan-time-only key there is no reason to defer: `compile()` builds the artifact once and the execute path's cached call must be a guaranteed hit. Guard it with a cache-miss regression test (see `test_dsl_sm100_thd_compile_key_plan_time_only`), not by inspection. + One bounded exception is allowed: the all-KV-zero clamp binds packed + batch-1 K/V views and mints a single additional entry on first hit. A + runtime-keyed entry is acceptable only when the number of distinct keys it + can produce is bounded by a small constant.🤖 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/AGENTS.md` around lines 91 - 95, Update the Rule 3 documentation in AGENTS.md to explicitly record the bounded all-KV-zero exception: the clamp in api_dsl.py may create one additional cached artifact when substituting packed batch-1 K/V views, but remains bounded rather than causing per-step cache misses. Keep the existing guaranteed-hit rule and regression-test requirement unchanged.python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py (1)
1920-1922: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReplace the ambiguous Unicode
∪in the docstring.Ruff RUF002 flags
∪(UNION) as confusable withU. Use an ASCII separator instead.♻️ Proposed docstring fix
- The Q∪V∪O SMEM alias slabs keep their full compile-time extents — only + The Q/V/O SMEM alias slabs keep their full compile-time extents — only the GMEM descriptor extents change. d * BPE must be a 16-byte multiple (TMA global-stride rule -> d % 8).🤖 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/sdpa/fwd/kernels/prefill_d512_f16_sm100.py` around lines 1920 - 1922, Update the docstring near the Q/V/O SMEM alias slab description to replace the Unicode union character “∪” with an unambiguous ASCII separator, preserving the existing meaning and surrounding text.Source: Linters/SAST tools
test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py (1)
1009-1013: 🩺 Stability & Availability | 🔵 TrivialConsider adding coverage for a launch stream that differs from torch's current stream.
The PR states no in-tree test covers this case, and Rule 4 in
python/cudnn/AGENTS.mdnotes the race hides behind the common path. A test that creates a sidetorch.cuda.Stream, passes its handle ascurrent_stream, and validates the THD output would exercise the ordering this PR adds. It would also catch the FP8 THD dummy-allocation defect flagged onpython/cudnn/sdpa/fwd/api_dsl.py.Do you want me to draft this test?
🤖 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/sdpa/frost/test_sdpa_fwd_dsl_sm120.py` around lines 1009 - 1013, Extend the THD forward test around _run_and_check to create a separate torch.cuda.Stream, pass its handle as current_stream, and validate the resulting output against the existing reference path. Ensure the side-stream case exercises the added ordering and retains the compile-cache assertions, including coverage for FP8 dummy allocation.
🤖 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.
Nitpick comments:
In `@python/cudnn/AGENTS.md`:
- Around line 91-95: Update the Rule 3 documentation in AGENTS.md to explicitly
record the bounded all-KV-zero exception: the clamp in api_dsl.py may create one
additional cached artifact when substituting packed batch-1 K/V views, but
remains bounded rather than causing per-step cache misses. Keep the existing
guaranteed-hit rule and regression-test requirement unchanged.
In `@python/cudnn/sdpa/fwd/api_dsl.py`:
- Line 1250: Rename the generator variable `l` in the `units` calculation to
`length`, and update its references within the expression so behavior remains
unchanged and Ruff E741 is resolved.
In `@python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py`:
- Around line 1920-1922: Update the docstring near the Q/V/O SMEM alias slab
description to replace the Unicode union character “∪” with an unambiguous ASCII
separator, preserving the existing meaning and surrounding text.
In `@test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py`:
- Around line 1009-1013: Extend the THD forward test around _run_and_check to
create a separate torch.cuda.Stream, pass its handle as current_stream, and
validate the resulting output against the existing reference path. Ensure the
side-stream case exercises the added ordering and retains the compile-cache
assertions, including coverage for FP8 dummy allocation.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 749d76f0-4bd7-46c3-984b-eef8ae7bac70
📒 Files selected for processing (10)
python/cudnn/AGENTS.mdpython/cudnn/sdpa/fwd/api_dsl.pypython/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.pypython/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.pytest/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.pytest/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py
37d0ebc to
e15d077
Compare
|
Numbering collision to resolve — #570 is now in
It lands at the same anchor this branch uses (immediately after Rule 2's "acceptance is a promise about the execute path" bullet), so expect a textual conflict there, and after resolving it the file would carry two Rule 3s. Suggested resolution — this branch's rules shift by one:
Three notes while you are in there, none of them blocking:
note to self: claude::774e8e99-23ad-4a94-be0d-53ed5ee4def9 — "审计前端Python API和Frost引擎代码复杂度" · cwd /home/scratch.yanxu_libs/cudnn_frontend · workspace /home/scratch.yanxu_gpu/fe_pr1 |
e15d077 to
3dcbda4
Compare
The THD execute paths do torch host work before the kernel launch — the [seq_kv | cu_q | cu_k] metadata allocation + D2H length reads + one-shot H2D upload (SM100 and the shared SM120 _thd_pack), the per-sequence O-descriptor buffer (SM100), the dummy-sink buffer (SM100), and the cached seq_q dummy's first-use allocation (SM120). These enqueued on torch's CURRENT stream while the kernel launches on the stream carried by the execute-time handle (ExecutionContext.stream): when the two differ, the prep and the kernel race. Run the prep inside _torch_stream_context (the same helper the fp8/mxfp8 amax paths already use), and resolve the launch stream BEFORE _thd_pack in both SM120 callers. Allocations happen inside the context too, so caching-allocator blocks are stream-tagged to the stream that uses them. Pre-existing since the FROST engines landed (NVIDIA#476); split out of the NVIDIA#526 review round to keep that PR scoped to native THD stride support. Only direct graph-API users with an explicit handle stream are affected — the PyTorch integration launches on torch's current stream. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…tor scratch (SM100) A per-execute fill kernel on the THD execute hot path initialized the per-sequence O TMA-descriptor buffer, whose contents provably do not matter: the kernel's builder pass copies every qword of each sequence's slot from the base descriptor (then patches address/extent) before the fence and before any consumer read — stale workspace bytes never survive to a read. The +16-qword tail is never read at all. The fill dates to the original FROST landing (NVIDIA#476) as belt-and-braces. Rule 1: no adapter-side fills on the execute hot path. (The matching dummy-sinks fill removal is split into its own PR.) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…compile keys The THD execute paths keyed the per-shape kernel compile on the packed token totals (sq=t_q, skv=t_kv, and max_sq on SM120). Under continuous batching the totals change every step, so the lru_cache degenerated into a fresh multi-second cute.compile per execute (issue NVIDIA#552's worst leg). - Kernel modules (SM100 d128/d192_d128/d256/d512 f16, SM120 f16/fp8): under THD the fake tensors' token extents are cute.sym_int symbols (one per ragged group — Q/O/LSE share t_q, K/V share t_kv) and the batch stride is rebuilt symbolically (the real view's batch stride is t * token_stride, a runtime value that never steps at batch extent 1). sq/skv are ignored under THD; SM100's _host reads the runtime totals from the dynamic tensor shapes. SM120's max_sq moves from a compile parameter to a runtime __call__ argument that sizes the per-sequence grid; trace-time shape checks compare only statically-known modes. - Adapter: the THD compile key is now derivable from the graph declaration alone, so compile() builds the artifact at PLAN time (the "thd-deferred" sentinel remains only for the unwired SM100 fp8 THD) and the execute paths' lru-cached compile calls are guaranteed hits; a shared _thd_compile_kwargs() keeps the two call sites identical. The all-KV-zero clamp's swapped K/V strides mint their own entry. - The D2H .tolist() round-trip still feeds the metadata upload, the ragged views' extents and the exact grid — removing it (and the CUDA- graph capture blocker) needs the plan-time-max grid + device cu_seqlens redesign tracked in NVIDIA#552. - New regression tests (SM100 + SM120) prove one compiled artifact serves different packed totals, checking numerics per total and asserting zero cache misses across executes. Verified on SM100 (B200-class): 487 passed / 4 skipped across the f16 dense+THD flavors, fp8, mxfp8, graph-level THD and sdpa op suites. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
3dcbda4 to
ac022a9
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-543-ac022a9 |
… 5 (launch-stream-ordered execute) Renumbered after NVIDIA#570 landed Rule 3 (no D2H reads on execute); these two complement it. Rule 4 codifies issue NVIDIA#552's compile-key lesson: never key a kernel compile on runtime data values — runtime extents compile dynamic (cute.sym_int), runtime launch scalars are call arguments, derived values (batch strides computed from totals) count as leaks, and with a plan-time-only key the compile belongs at plan time with a cache-miss regression test guarding the execute path. Rule 3 bans the read that feeds such a key; Rule 4 bans the key itself. The SM80 _compile_cached (NVIDIA#493) is flagged as the known open cleanup. Rule 3's THD known-violation entry is updated: the compile-side half is done (dynamic token extents), so t_q/t_kv now reach the host only for the metadata upload, ragged view extents and the launch grid. Rule 5 codifies this PR's stream-binding fix: every torch operation on the execute path (H2D uploads, buffer resets, allocator calls, post-kernel consumers) is ordered on the launch stream via _torch_stream_context, never implicitly on torch's current stream. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
ac022a9 to
accec9a
Compare
4 participants
Three fixes to the FROST THD (varlen) execute hot path, one commit each, plus the AGENTS.md rules that keep them fixed.
1. THD host prep raced the kernel on a non-default handle stream
execute()launches the kernel on the stream carried by the execute-time handle (ExecutionContext.stream). But the THD host prep before that launch — the[seq_kv | cu_q(B+1) | cu_k(B+1)]metadata build and its one-shot H2D upload, the O-descriptor and dummy-sink scratch (SM100), the cached seq-q dummy's first-use allocation (SM120) — enqueued on torch's current stream. When the two streams differ, the prep and the kernel that consumes it race. Pre-existing since the FROST engines landed (#476).Fix: resolve the launch stream first and run every piece of THD host prep inside
_torch_stream_context— the same helper the fp8/mxfp8 amax paths already use. Allocations happen inside the context too, so caching-allocator blocks are stream-tagged to the stream that actually uses them. Exposure is limited to direct graph-API users with an explicit handle stream (the PyTorch integration runs on torch's current stream, where the context is a no-op-equivalent) — which is exactly why no test had caught it.2. A dead per-execute fill kernel (SM100 THD)
Every SM100 THD execute zero-filled the per-sequence O-descriptor scratch. The kernel's builder pass copies every qword of each sequence's slot from the base descriptor (then patches address/extent) before the fence and before any consumer read, so stale bytes never survive; the
+16-qword tail is never read at all. The fill was #476 belt-and-braces and is a wasted kernel launch per execute (AGENTS.md Rule 1).Fix: allocate without the fill. The suites pass with deliberately uninitialized scratch, confirming nothing reads it. (The matching dummy-sinks no-fill is split into #573.)
3. The THD compile was keyed on the packed token totals (issue #552)
The per-shape kernel compile was
lru_cache-keyed onsq=t_q, skv=t_kv(SM100/SM120) andmax_sq(SM120) — values read out of the seq-length tensors at execute. Under continuous batching the packed totals change every step, so the cache never hits and every step pays a fresh multi-secondcute.compile. This is the compile-side half of the plan Rule 3 (#570) records for the THD host round-trip: makeTa runtime value instead of a compile-time constant.Fix, kernel side: under THD the fake tensors' token extents are
cute.sym_intsymbols — one per ragged group (Q/O/LSE sharet_q; K/V sharet_kv) — so one compiled artifact re-binds any total. The batch stride, which ist * token_stride(a runtime value that never steps, batch extent is 1), is rebuilt symbolically from the dynamic extent rather than keyed. SM100's_hostreads the runtime totals from the dynamic tensor shapes; SM120'smax_sqbecomes a runtime__call__argument that sizes the per-sequence grid, and its trace-time shape checks now compare only statically-known modes.Fix, adapter side: with every key component plan-time-derivable, the THD compile moves to
compile()— plan time, like every dense specialization — and the execute paths' cached calls are guaranteed hits (shared_thd_compile_kwargs(); the all-KV-zero clamp's swapped K/V strides mint their own entry on first hit). New regression tests on both arches execute twice with different packed totals and assert zero new cache misses, with numerics checked per total.Still open in #552 (unchanged here): the D2H
.tolist()itself — it feeds the metadata upload, the ragged views' extents and the exact launch grid — and with it CUDA-graph capture. That needs the plan-time-max grid + in-kernel dead-tile exit + devicecu_seqlensread evaluated there. The SM80 engines'_compile_cached(#493) has the same total-keyed pattern and is flagged as follow-up cleanup in Rule 4.4. AGENTS.md Hard Rules 4 and 5
_torch_stream_context, producers and consumers both — never implicitly on torch's current stream.Verification (all on SM100, B200-class)
sdpa_fwd_exop tests, and #570'stest_sdpa_execute_is_async.py(D2H-ban + CUDA-graph capture/replay)test_dsl_sm{100,120}_thd_compile_key_plan_time_onlyThe SM120 kernel changes are structurally identical to the SM100 ones but need an SM120 machine to execute; the existing SM120 THD suites plus the new cache test cover them on SM120 CI.
Related
else-branch this PR stream-binds (trivial one-hunk rebase for whichever lands second).🤖 Generated with Claude Code
Summary by CodeRabbit
Performance
Bug Fixes
API Changes
thd_max_sqfrom initialization/configuration to the execution call for SM120 attention kernels.