The fp8 A8 GEMM stages its operands through TMA: prefill +2.7% at chunk 1024 and +4.7% at 4096

[MichaelDementii]

A second implementation of the FP8 row-scaled A8 GEMM, and a cost model that routes to it where it
is cheaper. The arithmetic is untouched: same accumulator, same m16n8k32, same scale application,
same epilogue and output policies. Five call sites reach it — linear, linear_add,
linear_swiglu, attn_input_proj, gdn_input_proj — all at prefill widths.

This revision answers the four items in your review; the reply below says what changed for each. It
is rebased onto 487f8977 and every number here was re-taken on that base.

What actually changes

Not one operand-copy mechanism swapped for another. Four things move together and only the last is
about TMA:

	cp.async route	TMA route
output tile per CTA	64 tokens x 128 rows	256 tokens x 128 rows
consumer warp tile	32 x 32	64 x 64
accumulator registers per consumer thread	32	128
K pipeline	2 stages of K=128	4 stages of K=64
buffered K extent	256	256 (unchanged)
threads per CTA	256 (8 warps)	288 (8 consumer + 1 producer warp)
CTAs per SM	2	1
registers per thread	94	166
K-loop synchronisation	CTA-wide barrier per stage	producer/consumer mbarrier pair

The wider warp tile is the arithmetic change: an ldmatrix of A now feeds eight N fragments instead
of four and a B fragment feeds four M fragments instead of two, so MMA issued per operand load goes
from 8 to 32. The deeper, narrower pipeline and the dedicated producer are what let one CTA keep
that tile fed without the per-stage CTA-wide barrier.

The same table is the cost. Resident warps per SM fall from 16 to 9, so a shape with too few token
tiles to fill the machine pays that and gets none of the reuse back. That is why narrow widths are
declined, and it is what the cost model has no term for.

Per SM the route uses fewer of both budgeted resources than the one it replaces: 98816 B of shared
against 100352 B, and 166 x 288 registers against 94 x 512.

Measured

One RTX 5090, sm_120a, 170 SMs, 525 W. Clock locking is not available on this machine and the card
idles at 180 MHz, so every table is taken with the arms alternating inside each repetition.

Operator, ratio of this route to the route it replaces, three repetitions, per-repetition spread
under 0.012 throughout:

shape	call site	T=1024	T=4096	T=8192
5120x17408	linear_add	0.895	0.932	0.897
5120x6144	linear_add	0.933	0.978	1.000
5120x17408	linear	0.877	0.917	0.899
5120x6144	linear	0.891	0.922	1.000
14336x5120	linear	0.991	0.951	0.957
16384x5120	linear	1.000	0.981	0.955
34816x5120	linear	1.000	0.952	0.939

A cell at 1.000 is a width the predicate declines, so both arms run the same kernel — the bounds are
visible in the same numbers that measure the gain.

End to end. Of the published artifacts only qwen3.8-27b carries FP8_E4M3FN_ROW_BF16S
weights; qwen3.6-27b is nvfp4 and qwen3.6-35b-a3b is groupwise-int end to end, so this route does
not run on either and measuring there returns 1.000 by construction. ninfer_bench -p 16384, six
repetitions alternating:

prefill chunk	base, tok/s	this branch, tok/s	change
1024	8045.4 ± 69.1	8262.2 ± 38.8	+2.70%
4096	8620.5 ± 60.5	9024.6 ± 26.1	+4.69%
8192	8581.0 ± 53.9	8920.7 ± 24.4	+3.96%

± is the sample standard deviation over the six runs; each percentage is the mean of the six
paired ratios, so it does not reproduce exactly from the rounded cells.

Decode is not on this path, and that is a census rather than an assumption. -n 512 with
--cuda-graph-trace=node so graph contents are visible: 29 distinct kernels execute and
fp8_a8_tma_kernel is not among them, 0 launches. Decode throughput on the same artifact moves
+0.03% over four repetitions, ratios 0.9999 to 1.0006.

Suite and tools, on the tree this pushes:

