Lower non-MoE, workspace-free gemm flavors off one build-time recipe

[YangXu1990uiuc]

#558 has landed, so this is no longer stacked — the diff is this PR alone.

Rebased onto eb43b2c73. Clean: the three files both sides touched turned out to be disjoint regions, and #566's diff has no lines inside the execute path this PR rewrites. #566 is orthogonal to the recipe — every catalog entry keeps num_mma_m == 1, so select_config never picks the new geometries, and what the recipe reads off the compiled object (use_tma_store, vec_bytes_epi, the alignment requirements) follows automatically.

What

A compiled gemm's per-call path was re-deriving, on every execute, facts that a
runtime value cannot change: which axis of each operand carries M/N/K, each
major, the fp4 packing factor, every output's required alignment and its shape
rule, which outputs are reductions, and what order the kernel takes its
parameters in. This reads them once into gemm/frost/recipe.py when the kernel
compiles, and _lower captures that table into the closure that launches —
which by the end of this PR is the only thing that launches at all.

256×256×128 bf16, host enqueue, min over 25 reps of a 64-call burst from a
drained queue:

flavor	before	now	device ops
plain	17.73	19.76	1 kernel
epilogue (relu)	17.69	19.69	1 kernel
aux (bias tensor)	38.70	22.26	1 kernel
2 dense outputs	43.89	22.65	1 kernel
reduction output	49.90	26.12	1 kernel + 1 driver memset
multi-gemm	42.58	22.54	1 kernel

The "before" column is this branch's own earlier state, where only the plain
single-GEMM shape lowered. Against develop every row starts at 42–50.
Reduction sits ~4 µs above the rest because it still pays one
cuMemsetD32Async per tap; that goes when the kernel seeds itself.

The field that made six flavors one path

