Emit the FROST gemm parameter tables, and a demo that launches from them

[YangXu1990uiuc]

Before submitting

• I agree to license this contribution under the terms of LICENSE.txt.
• I ran pre-commit run and committed any formatting changes.
• I added GitHub labels: one cat-*, one or more mod-*, and one orig-* (see label list).

Affected area

FE OSS kernels or CuTeDSL / Benchmarks or performance

Summary

Don't take my word for it — run it:

CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 \
  python benchmark/gemm/frost/benchmark_bare_launch.py --vary-m 512

SM100 required. Every result is checked bit-identical against graph.execute() and the
script exits non-zero at the first mismatch, so any timing it prints was preceded by a
passing check.

bit-identical to graph.execute()  OK
rebound to a second buffer set    OK

=== host us/call (min over 25 bursts of 64) ===
  graph.execute(dict)                  20.28
  bare plan.execute(ptrs)               2.40
    building that ptr list              0.37
    cuLaunchKernelEx alone              2.11

=== serving M=512 from the M=128 block ===
  changed        : ['m', 'a_stride_l_0', 'out_stride_l_0']
  re-encoded     : descriptors [3, 8] of [3, 4, 8]
  bit-identical to a plan built at M=512    OK

20.3 us of host time per launch becomes 2.4, against 2.11 for the bare
cuLaunchKernelEx. Essentially all remaining host time is the driver.

Two changes:

1. FROST codegen emits its parameter tables — SLOT_TABLE, PATCH_GROUPS,
   PROBLEM_FIELDS — into the generated module. Purely additive: compiler.py is
   develop plus a table block and two call sites, +219/-2, and no generated kernel
   behaves differently.
2. A demo that builds the parameter block from those tables, takes the kernel out of
   the compiled cubin, computes the geometry from the generated module's own closed form,
   and launches.

Why

Codegen already decides, per kernel parameter, how wide it is and where its value comes
from — and then throws both away: the signature is a string join and the values are
unpacked positionally by the generated host. Anyone wanting to marshal the block has to
rediscover all of it.

The 20 us graph.execute() spends on this gemm is not the launch. It is re-deriving, per
call, facts settled when the plan was built: what shape each buffer is, which axis is
innermost, how to spell a CUtensorMap, which kernel to run.

The tables are read off the signature just rendered, not rebuilt alongside it, so they
cannot drift from the kernel they describe. A parameter codegen cannot classify comes out
as kind == 'unknown' with a null source and a consumer must refuse the kernel — a guess
would be worse than a refusal, and that is what makes this safe to extend one flavor at a
time.

The contract

Between build_plans() and any launch: dtype, rank, extents, strides, innermost axis and
base alignment are what the graph declared, and only the addresses may change. Never
checked — checking it is most of the 20 us. Documented in docs/frost_bare_launch.md.

An inference server keeps this without effort. --vary-m shows the next rung, where the
token count moves and the caller re-supplies problem_size; PATCH_GROUPS then prunes the
work to three stores, the descriptors that span M, and gridDimX.

Scope, and how it extends

The demo is deliberately one narrow case — dense bf16 matmul, TMA-store epilogue,
SM100 — so its claims can be checked in a minute.

What is classified today, measured by compiling each flavor and comparing the table's
widths against cuFuncGetParamInfo on the resulting cubin:

flavor	outcome
plain, relu epilogue, 1024³, nvfp4 block scale	every parameter classified
aux bias, two dense outputs, amax reduction, multi-gemm	refused: 1–2 unknown
MoE	no table emitted at all

Every refusal is the same parameter: an STG epilogue passes its output as
mC_tap_i: cute.Tensor, whose fake this does not model. That is the mechanism working,
not a gap in it, and it is the obvious next increment — codegen writes that fake's shape
and stride expressions a few lines from where the table is built.

MoE gets no table because its problem_size carries num_experts and num_groups where a
dense one carries batch, and its operands are unclassified. A per-call stream is also not
done; that is one more store into the launch configuration.

API and compatibility impact

None. SLOT_TABLE / PATCH_GROUPS / PROBLEM_FIELDS are new module-level constants in
generated FROST gemm kernels; nothing else changes. The benchmark needs CUTE_DSL_KEEP=cubin
before the first compile (it sets it itself) and points CUTE_DSL_DUMP_DIR at a temp
directory so it leaves nothing behind.

Testing

On one SM100 part.

• test/python/gemm/frost: 5767 passed, 2861 skipped, 0 failed.
• The demo at 128x256x64 --vary-m 512 and 1024x1024x512 --vary-m 2048: bit-identical to
  graph.execute() and to a plan built at the target M.