• ctest: 113 of 114 pass. The one failure is ninfer_attn_input_proj_test, and every one of
  its failing cases is W8 DFlash2 A16 ... graph phase=1; it fails the same way on clean
  487f8977. No FP8 case fails in either.
• compute-sanitizer, ten runs, 0 errors in every one: memcheck at 1024, 1345, 4001 and 4096
  — the last two are widths that are not a whole cp.async token tile, which is where the partial
  trailing tile reads past the token extent; racecheck, synccheck and initcheck at 256 and
  321, narrow because they instrument every shared access and a full-width run does not finish.
• greedy output against the previous route on the same artifact, six prompts at chunks 1024 and
  4096, both arms in both orders: 12 identical, 0 different, 0 invalid. Read that narrowly:
  these prompts are chunked into widths that are whole cp.async token tiles, which is where the two
  kernels agree byte for byte anyway. It is evidence that nothing else moved, not evidence of
  bit-identity in general — see item 3 in the reply.

Routing

A comparison of two quantised costs — waves needed times the work one SM carries through a wave —
with one empirical constant, plus a width floor and one per-geometry ceiling the model cannot
express. The constant is for this part and is not portable; what travels is the shape of the
comparison.

The sweeps behind the constant and both bounds are in docs/maintainer/fp8-a8-tma-route.md, added
by this change and indexed from docs/README.md. That document also carries the shared-memory
alignment measurement and the sweep that removed the multiple-of-tile condition, and it records two
things the bounds cost: the floor gives up four measured gains of 8 to 12 percent to avoid one loss
of 2, and the ceiling on 5120x6144 costs plain linear 3 to 5 percent while being correct for
linear_add, which is that geometry's production call site.

What this does not claim

• It does not claim the memory traffic is unchanged. Unchanged logical inputs do not establish
  unchanged physical traffic, and I did not measure traffic.
• It quotes no bandwidth derived from assumed bytes divided by time. Where a share of peak appears
  it is the column ninfer_linear_bench publishes, and that is a model, not a counter.
• It makes no claim about which resource binds. The peak denominator was not re-established on this
  machine, and the cold-operator bench does not reproduce the inference-time picture.
• It is not bitwise identical to the route it replaces. At widths that are a whole cp.async
  token tile the outputs still match byte for byte; at other widths the two kernels round
  differently by construction, and the numerical tests check those against a host reference.

Notes

The rebase kept the six-schedule ladder added to fp8_attn_input_a8.cu after this branch's original
base; the TMA route is the wide arm of that ladder, beside Prefill, which is the schedule the
predicate compares against. The three conflicting test files keep their existing cases and append.

ninfer_attn_input_proj_test fails on 487f8977 without this branch, in the W8 DFlash2 graph
replay. Reproduced from clean master with both arms run from one build directory; the number of
failing cases varies between runs, which is itself worth knowing. Details in the reply below.

🤖 Generated with Claude Code

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-03T10:44:53.716029Z	3d6f7f2	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[Neroued]

This review is posted by an agent on behalf of @Neroued. The comments below have been reviewed by the maintainer.

Reviewed commit: 3d6f7f2e71e729c9235ef35c8e1878224a59145b. This is a static review of the implementation and the submitted evidence; no local compilation, tests, or profiling were performed for this review.

The optimization direction is sound, and the TMA kernel is worth keeping. I would like the following addressed before merging.

1. Make the shared-memory alignment requirement of the swizzle indexing explicit.

   In fp8_a8_tma.cuh, fp8_tma_shared_byte() computes the 64-byte swizzle from relative row coordinates, implicitly assuming a zero swizzle base offset. The shared storage and dynamic shared backing only declare 128-byte alignment. The 64-byte swizzle pattern repeats at a 512-byte boundary; an otherwise valid 128-byte-aligned base can require a nonzero offset. Please either guarantee 512-byte alignment for the relevant backing, keeping the stage/subtile offsets consistent with it, or incorporate the actual shared base offset into the indexing.

   This is a missing source-level correctness precondition, not a claim that the submitted binary has already produced incorrect output. A particular compiler layout can satisfy the stronger alignment incidentally. See NVIDIA's swizzle offset description.