What actually differed between plain, aux, multi-output, multi-GEMM and block
scale was which buffers the launch passes and in what order, so that is data
now: arg_plan, one (operand index, aux ref) per positional argument after
problem_size. Three smaller fields carry the rest — stride_ins (whose
permuted strides ride in problem_size), shared_layout (block-scale
multi-GEMM collapses its A operands to one stride triple and requires the others
to match), seeds (a reduction's identity, packed as its output dtype).

Before that, each flavor had its own launcher assembling its own argument tuple,
which is why the emitted path served exactly one of them.

Loop, not unrolled — and the measurement behind that

The first version of _lower hand-wrote the plain flavor's checks as a straight
line with the structure unrolled. Measured three ways on the same plan and
buffers in one process, all bit-exact:

flavor	interpreted	flat-table loop	hand-unrolled
plain	35.82	19.72	17.45
aux (bias)	38.43	22.37	—
2 dense outputs	44.32	22.63	—

The loop is worth 42–49%; unrolling adds 12% on top and costs one closure body
per operand shape. Six flavors would have been six bodies to keep in agreement,
so the unrolled body is deleted. Source codegen off this same table is how to
buy the 12% back later, for every flavor at once rather than for the one that
was worth hand-writing.

One executor, and a checker beside it

The first version of this kept an interpreter behind the lowered path: anything
it was unsure of went to launch, which ran the call anyway. That made "the
fast path refused a legal call" a performance bug at worst, with nothing
asserting otherwise — so the last commit closes it.

The deferral counters split the eight hand-over sites into two populations that
reading the code does not. Six mean the call is illegal (wrong major, a
TMA-misaligned contiguous extent, a misaligned output base, an SF blob that is
not a dense run) and only ever needed a message. Two meant the call was legal
and the loop could not serve it, and both are gone:

• An operand in the graph's axis order. cuDNN declares B [b, K, N] where
  the kernel reads (b, N, K), so a bare device address — and any buffer
  reporting the declared (dim, stride) — arrived the other way round and paid
  a whole interpreted pass for being legal. It is re-labelled with one permute
  whose shape the recipe already carries, before any check runs, so everything
  below still reads one order. Plain gemm 19.95 → 19.76 µs: free on a call that
  does not use it.
• A strided reduction seed. fill_word_strided_async collapses the layout
  to the runs a memset can cover and issues cuMemsetD2D32Async, which takes a
  pitch — one call per remaining outer point (the batch, usually 1) rather than
  one per row. It was the last place the engine wrote through a buffer it does
  not own: tensor.fill_() works only while the caller passes a torch tensor,
  and queues on torch's stream rather than the one the kernel will run on.

Build-time declines also kept launch alive as an executor — no tvm-ffi front
door, a norm2 output. Those are _check_executable now, called from
probe_supported, so the engine declines the graph and it goes to the
backend, which is where it would have gone had this engine never been asked.
_lower reads that same function rather than restating it.

A multi-GEMM with no dense output was declined too and is supported again: the
decline existed because two multi-GEMM launchers read the batch off cs[0],
those launchers are deleted, and the recipe already carries
max(a_batch, b_batch) for the case. It is a flavor in the test table now.

Then launch loses its second half. explain() runs the same
recipe.problem / check_shapes / check_alignment and raises;
_call_positional, _call_multi_gemm, _call_block_scale_multi_gemm and
launch's assembly are deleted — 267 lines. MoE is untouched: it is ≥ 2
launches with its own workspace and no recipe.

The invariant tightens from "the fast path may accept a subset" to:

the set of calls the launch path refuses should EQUAL the set of illegal calls.

compiled.deferrals gets one direction of that and only one: a legal call of
every flavor must leave it empty, asserted per flavor and for the bare-address
form, so nothing legal is refused. It says nothing about the other direction,
because a call the fast path ACCEPTS never reaches the counter — that is covered
case by case in test_the_launch_path_accepts_exactly_the_legal_calls and by
the backend differentials. Closing it properly means generating the fused guards
and the readable diagnostics from one ordered list of checks; they are still
written twice by hand, and that is the open item.

Execution is no longer duplicated but the RULES still are, fused in lowered's
guards and readable in explain. If they drift, a call is refused and the
checker finds nothing wrong — so explain falling off the end raises a distinct
RuntimeError naming which guard fired. That guard has already earned its
place: it is what surfaced the shared-operand bug below, which would otherwise
have been a silent wrong answer.

Four defects review found in this, three of them one shape — a per-ROLE fact
treated as a per-SLOT one, or a write issued before every rule had its say:

• matmul(A, A) binds ONE slot as both operands (recipe.build keys the
  operand map on id(tensor)). Re-labelling that slot for the B role also
  transposed what the A role read. lowered builds one view per role now
  and never rewrites a slot; shared_layout moved to role positions for the
  same reason.
• The reduction seed wrote before the launch could reject the call. A fp32
  tap on a fp16 buffer took a 32-bit word × numel and ran past the allocation —
  the front door refuses the dtype, but after the fill. Two taps could also
  leave the first seeded and the second refused. And pitch >= width is the
  innermost pair only: it accepts shape (2,2) stride (2,2), whose two axes
  land on the same element. So the fill is planned before it is issued, the
  element width is a rule of the call, and every axis must clear the span of the
  one below it.
• Two fast guards had no rule in the checker — block-scale shared_layout,
  and a rank-2 operand, which additionally made problem raise IndexError
  from an axis index. Both have one, and problem checks the rank first.
• A scale factor's rank was nobody's rule. build() puts scale factors in
  heads alongside the inputs, so they reach permute(1, 2, 0) in the launch
  argument list — and the blob rules (16-byte base, one dense run, the size the
  template re-synthesizes) all pass for a flat rank-1 blob. It raised from
  inside the body whose contract is that it does not raise, and the checker
  could not name a cause. Same shape as the rank hole above: a rule reading
  three named axes off a buffer whose rank nothing had established. Inputs and
  outputs were covered; scale factors were the head left over.

What a differential does not buy, and why the interpreted one is gone. It
catches divergence, never a misconception the two paths share — which is how the
axis-order bug below survived one, with both paths reading Operand.axes and
both agreeing and both wrong. Two readings of the same wrong plan agreeing
proves nothing. What replaces it is fast-vs-backend
(test_a_degenerate_extent_is_refused_or_matches_the_backend,
test_an_operand_reporting_the_declaration_agrees_with_the_backend) plus the
per-flavor torch references in test_public_execute_flavors.py, which run
through the public graph.execute(). When a check discriminates between two
encodings, test it at the shape where the encodings coincide.

Two earlier divergences are what that drift looks like, both found by review
rather than by test: the hand-written path dropped the operand batch check, and
pinned an fp4 output's last axis at N where the graph says N/2. Both are single
facts in the recipe now (Operand.batch, Output.rule).

What it fixes

• Bare device addresses now work (the test_bare_address_operands xfail is
  gone). cuDNN declares a matmul's B as [batch, K, N] while a caller allocates
  it (batch, N, K) — the same memory, two axis orders, and at N == K the same
  (shape, stride) tuple. The tie-break is the backend's own rule: the
  descriptor defines the tensor and the pack supplies only a pointer, so a
  buffer that reports the declaration is read as the declaration. Measured, that
  agrees with the backend to bf16 tolerance where the other reading differs by
  65. This also fixes a pre-existing divergence — with FROST enabled,
  test_every_variant_pack_form_still_works fails all four forms on develop.
• The TMA 16-byte modulus is derived once (contiguous_modulus), read by
  both the graph-time gate and the per-call one. It was the one number with a
  packing factor in it, and so the one most likely to drift.

Where the rest of the win came from

• _check_plan_device did its from cudnn.frost.device import current_device
  per call: 1.1 µs of the 1.7 that function took. Moved to module scope.
• tensor_alignment's layout half is memoized on (shape, stride, elem_bytes, cap) — a function of values, with nothing to invalidate. The pointer half is
  still computed per call.
• The engine hands the launcher a positional operand list instead of building
  {id(tensor): buffer} and looking each operand back up.

Naming

slot was doing two jobs — an int POSITION and the OBJECT at it — so neither
read. pack.slot(t) → pack.index_of(t), pack.views(...) →
pack.operands(...), and the C++ VariantPackSlot → OperandBuffer (with its
set_slot/read_slot/… accessors renamed to match).

What it does not do

• Declined at the support gate, each by a named rule: no tvm-ffi front door,
  a norm2 output. The graph goes to the backend rather than being compiled
  into a kernel only a second executor could run.
• norm2 is the one capability this removes, and it was direct-call only.
  Its taps land through cross-CTA atomics, so the square root cannot go in the
  epilogue — it has to run after every CTA has contributed, which is a second
  launch. Nothing public could reach it anyway: the backend refuses a norm2
  reduction descriptor at build_operation_graph, so no execute() ever gets a
  plan for one (that is test_norm2_reduction_is_refused_at_build, already
  here). The two test_reduction_mode_coverage cases that computed it by
  calling the compiled object with torch tensors are now one test asserting the
  decline and saying why.
• MoE keeps its own launchers. CompiledMoeGemm is a separate class with no
  recipe; it is ≥ 2 launches with its own workspace and a different problem,
  which is why _maybe_wrap_layout, _wrap_raw_tensor,
  _initialize_reduction_outputs and _finalize_reductions all survive.
• Batch is still pinned to the graph's declaration, so an override_shapes
  that changes it is refused. The kernel does not need that — matmul_a_batch /
  matmul_b_batch reach the template only inside a cutlass.const_expr(… == 1),
  so what is baked is the broadcast predicate and not the value. develop pins
  it the same way, so relaxing it is a widening, not a regression to fix here.

Test

test/python/gemm/frost on the tip: 5768 passed / 2860 skipped / 0 failed
(4212 before the rebase — #566 added its num_mma_m=2 configs to the test
tables).

pytest test/python -q -n 8 --ignore=test/python/test_mhas.py on SM100 (B200):
89 failed, 15829 passed, 7123 skipped, and zero failures anywhere under
test/python/gemm/ — the directory this PR changes.

Pin CUDA_VISIBLE_DEVICES to the Blackwell for the full run: unpinned, xdist
spreads across every GPU on the box and lands most of it on an L40S, which gives
238 failures by 68% and reads exactly like a large regression.

All 89 are pre-existing and outside this PR's surface: 87 are numerical
tolerance failures in the cuteDSL OSS kernels (fe_api/gemm/test_gemm_swiglu 64,
fe_api/sdpa/test_sdpa_bwd 12, fe_api/test_rubin_kernel_dispatch 9,
fe_api/dsa/test_DSA_indexer_top_k 2) and 2 are test_mxfp8_quant::test_te_parity
failing on a missing transformer_engine. Established rather than assumed: those
tests call GemmSwigluSm100.execute and friends directly, reference no
pygraph/variant_pack/_normalize, and an import trace over three of them
shows no cudnn.gemm.frost.* module is ever loaded.

test/python/test_mhas.py is excluded on the maintainer's instruction —
test_mhas_v2 is the SDPA suite that counts. Worth recording why: the old file
writes shared pytest config state mid-test (request.config.option.left_bound = None, ...dropout = None) which every later test in the same xdist worker reads
back as its own parameters, so its outcome depends on ordering and worker
assignment. Three runs of one unchanged commit gave 331 / 155 / 173 failures, all
from that file.

One tolerance is not bit-exact and the reason is measured: a reduction output is
not bit-reproducible against itself, because its taps land through cross-CTA
float atomics. The same path run six times spreads 0.0049 while the two paths
differ by 0.00024, so that comparison is at the noise floor rather than at 0.

---

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

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR replaces slot-oriented variant-pack APIs with operand-buffer APIs. It adds immutable GEMM recipes, cached lowered execution, stream-ordered reduction initialization, workspace unknown-size handling, alignment and layout validation, and broad FROST execution tests.

Changes

FROST execution and operand buffers

Layer / File(s)	Summary
Operand-buffer contract and native bindings python/pygraph/variant_pack.cpp, python/cudnn/engines/base.py, python/cudnn/_pygraph.py, python/cudnn/frost/buffers.py, python/cudnn/frost/workspace.py, python/cudnn/linear_attention/engine_utils.py	Variant-pack slot APIs become operand-position and OperandBuffer APIs. UID binding, DLPack exchange, workspace carving, and raw-address handling use the new terminology.
Recipe metadata and buffer initialization python/cudnn/gemm/frost/recipe.py, python/cudnn/frost/buffers.py, python/cudnn/gemm/frost/dtypes.py	GEMM recipes describe layouts, outputs, launch arguments, alignments, strides, reductions, and workspace data. Stream-ordered fills validate overlapping layouts.
Recipe-driven lowered execution python/cudnn/gemm/frost/compiler.py, python/cudnn/gemm/frost/engine.py, docs/python_graph_and_execution_backends.md, python/cudnn/frost/README.md, python/cudnn/linear_attention/frost/kernel/*	FROST builds recipes and lowered launchers during initialization. Runtime checks cover shapes, axes, strides, alignment, scale factors, shared layouts, and reduction seeds. Unsupported calls report diagnostics.
Execution integration and validation coverage test/python/gemm/frost/test_execute_recipe.py, test/python/gemm/frost/test_public_execute_flavors.py, test/python/gemm/frost/test_matmul_epilogue_fusion.py	Tests cover public execution, recipes, reductions, FP4, multiple GEMMs, raw addresses, workspace sizes, layout validation, and NORM2 rejection.

Estimated code review effort: 5 (Critical) | ~90 minutes

Mergeability Score: 🟡 Moderate · up to 1667f

The PR consolidates GEMM execution and validation across several flavors, but merge readiness is reduced by a rank-validation test that may pass without exercising the intended execution path and by a large parameter sweep placed in the fastest test tier; these should be corrected or explicitly accepted before merging.

Sequence Diagram(s)

sequenceDiagram
  participant graph_execute
  participant _FrostGemmPlan
  participant CompiledFusedGemm
  participant GemmRecipe
  graph_execute->>_FrostGemmPlan: execute variant pack
  _FrostGemmPlan->>CompiledFusedGemm: invoke lowered launcher
  CompiledFusedGemm->>GemmRecipe: validate runtime operands
  GemmRecipe-->>CompiledFusedGemm: return validation result
  CompiledFusedGemm-->>graph_execute: launch or report rejection

Loading

Possibly related PRs

• NVIDIA/cudnn-frontend#558: The current PR extends its FROST buffer filling, workspace handling, compiler execution, and public GEMM tests.

Suggested reviewers: yanqinz2

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 58.76% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the primary change: lowering non-MoE, workspace-free GEMM flavors from a shared build-time recipe.
Description check	✅ Passed	The description is comprehensive and explains the changes, rationale, compatibility impact, implementation details, and targeted and full-suite testing results.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[YangXu1990uiuc]

Thanks — all three taken. Pushed ea476d6c8.

1. The aux-rank bug is real, and wider than rank-1. It is not rank-1 specific: _reshape_aux_to_fake asked for the rank with getattr(t, "ndim", <the fake's rank>), and OperandBuffer carries no ndim (__len__ is the first extent, not the rank). So the default meant "already matches" and every aux arriving through graph.execute() skipped the reshape — the existing public aux test passed only because a rank-3 bias needs no reshape.

Same graph, same data, only the handed-over buffer's shape differs:

bias declared	handed over as	backend	frost
[1,1,N]	(1,1,N)	OK	OK, identical
[1,1,N]	(1,N)	OK	ValueError
[1,1,N]	(N,)	OK	ValueError

Fixed with len(t.shape) — every buffer answers that, and there is no default to take. OperandBuffer also grows an ndim, so the next getattr for one cannot repeat this. The regression test is a differential against the backend across ranks 3/2/1, because the rule being broken is the backend's own: the descriptor defines the tensor, the pack supplies only a pointer.

(The C++ diff is larger than the one line it should be: the longer ndim line pushes the .def(...) chain past the column limit, so clang-format re-broke auto operand_class = and re-indented all 20 chained calls. Mechanical, no semantic change.)

2. Block-scale evidence added. An nvfp4 lowered-vs-interpreted differential for one and two GEMMs — the two-GEMM case also asserting shared_layout, which only it populates — plus an nvfp4 matmul through the public graph.execute() that asserts _lowered is not None. Two GEMMs do not fit the auto-selected cta_n=256 in TMEM, so that case pins a geometry that does.

3. Full suite. Still running; I will post counts here. Worth recording that my first two attempts were invalidated by my own runner flags, not by code:

• export CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 → 331 failed. The frost suites opt in per-test; forcing it globally breaks test_native_backend_lowering (which asserts a graph lowers to the backend), MoE, block-scale, silu. test_dispatch.py::test_frost_opt_in_does_not_leak_out_of_the_frost_suites exists to catch exactly this and was itself in the failure list.
• -p no:randomly → 155 failed, all test_mhas. pytest-randomly reseeds before each test and these SDPA tests draw their config randomly, so disabling it changes which configs run. Proven on one id, same commit, same GPU: default → skipped, -p no:randomly → numeric failure at test_mhas.py:1527.

So the run in flight is plain pytest test/python -q -n 8. Targeted test_execute_recipe.py + test_public_execute_flavors.py: 66/66.

Title narrowed to "non-MoE, workspace-free" as suggested.

On the two-track structure and the lazy UID→index: agreed, both are intermediate. Not touching them here — the end state is the interpreter consuming the same tagged plan (or a codegen'd native PreparedLaunch) and then deleting _call_*, which is a bigger change than this PR should carry. Also agreed this does nothing for cold start: first execute is still ~733–933 µs and plan build ~1.1–1.4 s, both untouched.

Thanks for measuring FP4 block-scale (44.77 → 26.14) — that is the one flavor I had not timed, and for confirming GPU event time is unchanged, which is the check that says the win is host-side.

[YangXu1990uiuc]

Full-suite result, plus a correction to what I said above.

Result

pytest test/python -q -n 8 --ignore=test/python/test_mhas.py on SM100 (B200):
89 failed, 14257 passed, 6325 skipped — and zero failures anywhere under test/python/gemm/, the directory this PR changes. Targeted test_execute_recipe.py + test_public_execute_flavors.py: 66/66.

The 89 are pre-existing and outside this PR's surface — 87 numerical tolerance failures in the cuteDSL OSS kernels (test_gemm_swiglu 64, test_sdpa_bwd 12, test_rubin_kernel_dispatch 9, test_DSA_indexer_top_k 2), 2 from a missing transformer_engine. Shown rather than assumed: those tests call GemmSwigluSm100.execute and friends directly, reference no pygraph/variant_pack/_normalize, and an import trace over three of them shows no cudnn.gemm.frost.* module is ever loaded (only cudnn._pygraph, which bare import cudnn pulls in regardless).

Correction

Above I wrote that -p no:randomly mattered because "pytest-randomly reseeds before each test". That was wrong — the plugin is not installed here, so the flag was inert, and I attributed cause from a single-sample A/B. Withdrawn.

What is actually true about test_mhas.py: it writes shared pytest config state mid-test (request.config.option.left_bound = None, ...dropout = None), which every later test in the same xdist worker reads back as its own parameters — so its outcome depends on ordering and worker assignment. Three runs of one unchanged commit gave 331 / 155 / 173 failures, all from that file. It is excluded here on the maintainer's instruction; test_mhas_v2 is the SDPA suite that counts, and it is green.

One more commit: 98578b5b6

Taking the "keep the fallback, but do not let it become a second permanent path" point seriously enough to measure it. Both halves of "why not the fast path" are named data now:

• compiled.declined — which rule denied this graph a fast path at build (needs workspace, post-kernel sqrt, multi-GEMM without a dense output, …), instead of a bare None.
• compiled.deferrals — the eight per-call hand-over sites, counted by reason, incremented only on the path already giving up so the fast path pays nothing.

The counts separate two populations that reading the code does not. Six of the eight sites mean the call is illegal (wrong major, TMA-misaligned extent, misaligned output base, SF blob that is not a dense run) and exist only so the interpreter can produce the message — those never justify widening the plan, and a fused fast check plus a slow re-walk to name the failing check serves them completely. The other two are calls that are legal and the fast path cannot serve: an operand described from the graph (bare device address, or a buffer reporting the declaration), and a strided reduction output. Only that second population is a reason to widen.

Three tests: every flavor's legal call defers for nothing, a decline names its rule, and a bare address is counted under exactly that reason and no other. That is what makes deleting the fallback later a measurement rather than a leap.

[YangXu1990uiuc]

Pushed 943027276 — this closes the two-track structure that the review flagged as an accepted intermediate. There is one executor now. The PR description is updated; the short version:

The deferral counters added in 98578b5b6 split the eight hand-over sites into two populations that reading the code does not. Six mean the call is illegal and only ever needed a message. Two meant the call was legal and the loop could not serve it — and while an interpreter sat behind the fast path, "the fast path refused a legal call" was a performance bug at worst with nothing asserting otherwise. Both are gone:

• An operand in the graph's axis order. cuDNN declares B [b, K, N] where the kernel reads (b, N, K), so a bare device address — and any buffer reporting the declared (dim, stride) — arrived the other way round and paid a whole interpreted pass for being legal. It is re-labelled with one permute whose shape the recipe already carries, before any check runs. Plain gemm 19.95 → 19.76 µs, i.e. free on a call that does not use it. This also retires the declared layout is indistinguishable from the kernel's build-time decline, which existed only because the stride guard alone cannot tell the two orders apart at a unit extent.
• A strided reduction seed. fill_word_strided_async collapses the layout to the runs a memset can cover and issues cuMemsetD2D32Async, which takes a pitch — one call per remaining outer point (the batch, usually 1) rather than one per row. That was the last place the engine wrote through a buffer it does not own: tensor.fill_() works only while the caller passes a torch tensor, and queues on torch's stream rather than the one the kernel will run on.

Three build-time declines also kept launch alive as an executor — no tvm-ffi front door, a norm2 output, a multi-GEMM with no dense output. They are _check_executable now, called from probe_supported, so the engine declines the graph and it falls to the backend. _lower reads that same function instead of restating it.

launch then loses its second half: explain() runs the same recipe.problem / check_shapes / check_alignment and raises. _call_positional, _call_multi_gemm, _call_block_scale_multi_gemm and launch's assembly are deleted — 267 lines. MoE is untouched (≥ 2 launches, its own workspace, no recipe).

The invariant tightens from "the fast path may accept a subset" to: the set of calls the launch path refuses must equal the set of illegal calls. compiled.deferrals enforces it — a legal call of every flavor must leave it empty, asserted per flavor and for the bare-address form. The counter stopped being observability and became the test.

The rules are still written twice (fused in the guards, readable in explain), so explain falling off the end raises a distinct RuntimeError naming which guard fired, rather than returning quietly.

One capability removed, deliberately: norm2, and it was direct-call only. Its taps land through cross-CTA atomics, so the square root has to run after every CTA has contributed — a second launch, not an epilogue op. Nothing public could reach it: the backend refuses a norm2 reduction descriptor at build_operation_graph, so no execute() ever gets a plan for one (test_norm2_reduction_is_refused_at_build, already in this PR). The two test_reduction_mode_coverage cases that computed it by calling the compiled object with torch tensors are now one test asserting the decline and saying why.

The lowered-vs-interpreted differential goes with the interpreter, which I think is the right trade rather than a loss: two readings of the same wrong plan agreeing proves nothing, and that is exactly how the axis-order bug survived one — both paths read Operand.axes, both agreed, both were wrong. What replaces it is fast-vs-backend (test_a_degenerate_extent_is_refused_or_matches_the_backend, test_an_operand_reporting_the_declaration_agrees_with_the_backend) plus the per-flavor torch references in test_public_execute_flavors.py, which run through the public graph.execute() rather than beside it.

Cost, unchanged (min over 25 reps of a 64-call burst, 256×256×128 bf16, SM100): plain 19.76, epilogue 19.69, aux 22.26, 2 outputs 22.65, reduction 26.12, multi-gemm 22.54.

Test: test/python/gemm/frost — 4204 passed / 2062 skipped, zero failures. Targeted test_execute_recipe.py + test_public_execute_flavors.py: 84/84. Full test/python is running now and I will post the number when it lands.

[YangXu1990uiuc]

@cudnn-ci-bot run frost

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-559-02a47d2
Pipeline: 62424366
Targets: frost

[YangXu1990uiuc]

All three blockers fixed in 02a47d2c6, plus the fourth point — thank you, these were all real and the first one was wrong numbers, not just a bad message.

Your framing is right that they are one shape: a per-ROLE fact treated as a per-SLOT one, or a write issued before every rule has had its say.

1. Shared operand — fixed, and it was mine

recipe.build keys the operand map on id(tensor), so matmul(A, A) gives the A role and the B role the same index; re-labelling ops[idx] for B therefore handed A a transposed view of the same memory. lowered now builds one view per role (vs = [operands[i] for i in in_slots]) and never rewrites a slot. shared_layout moved to role positions for the same reason. No measurable cost at this operand count.

The test is test_one_buffer_in_two_roles_reads_each_role_s_own_axis_order: matmul(A, A) at d×d, asserting the recipe really does resolve both roles to one slot, then that the product is A @ A (which the graph's [b, K, N] B declaration makes asymmetric, so a transposed read cannot pass).

Worth noting: your repro surfaced as the drift RuntimeError, which is the guard doing exactly the job it was added for. Without it this would have been a silent wrong answer.

2. Seed before validation — fixed, three separate holes

• Element width is now a rule of the call. The word is 32 bits and the count is the buffer's numel, so a fp32 tap on a fp16 buffer wrote twice the bytes it owned. The launch's dtype check is correct but arrives after the fill, so the width is checked with the caller's memory still untouched. Test asserts the canaries either side are intact, not merely that it raised.
• Plan-all, then fill-all. strided_fill_plan returns a verified plan or None; every seed is planned before any is issued, so a second unseedable tap cannot find the first already filled.
• The overlap rule was too weak. pitch >= width is the innermost pair only and accepts shape (2,2) stride (2,2) — width 1, both axes on element 2. The rule is now that each axis clears the whole span of the one below it, and that case is in the table.

One ordering I did not change and want to flag rather than hide: a call the fast guards accept but the tvm-ffi front door then rejects (your TMA outer-stride case) still seeds first. Seeding has to precede the launch, so closing that means checking what the front door checks — which is the same "one ordered list" work as below, not a local fix.

3. Guard/checker coverage — fixed

_shared_layout_reject and _seed_reject added; recipe.problem checks the rank before it indexes an axis, so a rank-2 operand gets a message instead of an IndexError.

4. multi-GEMM without a dense output — supported, not narrowed

You are right that this was a second capability removal and I described only norm2. It is supported again: the decline was written when two multi-GEMM launchers read the batch off cs[0], those launchers are deleted, and the recipe already carries max(a_batch, b_batch) for the case. It is a seventh entry in _FLAVORS now, so it goes through the lower/legal-call/agreement tests like the rest. norm2 is the only removal, and the PR text says so accurately now.

On the invariant claim — you are right, and I have corrected it

deferrals empty on a legal call proves only that nothing legal is refused. It cannot prove nothing illegal is accepted, because an accepted call never reaches the counter. That direction is covered case by case in test_the_launch_path_accepts_exactly_the_legal_calls and by the two backend differentials, which is weaker than what I wrote. Corrected in _lower's docstring, the test-file docstring, HANDOFF.md and the PR body.

I agree the real answer is generating both the fused guards and the readable diagnostics from one ordered CheckSpec list, and I have recorded that as the open item rather than claiming the split is safe. It is also what would close the front-door ordering above.

Both design docs now carry a SUPERSEDED header: FROST_GEMM_UNIFIED_EXECUTION_HANDOFF.md's "Semantics versus executors" (no ReferenceExecutor shipped, and a build without tvm-ffi declines the graph rather than getting a LegacyCuteExecutor), and frost_gemm_execute_design.md, whose every _call_positional reference is to deleted code.

Test

test/python/gemm/frost: 4212 passed / 2062 skipped / 0 failed (was 4204 — the new tests plus the restored flavor). Targeted test_execute_recipe.py + test_public_execute_flavors.py: 91/91. Full test/python on the fixed head is running; the head before these fixes gave 89 failed / 14266 passed / 6325 skipped, matching your run and the pre-existing baseline exactly, with zero failures under test/python/gemm/.

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (6)

test/python/gemm/frost/test_execute_recipe.py (4)

513-516: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Drop the unused unpack target.

c is never used in _padded_rows. Ruff reports RUF059.

♻️ Proposed change

-    a, b, c = _operands()
+    a, b, _ = _operands()
     return a, b, torch.empty(1, M, N * 2, dtype=torch.bfloat16, device="cuda")[:, :, :N]

🤖 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/gemm/frost/test_execute_recipe.py` around lines 513 - 516, Update
_padded_rows to unpack only the operands it uses, removing the unused c target
that triggers Ruff RUF059 while preserving the returned tensor construction.

Source: Linters/SAST tools

---

160-165: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Extract the "last node output" lookup into a helper.

The idiom [t for t in g._nodes[-1].outputs.values()][0] repeats at lines 162, 185, 204, 212, and 381. Ruff reports RUF015 on each occurrence. A single helper removes the duplication and the hints.

♻️ Proposed helper

+def _last_output(g):
+    return next(iter(g._nodes[-1].outputs.values()))
+
+
 def _reduction_graph():
     g = _plain_graph()
-    Y = [t for t in g._nodes[-1].outputs.values()][0]
+    Y = _last_output(g)

🤖 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/gemm/frost/test_execute_recipe.py` around lines 160 - 165,
Extract the repeated last-node output lookup into a shared helper near the graph
test utilities, then replace each occurrence in _reduction_graph and the other
affected test helpers with that helper. Ensure the helper preserves the current
behavior of returning the first output from g._nodes[-1].outputs.

Source: Linters/SAST tools

---

838-848: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Assert the deferral reason for the batch mismatch.

Other rejection tests pin the counter, for example line 334 and line 607. Add the same assertion here so a rejection raised by an unrelated rule cannot pass this test.

♻️ Proposed change

     with pytest.raises(ValueError):
         compiled.lowered(_bound_buffers(compiled, a, one_batch_b, c), stream=None)
+    assert dict(compiled.deferrals) == {"operand batch": 1}

Use the exact reason string the launch path records for a batch mismatch.

🤖 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/gemm/frost/test_execute_recipe.py` around lines 838 - 848, Update
test_operand_batch_is_checked to capture the ValueError from compiled.lowered
and assert its exact message matches the launch path’s recorded batch-mismatch
deferral reason, following the counter assertions used by the other rejection
tests.

---

716-719: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Move the degenerate-extent sweep to a higher test level.

This sweep runs 16 cases, and each case builds and executes both a backend plan and a FROST plan. The file mark is L0 (line 40). Keep a small subset at L0 and mark the full sweep at a higher level.

As per coding guidelines: "Mark every new Python test with a level from L0 through L4; keep L0 tests fast and place large parameter sweeps at higher levels."

♻️ Proposed change

 `@requires_sm100`
+@pytest.mark.L2
 `@pytest.mark.parametrize`("batch", (1, 2), ids=("b1", "b2"))
 `@pytest.mark.parametrize`("m,n,k", [(m, n, k) for m in (1, 128) for n in (1, 128) for k in (1, 128)], ids=str)
 def test_a_degenerate_extent_is_refused_or_matches_the_backend(monkeypatch, batch, m, n, k):

🤖 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/gemm/frost/test_execute_recipe.py` around lines 716 - 719, Update
test_a_degenerate_extent_is_refused_or_matches_the_backend so the full 16-case
parameter sweep runs at a higher test level than L0, while retaining a small
representative subset marked L0. Ensure the test-level markers follow the
project’s L0–L4 conventions without changing the test’s backend/FROST validation
behavior.

Source: Coding guidelines

test/python/gemm/frost/test_public_execute_flavors.py (1)

26-29: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Reuse the shared capability mark.

gemm_test_utils.requires_sm100 already gates SM100 tests, and test/python/gemm/frost/test_execute_recipe.py imports it at line 30. A second local definition can drift from the shared one.

♻️ Proposed change

-_GPU = pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
-    reason="the FROST gemm engine claims SM100",
-)
+from gemm_test_utils import requires_sm100 as _GPU

Keep the import at the top of the module.

🤖 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/gemm/frost/test_public_execute_flavors.py` around lines 26 - 29,
Replace the local _GPU skipif definition with the shared
gemm_test_utils.requires_sm100 capability mark, keeping the existing top-level
import and applying the shared mark to the same tests.

python/cudnn/frost/buffers.py (1)

348-407: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value

The fill planner looks correct, including the overlap rule.

collapse_layout merges an outer axis into an inner run only when the outer stride equals the inner span, and it keeps the inner stride for the merged run. strided_fill_plan rejects a zero stride over a real extent and rejects any outer stride that does not clear the axis below it, so shape (2, 2) stride (2, 2) is refused as documented. The width, pitch, and height selection maps a rank-3 padded tap onto one 2D memset per outer point.

One note for a follow-up: fill_word_async and _fill_word_2d_async import cuda.bindings.driver on every call. The seeding path runs per launch, so binding the driver module once at import time would remove that lookup from the hot path.

Also applies to: 410-435

🤖 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 348 - 407, Move the
cuda.bindings.driver import out of the per-call paths in fill_word_async and
_fill_word_2d_async, binding the driver module once at module import time and
reusing that module in both functions.

🤖 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 `@docs/python_graph_and_execution_backends.md`:
- Around line 181-188: Update the fallback descriptions in
docs/python_graph_and_execution_backends.md lines 181-188 and
python/cudnn/frost/README.md lines 243-249: replace claims that unsupported
graphs are handed to an interpreter with wording that CompiledFusedGemm.explain
provides diagnostics, names the reason, and raises; preserve the surrounding
explanation of lowered-path behavior and rejection ownership.

In `@python/cudnn/gemm/frost/compiler.py`:
- Around line 2060-2071: Update the sfs validation loop around operands[idx] to
reject any scale-factor operand whose rank is not exactly 3 before the existing
alignment, density, and size checks. Add the same rank-3 requirement to
recipe._sf_blob_reject so explain reports the scale-factor blob rejection
instead of falling through to the drifted-apart error, while preserving the
existing MoE launcher contract.

In `@python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py`:
- Around line 2620-2622: Insert a sentence break in the docstring for the
descriptor-building kernel in
python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py lines 2620-2622,
placing “Launched” on a new line after “tensormap_workspace”. Apply the same
docstring-only correction in
python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py lines 3960-3962 and
python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py lines 2543-2545.

---

Nitpick comments:
In `@python/cudnn/frost/buffers.py`:
- Around line 348-407: Move the cuda.bindings.driver import out of the per-call
paths in fill_word_async and _fill_word_2d_async, binding the driver module once
at module import time and reusing that module in both functions.

In `@test/python/gemm/frost/test_execute_recipe.py`:
- Around line 513-516: Update _padded_rows to unpack only the operands it uses,
removing the unused c target that triggers Ruff RUF059 while preserving the
returned tensor construction.
- Around line 160-165: Extract the repeated last-node output lookup into a
shared helper near the graph test utilities, then replace each occurrence in
_reduction_graph and the other affected test helpers with that helper. Ensure
the helper preserves the current behavior of returning the first output from
g._nodes[-1].outputs.
- Around line 838-848: Update test_operand_batch_is_checked to capture the
ValueError from compiled.lowered and assert its exact message matches the launch
path’s recorded batch-mismatch deferral reason, following the counter assertions
used by the other rejection tests.
- Around line 716-719: Update
test_a_degenerate_extent_is_refused_or_matches_the_backend so the full 16-case
parameter sweep runs at a higher test level than L0, while retaining a small
representative subset marked L0. Ensure the test-level markers follow the
project’s L0–L4 conventions without changing the test’s backend/FROST validation
behavior.

In `@test/python/gemm/frost/test_public_execute_flavors.py`:
- Around line 26-29: Replace the local _GPU skipif definition with the shared
gemm_test_utils.requires_sm100 capability mark, keeping the existing top-level
import and applying the shared mark to the same tests.

🪄 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: f6e7e898-1b82-4479-9f45-7bd94a0f15db

📥 Commits

Reviewing files that changed from the base of the PR and between f06f1ef and 5d41d3e.

📒 Files selected for processing (18)

• docs/python_graph_and_execution_backends.md
• python/cudnn/_pygraph.py
• python/cudnn/engines/base.py
• python/cudnn/frost/README.md
• python/cudnn/frost/buffers.py
• python/cudnn/frost/workspace.py
• python/cudnn/gemm/frost/compiler.py
• python/cudnn/gemm/frost/dtypes.py
• python/cudnn/gemm/frost/engine.py
• python/cudnn/gemm/frost/recipe.py
• python/cudnn/linear_attention/engine_utils.py
• python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py
• python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py
• python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py
• python/pygraph/variant_pack.cpp
• test/python/gemm/frost/test_execute_recipe.py
• test/python/gemm/frost/test_matmul_epilogue_fusion.py
• test/python/gemm/frost/test_public_execute_flavors.py

[yanqinz2]

lgtm and my agent

[coderabbitai]

🧹 Nitpick comments (4)

test/python/gemm/frost/test_execute_recipe.py (3)

449-452: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Discard the unused unpacked value.

Ruff reports RUF059 at Line 451 because c is never read. The neighboring builders at Lines 422, 427, 432, and 437 already use _ for this.

♻️ Proposed change

-    a, b, c = _operands()
+    a, b, _ = _operands()
     return a, b, torch.empty(1, M, N * 2, dtype=torch.bfloat16, device="cuda")[:, :, :N]

🤖 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/gemm/frost/test_execute_recipe.py` around lines 449 - 452, Update
the _padded_rows function to unpack the unused second value from _operands()
into `_`, matching the neighboring operand-builder functions and resolving
RUF059.

Source: Linters/SAST tools

---

98-98: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Replace the single-element list comprehensions.

Ruff reports RUF015 at Lines 98, 121, 140, 148, and 317. Use next(iter(...)) to read the first output tensor. This keeps the lint clean if the repository gate runs Ruff on tests.

♻️ Proposed change (apply the same edit at each site)

-    Y = [t for t in g._nodes[-1].outputs.values()][0]
+    Y = next(iter(g._nodes[-1].outputs.values()))

Also applies to: 121-121, 140-140, 148-148

🤖 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/gemm/frost/test_execute_recipe.py` at line 98, Replace the
single-element list comprehensions at the affected output-tensor assignments,
including the sites around lines 98, 121, 140, 148, and 317, with
next(iter(...)) over the corresponding outputs collection so each assignment
retrieves the first tensor without constructing a list.

Source: Linters/SAST tools

---

652-654: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Move this sweep above L0.

Line 39 marks the whole module L0. This test expands to 16 parameter combinations, and each combination builds and compiles up to two execution plans through _matmul_on. That is a large sweep for the fast tier.

Add a higher level marker to this test, or reduce the extent set at L0 and keep the full sweep at a higher level.

♻️ Proposed change

 `@requires_sm100`
+@pytest.mark.L2
 `@pytest.mark.parametrize`("batch", (1, 2), ids=("b1", "b2"))
 `@pytest.mark.parametrize`("m,n,k", [(m, n, k) for m in (1, 128) for n in (1, 128) for k in (1, 128)], ids=str)

As per coding guidelines: "Mark every new Python test with a level from L0 through L4; keep L0 tests fast and place large parameter sweeps at higher levels."

🤖 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/gemm/frost/test_execute_recipe.py` around lines 652 - 654, Move
the parameterized test using the m/n/k sweep above the module-level L0 marker by
assigning it a higher test level, or restrict its L0 parameter set while
retaining the full sweep at that higher level. Update the test’s existing
markers near the requires_sm100 and parametrization declarations, preserving
coverage for all 16 combinations outside L0.

Source: Coding guidelines

test/python/gemm/frost/test_public_execute_flavors.py (1)

141-155: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Assert that the FROST iteration used the lowered path.

The test compares numbers between the two plan kinds. It does not check that the FROST plan accepted the operand. The docstring at Lines 125-128 names the defect: an operand buffer with no ndim looked like it already matched, so these calls were refused. A refusal is reported through deferrals, and the numeric comparison would still pass. The sibling test at Line 297 already asserts the lowered path.

Add the same assertion for the want_frost iteration so a regression fails here.

♻️ Proposed change

         y = torch.zeros(1, M, N, dtype=torch.bfloat16, device="cuda")
         _run(g, {A: a, B: b, bias: bias_buf, Y: y})
+        if want_frost:
+            plan = g._compiled_plans[g._plan_index]
+            assert plan._lowered is not None
+            assert dict(plan._compiled.deferrals) == {}
         out[want_frost] = y

🤖 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/gemm/frost/test_public_execute_flavors.py` around lines 141 -
155, Update the want_frost iteration in the test around g.select_plan and _run
to assert that the selected FROST plan reports the expected lowered-path
execution through deferrals, matching the assertion used by the sibling test.
Keep the existing numeric comparison and output collection unchanged.

🤖 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 `@test/python/gemm/frost/test_execute_recipe.py`:
- Around line 449-452: Update the _padded_rows function to unpack the unused
second value from _operands() into `_`, matching the neighboring operand-builder
functions and resolving RUF059.
- Line 98: Replace the single-element list comprehensions at the affected
output-tensor assignments, including the sites around lines 98, 121, 140, 148,
and 317, with next(iter(...)) over the corresponding outputs collection so each
assignment retrieves the first tensor without constructing a list.
- Around line 652-654: Move the parameterized test using the m/n/k sweep above
the module-level L0 marker by assigning it a higher test level, or restrict its
L0 parameter set while retaining the full sweep at that higher level. Update the
test’s existing markers near the requires_sm100 and parametrization
declarations, preserving coverage for all 16 combinations outside L0.

In `@test/python/gemm/frost/test_public_execute_flavors.py`:
- Around line 141-155: Update the want_frost iteration in the test around
g.select_plan and _run to assert that the selected FROST plan reports the
expected lowered-path execution through deferrals, matching the assertion used
by the sibling test. Keep the existing numeric comparison and output collection
unchanged.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: a5341e07-5cab-4733-af60-40d956a83429

📥 Commits

Reviewing files that changed from the base of the PR and between 5d41d3e and 5dfa3c1.

📒 Files selected for processing (6)

• python/cudnn/gemm/frost/compiler.py
• python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py
• python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py
• python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py
• test/python/gemm/frost/test_execute_recipe.py
• test/python/gemm/frost/test_public_execute_flavors.py

🚧 Files skipped from review as they are similar to previous changes (4)

• python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py
• python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py
• python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py
• python/cudnn/gemm/frost/compiler.py

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

test/python/gemm/frost/test_execute_recipe.py (1)

449-452: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Optional: clear the Ruff hints in this file.

Line 451 unpacks c and never uses it (RUF059). Lines 98, 121, 140, 148, 238, 317, and 536 take a single element through a list comprehension and an index (RUF015). Confirm whether Ruff gates CI for test/python; if it does, these fail the lint job.

♻️ Proposed fix for line 451

 def _padded_rows():
     """Legal: the outer stride is free, only the contiguous extent is pinned."""
-    a, b, c = _operands()
+    a, b, _ = _operands()
     return a, b, torch.empty(1, M, N * 2, dtype=torch.bfloat16, device="cuda")[:, :, :N]

🤖 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/gemm/frost/test_execute_recipe.py` around lines 449 - 452, Clear
the Ruff warnings in test_execute_recipe.py: update _padded_rows to avoid
binding the unused c value, and replace the single-element list comprehensions
at the referenced locations with direct element access. Confirm whether Ruff
gates test/python in CI and ensure these changes satisfy that lint
configuration.

Source: Linters/SAST tools

🤖 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 `@docs/python_graph_and_execution_backends.md`:
- Around line 149-178: Align the staged optimization figures in the “What a
per-execute path costs” section with the headline baseline: make the first-step
numbers start at the stated 40 µs baseline, or explicitly explain why the 44 µs
measurement differs. Ensure the sequence maps consistently to the reported 20 µs
target.

In `@test/python/gemm/frost/test_execute_recipe.py`:
- Around line 674-708: Raise the test level for
test_a_degenerate_extent_is_refused_or_matches_the_backend above L0 so its full
16-case parameter sweep is not included in the fast tier. Preserve the existing
parameterization and test behavior.

---

Nitpick comments:
In `@test/python/gemm/frost/test_execute_recipe.py`:
- Around line 449-452: Clear the Ruff warnings in test_execute_recipe.py: update
_padded_rows to avoid binding the unused c value, and replace the single-element
list comprehensions at the referenced locations with direct element access.
Confirm whether Ruff gates test/python in CI and ensure these changes satisfy
that lint configuration.

🪄 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: 73b6fb8d-97c9-43fc-aca9-50ce1df67be7

📥 Commits

Reviewing files that changed from the base of the PR and between 5dfa3c1 and 1667f73.

📒 Files selected for processing (8)

• docs/python_graph_and_execution_backends.md
• python/cudnn/frost/README.md
• python/cudnn/gemm/frost/compiler.py
• python/cudnn/gemm/frost/recipe.py
• python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py
• python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py
• python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py
• test/python/gemm/frost/test_execute_recipe.py

🚧 Files skipped from review as they are similar to previous changes (6)

• python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py
• python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py
• python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py
• python/cudnn/frost/README.md
• python/cudnn/gemm/frost/compiler.py
• python/cudnn/gemm/frost/recipe.py

[coderabbitai]

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Align the stated per-execute baseline.

Lines 153-155 give the re-deriving baseline as 40 µs and the target as 20 µs. Line 174 gives the first step as "44 → 35". The staged numbers (44 → 35 → 20) do not start from the headline number, so a reader cannot map the steps onto the summary. Use one baseline in both places, or state why the two measurements differ.

📝 Proposed wording fix

-Read the first set into a table at build (`gemm/frost/recipe.py` is the worked
-example) and let the call read the table. That alone is 44 → 35.
+Read the first set into a table at build (`gemm/frost/recipe.py` is the worked
+example) and let the call read the table. That alone is 40 → 35.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	#### What a per-execute path costs
	An engine owns its internals, and this section does not change that. It exists
	because the default outcome is expensive: an engine that re-derives its per-call
	facts lands around **40 µs of host time per execute**, and for a single-kernel
	op that is most of what the caller pays. The same kernel with those facts read
	once is **20**. Both numbers are `frost_gemm` at 256×256×128 bf16, host
	enqueue, min over 25 reps of a 64-call burst from a drained queue.
	The budget it has to fit in, all measured on SM100:
	| | µs |
	|---|---|
	| `cuLaunchKernelEx`, untraced | 1.85 |
	| one CuTe-DSL entry | ~3.6 |
	| `graph.execute()` entry + `_normalize` | ~8 |
	| **everything else is the engine's** | |
	Do not read a per-call cost out of an nsys trace: CUPTI adds ~2.2 µs per traced
	API call, which is more than the call.
	**Split the facts by when they are decided.** Operand roles and majors, packing
	factors, alignment requirements, output shape rules, which outputs need a seed —
	all fixed when the kernel compiled. M/N/K, strides and pointers arrive per call.
	Read the first set into a table at build (`gemm/frost/recipe.py` is the worked
	example) and let the call read the table. That alone is 44 → 35.
	**Then lower the table into one closure per plan**, with its constants captured
	and the operand structure flattened into the loop headers, so the call does no
	attribute lookup and takes no branch the build already settled. That is 35 → 20.
	#### What a per-execute path costs
	An engine owns its internals, and this section does not change that. It exists
	because the default outcome is expensive: an engine that re-derives its per-call
	facts lands around **40 µs of host time per execute**, and for a single-kernel
	op that is most of what the caller pays. The same kernel with those facts read
	once is **20**. Both numbers are `frost_gemm` at 256×256×128 bf16, host
	enqueue, min over 25 reps of a 64-call burst from a drained queue.
	The budget it has to fit in, all measured on SM100:
	| | µs |
	|---|---|
	| `cuLaunchKernelEx`, untraced | 1.85 |
	| one CuTe-DSL entry | ~3.6 |
	| `graph.execute()` entry + `_normalize` | ~8 |
	| **everything else is the engine's** | |
	Do not read a per-call cost out of an nsys trace: CUPTI adds ~2.2 µs per traced
	API call, which is more than the call.
	**Split the facts by when they are decided.** Operand roles and majors, packing
	factors, alignment requirements, output shape rules, which outputs need a seed —
	all fixed when the kernel compiled. M/N/K, strides and pointers arrive per call.
	Read the first set into a table at build (`gemm/frost/recipe.py` is the worked
	example) and let the call read the table. That alone is 40 → 35.
	**Then lower the table into one closure per plan**, with its constants captured
	and the operand structure flattened into the loop headers, so the call does no
	attribute lookup and takes no branch the build already settled. That is 35 → 20.

🤖 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 `@docs/python_graph_and_execution_backends.md` around lines 149 - 178, Align
the staged optimization figures in the “What a per-execute path costs” section
with the headline baseline: make the first-step numbers start at the stated 40
µs baseline, or explicitly explain why the 44 µs measurement differs. Ensure the
sequence maps consistently to the reported 20 µs target.

[coderabbitai]

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Move this sweep above L0.

The module marks every test L0 at line 39. This case runs 16 combinations (2 batches x 8 shape triples), and each combination builds backend plans and FROST plans and executes both. That is a large parameter sweep in the fast tier.

Mark this test at a higher level, or reduce the L0 set to a few representative degenerate shapes and move the full sweep up.

📝 Proposed level override

 `@requires_sm100`
+@pytest.mark.L2
 `@pytest.mark.parametrize`("batch", (1, 2), ids=("b1", "b2"))
 `@pytest.mark.parametrize`("m,n,k", [(m, n, k) for m in (1, 128) for n in (1, 128) for k in (1, 128)], ids=str)
 def test_a_degenerate_extent_is_refused_or_matches_the_backend(monkeypatch, batch, m, n, k):

As per coding guidelines: "Mark every new Python test with a level from L0 through L4; keep L0 tests fast and place large parameter sweeps at higher levels."

🤖 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/gemm/frost/test_execute_recipe.py` around lines 674 - 708, Raise
the test level for test_a_degenerate_extent_is_refused_or_matches_the_backend
above L0 so its full 16-case parameter sweep is not included in the fast tier.
Preserve the existing parameterization and test behavior.

Source: Coding guidelines

[YangXu1990uiuc]

@cudnn-ci-bot run frost

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-559-1667f73
Pipeline: 62438689
Targets: frost