• Earlier sweeps (before the scope narrowed) covered six shapes / four tile configs / three
  distinct dynamic-SMEM sizes, all bit-identical; and the SMEM number the demo recovers was
  checked against what a captured launch actually passed, 8/8 shapes exact.
• Parameter block rebuilt from SLOT_TABLE and diffed against a real launch byte for byte.
• pre-commit / black clean.

The comments below record what changed after the first push and why, including one part of
the original PR that I reverted after an A/B showed it removed a real check.

---

_{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_demo}

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 83a9f40b-3399-4707-b1af-5a7788d73905

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

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

[YangXu1990uiuc]

Test result

test/python/gemm/frost on one SM100 part, with this diff applied:

5767 passed, 2861 skipped in 190.94s (0:03:10)

That is the suite that covers the tap cute.Tensor → cute.Pointer change, which is the
only part of this PR that alters a generated kernel — the tables are additive. It includes
the MoE grouped-matmul launchers, block-scale, swiglu, multi-gemm and the epilogue-fusion
flavors.

CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 CUDA_VISIBLE_DEVICES=<sm100> \
  python -m pytest test/python/gemm/frost -q -n 8

[YangXu1990uiuc]

Shape sweep

The first push only worked at the default shape — at almost any other M/N/K the epilogue
picks a TMA store, so the output arrives as tma_c_desc_0 rather than mC_tap_0 and the
demo died on KeyError: ('c', 0). Fixed in 2d05d4b: roles come from SLOT_TABLE instead
of a hardcoded list, and the C descriptor is built too, including the
a_is_m_major / b_is_n_major / cd_out_is_m_major branches — which select different
stride components, so the wrong branch would encode a valid-looking descriptor pointing at
the wrong memory.

Six shapes, four tile configs, three distinct dynamic-SMEM sizes, both epilogue paths,
each also run through --vary-m:

M	N	K	grid	cluster	smem	epilogue	graph.execute	bare	--vary-m
256	256	128	(2,8,1)	2x1	149504	tap	19.52	2.42	512 OK
128	256	64	(1,8,1)	1x1	182272	TMA store	20.25	2.44	256 OK
1024	1024	512	(8,8,1)	2x1	215040	TMA store	20.53	2.40	2048 OK
512	2048	256	(4,16,1)	2x1	215040	TMA store	20.45	2.38	1024 OK
4096	4096	1024	(32,32,1)	2x1	215040	TMA store	20.48	2.36	8192 OK
128	128	64	(1,4,1)	1x1	182272	TMA store	20.15	2.40	256 OK

Every row printed bit-identical to graph.execute() OK, rebound to a second buffer set OK, and bit-identical to a plan built at M=<2x> OK.

Three distinct SMEM values is the interesting column: the demo recovers that number by
finding the single unique constant in the compiled host module, and it stayed correct
across every tile config here. A wrong value would fault or corrupt, so the bit-identity
check is the guard.

[YangXu1990uiuc]

The dynamic-SMEM number, checked against ground truth

This is the one field in the demo with no first-class source, and the one a bit-identity
check cannot fully police: too small faults, but too large is legal — shared memory can
be over-requested — so a wrong-but-workable value would pass every correctness check while
quietly costing occupancy.

So it is checked directly: capture a real graph.execute into a CUDA graph, read the
kernel node's sharedMemBytes (what the DSL's own host function passed), and compare.

  shape                      read  launched   agree
  256x256x128              149504    149504   OK
  128x256x64               182272    182272   OK
  1024x1024x512            215040    215040   OK
  512x2048x256             215040    215040   OK
  4096x4096x1024           215040    215040   OK
  128x128x64               182272    182272   OK
  2048x512x128             215040    215040   OK
  256x4096x512             215040    215040   OK

8/8 shapes read the value the DSL actually launched with

It still deserves the upstream fix the doc asks for — a property next to __cubin__ —
because "the single unique constant in that range" is a property of today's generated host,
not a contract. Until then this is the authoritative source rather than a recomputation,
which would duplicate the kernel body's SMEM layout and drift the first time a buffer is
added.

[YangXu1990uiuc]

One more, found by reading the templates rather than running them

The table computed where output strides start as 4 + 3*na + 3*nb. That is what the dense
and mainloop templates do — _stride_idx = 4, advanced three per A and per B operand — but
two other paths disagree, and a table has to describe the kernel as generated, not as
the layout implies.

MoE puts num_experts and num_groups where a dense problem_size carries batch,
and starts its stride triples at 5:

# sm100_moe_grouped_matmul_fwd_1ctamma.py
num_experts = problem_size[3]
num_groups  = problem_size[4]
_stride_idx = 5

So the table labelled index 3 'batch' and read every stride one slot early. It was inert —
a MoE kernel's mA_i parameters are unclassified, so every such table already carried
unknown and a consumer had to refuse it — but shipping wrong-and-inert data is exactly
what this mechanism exists to avoid. MoE now gets no table.

Block-scale hands its host a literal 10 while its templates set
supports_multi_gemm=True:

red_host_stride_unpack = _reduction_stride_host_unpack(chain) if chain.has_moe \
    else _reduction_stride_host_unpack_from(chain, 10)

For na = nb = 1 that agrees with the formula. For a multi-GEMM block-scale chain it does
not, and the table would have named the output strides where the operands' triples still
are. Each renderer now passes the base it actually used, and the table refuses to describe
a kernel whose operand triples do not end exactly there.

Separately, and untouched here: whether that literal 10 is itself a bug in the
block-scale host for multi-GEMM chains is worth a look by whoever owns that renderer. The
dense template computes the position; the block-scale one asserts it. I have not built such
a graph to confirm it is reachable, so I am flagging rather than claiming.

Re-verified after the change: eight gemm flavors still classify with zero unknowns and
widths matching cuFuncGetParamInfo; the demo is still bit-identical at M=256 and
--vary-m 512.

[YangXu1990uiuc]

I reverted part of this PR: the tap change removed a real check

While sweeping the demo I had an adversarial pass read the diff, and it flagged the
cute.Tensor → cute.Pointer tap change as dropping the only check on a dense output's
layout. That reproduces. Same graph, same output shape, three runtime layouts:

output layout	gh/develop	with the tap change
contiguous (the declaration)	accepted, correct	accepted, correct
padded rows, N unit-stride	rejected: alignment	rejected: alignment
transposed, N stride = M	rejected: Mismatched c_tap_0.strides[1]	CUDA misaligned address

The old compile fake declared stride=(cute.sym_int64(), 1, cute.sym_int64()) for an
N-major dense tap. That literal 1 is static, so the DSL's front door had something to
compare a runtime tensor's layout against, and it rejected a mis-strided output on the
host. A cute.Pointer carries no layout, so nothing compares and the epilogue's vectorised
store faults on the device instead.

Not silent corruption in this case, but a strictly worse failure mode — and nothing rules
out an aligned-but-mis-strided layout that would be silently wrong. It also fires on the
ordinary graph.execute path, which the demo's contract does not cover: that contract
is between a caller and a bare launch, not between a caller and the public API. So the
optimisation comes out. The tables are the deliverable.

What that costs, stated honestly

mC_tap_i is a cute.Tensor again and is not classified, so four flavors now emit a
table containing unknown and a consumer must refuse them:

flavor	outcome
plain, relu epilogue, 1024³, nvfp4 block scale	every parameter classified
aux bias, two dense outputs, amax reduction, multi-gemm	refused: 1–2 unknown
MoE	no table emitted at all

That is the mechanism working rather than a gap in it, and it is the obvious next
increment — codegen writes that fake's shape and stride expressions a few lines from where
the table is built. The earlier "eight flavors, zero unknown" comment was measured with
the tap change; this supersedes it.

The headline is unchanged. The demo's default moves to 128x256x64, which takes the
TMA-store epilogue and has no tap at all:

bit-identical to graph.execute()  OK
rebound to a second buffer set    OK
  graph.execute(dict)                  20.10
  bare plan.execute(ptrs)               2.39
    cuLaunchKernelEx alone              2.12

Two doc claims corrected

• "Changing M re-encodes A only" was true only for the STG flavor. Under a TMA store
  the output descriptor spans M too, so --vary-m re-encodes A and C — the demo now
  prints re-encoded: descriptors [3, 8] of [3, 4, 8]. B still never moves, because N and
  K did not.
• The SLOT_TABLE example showed a tap and a C descriptor in the same kernel. No kernel
  has both; it is now verbatim output.

And one correction to my own earlier comment

I suggested the block-scale renderer's literal 10 might be a bug for multi-GEMM chains.
It is not. The block-scale templates read exactly one A triple and one B triple
whatever na and nb are, and recipe.py sends only those two — so 10 is right and the
formula was the wrong thing to compare against. The guard I added is still correct (it
refuses to describe a kernel whose operand triples do not end where the host starts
reading), but my reasoning for it was not. Nothing to look at in that renderer.

Housekeeping

Four .cubin files got swept into an earlier commit by git add -A. They are removed at
HEAD but remain in this branch's history; taking them out needs a force push, which I have
not done.