2. Keep the routing policy compact and distinguish calibration from a performance guarantee.

   The wave-count comparison is a reasonable heuristic, and a non-monotonic admission frontier is not inherently a problem. However, the 0.936 coefficient is fitted at one shape/width and shared across different K dimensions and epilogues. The model does not represent pipeline fill or epilogue cost, and the floor plus the 5120x6144 ceiling compensate for measured cases it cannot distinguish.

   Empirical tuning for the RTX 5090 fits this project's scope; this is not a request for a general model covering other GPUs. Please simplify the large calibration/assertion section in fp8_a8_schedule.cuh: retain geometry/resource invariants and necessary coverage checks, and move the detailed measurement history to a maintained performance reference. Assertions that a predicate keeps returning a particular answer do not establish that the selected route remains faster. Express the applicable calibration domain and known bounds clearly.

3. Reconsider making the old kernel's bitwise behavior a permanent routing constraint.

   The new kernel handles a partial trailing token tile, but fp8_a8_tma_applies() additionally requires tokens % MmaSchedule::kBlockTokens == 0 to match the old FullTokens branch's observed floating-point contraction behavior. That couples the new route's performance coverage to a private property of the old implementation.

   Please distinguish a measured performance reason for excluding these widths from a restriction imposed only to simplify memcmp. The long-term qualification criterion is the production Op against its independent mathematical oracle at the defined semantic boundaries. Bitwise comparison with the previous kernel is useful supporting evidence, but should not become an additional public numerical contract. I am not asking to remove the gate without evidence; I am asking for its long-term justification to follow the Op contract and performance requirements. The comments also currently describe bit identity more strongly than the PR body's toolchain-dependent caveat warrants.

4. Reuse the existing mbarrier helpers.

   The PR adds fp8_mbarrier_init, fp8_mbarrier_wait, fp8_mbarrier_arrive, and fp8_mbarrier_arrive_expect_tx, while its base already provides the corresponding functionality in src/ops/common/mbarrier.cuh, used by the NVFP4 TMA implementation. These operations have no FP8-specific semantics. Please reuse the existing helpers so synchronization fixes and maintenance have one owner. This does not require introducing a generic GEMM framework.

The performance explanation should also describe the actual combined change: TMA delivery, a larger output tile, and a different warp decomposition. It is more than an operand-copy substitution:

• The output tile changes from 64x128 to 256x128.
• The consumer warp tile changes from 32x32 to 64x64, increasing accumulator storage and operand reuse.
• Two stages of K=128 become four stages of K=64: the nominal buffered K extent remains 256.
• Dedicated producer/consumer synchronization replaces the old K-loop's repeated CTA-wide synchronization pattern.

These changes provide credible mechanisms for the reported speedup, while also explaining the occupancy and small-shape tradeoffs. They can change repeated operand requests and cache behavior, so unchanged logical inputs do not prove unchanged physical memory traffic. Likewise, bandwidth derived from assumed bytes divided by time is not a DRAM counter measurement; the unreestablished peak denominator and the cold-operator versus inference discrepancy limit the roofline/bottleneck conclusions. Please narrow those claims accordingly. This does not by itself invalidate the reported elapsed-time improvements.

The reported default-chunk prefill gain is worth pursuing. The shared kernel body with Op-specific epilogue/output policies, the explicit paired SwiGLU handling, guarded tail accesses, and reuse of the existing workspace are good aspects of the implementation. Keeping the old route for widths where it is needed is justified. Per-launch descriptor construction and repeated device queries can remain follow-up optimization opportunities; this review has no evidence that they should block this change.

[MichaelDementii]

Thank you. All four are addressed, and the branch is rebased onto 487f8977. Every number below
was re-taken on that base: the previous table came from a machine we no longer have, so none of it
is carried over.

The first item turned out to be a defect rather than a documentation gap, so I will lead with that.

1. Shared-memory alignment of the swizzle indexing

You were right, and it is stronger than "the source does not state the precondition": the swizzle
is a function of the shared-memory address, not of an offset relative to the destination. The same
tile, the same descriptor, two destinations, both de-swizzled with fp8_tma_shared_byte:

dynamic shared base       : 0x600   (mod 512 = 0)
tile at atom-aligned base : 0 / 1024 bytes wrong
tile at base + 128        : 1024 / 1024 bytes wrong

And __align__ on the extern __shared__ array is what places that base. With 16 bytes declared it
follows whatever static shared the kernel carries — measured at 16, 64, 128, 0 and 128 mod 512 for
16, 64, 128, 1024 and 1152 bytes of static shared. With 128 declared, CUDA 13.1 puts it at 0x400
in all five cases, which is 512-aligned. That is why the submitted code computed correct results: a
layout accident, exactly as you said.

kFp8A8TmaSwizzleAtomBytes = 512 now sits on the tensor storage, the union around it, the outer
storage and the allocation, with static asserts that the two stage strides and a paired block's
second branch are whole atoms. It costs 384 bytes of padding, sizeof 98432 to 98816 against the
101376 cap, and no time — the base is 0x400 either way.

2. Routing policy: calibration separated from the code

The measurement history is out of the header and into docs/maintainer/fp8-a8-tma-route.md, indexed
from docs/README.md. The header keeps the schedule registrations, the model, the two bounds with
the reason each exists, and the coverage checks. The section this branch adds to it went from 387
lines to 284.

You are right that the assertions could not carry what they looked like they carried. The ones that
remain assert coverage and say so: that the predicate still sends the widths the numerical tests run
at down this path, so those tests exercise the kernel they were extended for. Nothing there claims
the route is faster; that claim lives in the document next to the numbers, where re-measuring can
amend it.

The constant's description now says what the model represents — wave quantisation, and nothing else:
no term for pipeline fill, none for K, none for the epilogue — and that both bounds exist because of
shapes it cannot tell apart. One consequence is written down that I had not stated before: the
ceiling on 5120x6144 is correct for linear_add, which is that geometry's production call site,
and costs plain linear 3 to 5 percent at the same widths. If linear ever becomes a production
path there, that bound needs re-taking.

3. The multiple-of-tile condition is gone

Measured, and the answer is not the one I expected. Swept over 1024..1520 in steps of 16 on
14336x5120, three arms on one card — the route being replaced, the shipped predicate, and a
scratch build with every bound removed:

32 widths, 8 of them whole cp.async tiles. The condition declined the other 24:

what it declined	widths	model would have	route measured	so the condition
1040..1328, not whole tiles	15	declined anyway	1.199 to 1.039	repeated a decision
1360..1520, not whole tiles	9	taken	0.980 to 0.920	blocked a real gain

The 8 whole tiles were never its business: the model takes 1024, 1408 and 1472 and declines 1088,
1152, 1216, 1280 and 1344 on its own.

It never rescued a width the model would have wrongly taken, and it cannot: the model's only
width-dependent inputs are ceil(T/256) and ceil(T/64), both constant between adjacent multiples
of 64, so the verdict is already fixed across each band and the condition can only subtract from a
decision already made. So it was a restriction imposed to simplify memcmp, with no performance
reason behind it, and I have removed it rather than re-justified it.

Confirmed on the shipped predicate rather than inferred. Re-swept after the removal, the same 32
widths split 20 declined and 12 taken: the 20 the model declines measure 0.997 to 1.001, which is
both arms running the same kernel, and the 12 it takes measure 0.9195 to 0.9926 — 1024 as before,
and 1360 through 1520, which the condition used to block.

That admits between 1134 and 5292 further widths per geometry over 1024..8192. They are not
byte-comparable with the previous kernel by construction, so the numerical tests cover them against
their host reference instead: 1345 on 14336x5120, 1153 on 16384x5120 and 34816x5120, 4001 on
both residual shapes, added to test_fp8_a8.cpp, linear_add/test_fp8.cpp,
linear_swiglu/test_fp8.cpp, test_attn_input_proj.cpp and test_gdn_input_proj.cpp.