Final state

compiler.py is now gh/develop plus the table block and the two renderer call sites,
nothing else — +219/-2, and the only line mentioning a pointer is the cute.Pointer entry
in _PARAM_ABI, a table entry rather than a codegen change.

test/python/gemm/frost      5767 passed, 2861 skipped, 0 failed

128x256x64   --vary-m 512    bit-identical OK   20.28 -> 2.40   (launch alone 2.11)
1024x1024x512 --vary-m 2048  bit-identical OK   20.39 -> 2.38

(An intermediate commit had me hand-write the MoE launcher lines back from memory instead
of reading them; that failed 203 tests and is fixed by taking the file from base and
re-applying additively. Left in history rather than squashed, since squashing needs a
force push.)

[YangXu1990uiuc]

Doc corrections: three claims the demo could not support

The adversarial pass over the diff also flagged the documentation, and it was right three
times. All the same shape — evidence gathered by probes that are not in this PR,
written as though the demo produced it.

1. "the demo does, and diffs it against what a real launch passed, byte for byte."
   It does not. It builds the block from SLOT_TABLE and checks it two ways: slot widths
   against cuFuncGetParamInfo on the loaded cubin, and output bit-identity against
   graph.execute(). The byte-for-byte diff is real but lives in a separate probe.

2. The 10.7 us row ("graph.execute with only the engine's launch closure replaced")
   sat in the measurement table directly under Run it:, and nothing here can produce it —
   it needs the engine's launch closure monkey-patched. It is a real measurement and the
   most useful one for deciding what to fix first, so it stays, now labelled with where
   it came from and that this file does not reproduce it.

3. "the TMA descriptors are built by device kernels into a workspace rather than encoded
   on the host" overstates linear attention. The base descriptors are host-encoded;
   only the per-(batch × head) arrays are device-built.

The measurement table also still carried M=N=256 K=128, the shape from before the scope
narrowed. It now carries the demo's own default and its own numbers.

One code comment added: _MAP_DEPS names a stride triple per operand index, which is
right only where the host sends one per operand. Block-scale sends a single shared A and B
triple whatever na and nb are — the stride-base guard is what keeps those chains from
reaching it, and that is now written down rather than incidental.

Nothing about the headline changes; the demo is unaffected and still bit-identical.

[YangXu1990uiuc]

Review closed: 12 confirmed findings, all addressed

The adversarial pass finished — 5 dimensions, every finding put through 2 independent
refutation attempts. 19 raised, 12 survived, and they cluster into four things, all now
fixed:

#	finding	resolution
1–3	stride_base = 4 + 3*na + 3*nb is wrong for block-scale multi-GEMM — every stride slot off by 3*(na+nb-2), with nothing marked unknown, so the "refuse rather than guess" valve never fired	each renderer passes its real base; a chain whose operand triples don't end there gets no table
4	_MAP_DEPS keys descriptor deps by operand index, but block-scale shares one A/B triple	same guard; now written down in a comment rather than incidental
5, 8, 9	tap-as-pointer drops the N-contiguity and dtype checks	reverted
6, 7, 10–12	four doc claims citing evidence not in this PR	corrected

The one worth reading

Findings 1–3 are the sharpest thing the review produced, and they are not the bug I
would have guessed. The table was right about the dense templates and wrong about
block-scale, and it was wrong silently — no unknown, so a consumer following the
documented contract would have built a block with output strides read from where operand
strides live. For the two-output reduction variants that is not even an IndexError; it
is output 0 stored with output 1's strides.

I had fixed this before the review returned, but from the wrong premise — I thought the
block-scale renderer's literal 10 was a bug. It is not: those templates read exactly one
A triple and one B triple whatever na and nb are. Right guard, wrong reason, and the
reason matters because it decides whether anyone should go change that renderer. Nobody
should.

Verified by building the chain the review named rather than by arithmetic:

=== dual block-scale nvfp4: na=1 nb=2 ===
  SLOT_TABLE emitted   : False
  PATCH_GROUPS emitted : False
  4 + 3*na + 3*nb = 13, renderer's base = 10
  OK: refused, as the guard intends

=== dual block-scale mxfp8: na=1 nb=2 ===   (same)

Where that leaves coverage

flavor	outcome
plain, relu epilogue, 1024³, single-GEMM nvfp4 block scale	every parameter classified
aux bias, two dense outputs, amax reduction, multi-gemm	refused: unknown tap
block-scale multi-GEMM, MoE	no table emitted

Narrower than the first push claimed, and every boundary is now a refusal rather than a
guess — which is the property the whole design rests on.