The consequence for the description: this change is no longer bitwise identical to the route it
replaces, and the title no longer says so. Where the width is a whole cp.async tile the outputs
still match byte for byte, and I report that below as supporting evidence rather than as a contract.

4. mbarrier helpers

Removed. The route uses cta_mbarrier_init / _wait / _arrive / _arrive_expect_tx and
cta_mbarrier_fence_init from ops/common/mbarrier.cuh.

Worth recording why this was more than redundancy: the copies used plain labels (wait_loop:,
wait_done:) where cta_mbarrier_wait uses %=-suffixed ones. Two waits inlined into the same
function would have been a duplicate-label assembly error, waiting on an inlining decision to
change.

The description of what changes

You are right that "the operand copy is staged differently" is not the change. Four things move
together and only the last is about TMA:

	cp.async route	TMA route
output tile per CTA	64 tokens x 128 rows	256 tokens x 128 rows
consumer warp tile	32 x 32	64 x 64
accumulator registers per consumer thread	32	128
K pipeline	2 stages of K=128	4 stages of K=64
buffered K extent	256	256 (unchanged)
threads per CTA	256	288 (8 consumer warps + 1 producer)
CTAs per SM	2	1
registers per thread	94	166
K-loop synchronisation	CTA-wide barrier per stage	producer/consumer mbarrier pair

The wider warp tile is the arithmetic change: an ldmatrix of A now feeds eight N fragments instead
of four and a B fragment feeds four M fragments instead of two, so the MMA issued per operand load
goes from 8 to 32. The deeper, narrower pipeline and the dedicated producer are what let one CTA
keep that tile fed without the per-stage CTA-wide barrier.

The same table is the cost. Resident warps per SM fall from 16 to 9, so a shape with too few token
tiles to fill the machine pays that and gets none of the reuse back — which is the mechanism behind
the width floor, and what the model has no term for.

Claims withdrawn

• "the traffic is identical by construction" — withdrawn. Unchanged logical inputs do not
  establish unchanged physical traffic, and I did not measure traffic.
• derived bandwidth — the GB/s figures were assumed bytes divided by time, not a DRAM counter.
  They are gone; where a share of peak appears it is the column ninfer_linear_bench publishes,
  named as such, and still a model rather than a measurement.
• the roofline conclusion — withdrawn. The peak denominator was not re-established on this
  machine and the cold-operator bench does not reproduce the inference-time picture, so I make no
  claim about which resource binds.
• "the kernel SASS hash is identical across both rebases" — withdrawn, because the instrument
  does not measure that. Building one unchanged source three times gives three different hashes of
  the cuobjdump -sass output for fp8_a8.cu, so the comparison cannot separate a source change
  from a recompile. Where the previous revision used it as evidence, this one re-runs the suite.

None of that touches the elapsed times.

Rebase

Four files conflicted on 699adfce and all four were resolved by taking your version and
re-applying only this branch's additions. One is worth flagging: the six-schedule ladder you added
to fp8_attn_input_a8.cu did not exist on the old base and the submitted branch would have deleted
it. It is kept in full, and the TMA route is now the wide arm of that ladder, placed beside
Prefill — the schedule the predicate compares against. The three test files likewise keep your
cases and append; the submitted branch had narrowed the linear_swiglu case lists, which was wrong
of it. 487f8977 then merged cleanly: its two files and this branch's fifteen are disjoint.

One thing that is not mine to fix

ninfer_attn_input_proj_test fails on 487f8977 without this branch. Every failing case has the
form attn q|k|value W8 DFlash2 A16 T=<N> graph phase=1: reduction criterion failed, for T drawn
from the list run_w8_dflash2() replays, and only in phase=1 — the second pass, where the
activation changes sign and the captured graph is replayed again. The first pass is clean.

The number of failing cases is not stable across runs — 34 in one, 36 in another — while their form
never varies. That instability is itself a fact about the defect rather than noise in the harness,
and it points at the replay rather than at the arithmetic. --dflash2-only reproduces it in about
two minutes. I built the test from clean master and ran both arms from one build directory to be
sure it is not this branch; no FP8 case fails in either.