The fp8 A8 GEMM stages its operands through TMA: prefill +2.4% at chunk 1024 and +4.9% 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 — linear, linear_add, linear_swiglu,
attn_input_proj, gdn_input_proj — all at prefill widths.

Rebased onto b88c0f6f and re-measured there, because the base moved: ee9d5192 speeds up the
NVFP4 part of qwen3.8-27b, and that mixed artifact is the one this route runs on.

The four items

1. Swizzle alignment — a defect, not a missing sentence. The swizzle is a function of the shared
address, not of an offset from the destination. Same tile, same descriptor: an atom-aligned base is
correct, base+128 is 1024 of 1024 bytes wrong. __align__(128) happens to land on 0x400 under
CUDA 13.1, which is why the submitted code computed correct results. 512 now sits on the tensor storage, the union, the outer storage and the allocation, with
static asserts that all three tile bases the producer writes — both stage strides and a paired
block's second branch — are whole atoms. It costs 384 B of shared, sizeof 98432 → 98816 against a
101376 cap, and no time.

2. Calibration out of the header. The measurement history moved to
docs/maintainer/fp8-a8-tma-route.md, indexed from docs/README.md; the header section went 387 →
284 lines. What is left asserts coverage and says so — that the predicate still sends the test
widths down this path, so the numerical tests exercise the kernel they were extended for. Nothing
there claims the route is faster.

3. The multiple-of-tile condition is gone. Swept 1024..1520 in steps of 16 on 14336x5120: it
declined 24 of 32 widths — repeating the model's own decision on 15, blocking a measured gain on 9
(0.975 down to 0.915). It could not have rescued a width the model wrongly takes: the model's only
width-dependent inputs are ceil(T/256) and ceil(T/64), both constant across each 64-token band.
So it was a memcmp convenience, and this route is no longer bitwise identical to the one it
replaces. Widths that are a whole cp.async token tile still match byte for byte; the rest are
checked against the host reference — 1345, 1153 and 4001 added across the five test files.

4. mbarrier helpers. Removed; the route uses cta_mbarrier_* from ops/common/mbarrier.cuh.

New since your review: token-fastest rasterisation, matching ee9d5192 and the cp.async
schedules this route falls back to, which are already registered Fp8MmaRaster::TokenFast. Each CTA
computes the same tile, so the output is unchanged by construction. It is not uniformly free: at the
operator it is worth 1.5–3.3% on 34816x5120 and costs 0.2–0.5% on 5120x17408 through
linear_add; end to end it is −0.12% at chunk 1024, +0.71% at 4096, +0.97% at 8192. The direction
the numbers point in is the concurrent working set against L2 — 34816x5120 is the only geometry
whose weights, at 178 MB, do not fit. Choosing the order per geometry from that working set would
recover the 0.2–0.5%; it is written up in the document rather than taken here, because it is a
decision for three routes and not only this one.

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 feeds eight N fragments instead of
four, 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 — which is why narrow widths are
declined, and what the cost model has no term for. Per SM the route uses less of both budgeted
resources than the one it replaces: 98816 B of shared against 100352, and 166 x 288 registers
against 94 x 512.

Measured

One RTX 5090, sm_120a, 170 SMs, 525 W, CUDA 13.1. Clock locking is not available and the card idles
at 180 MHz, so the arms alternate inside each pass and a fourth arm — the base under a second label
— carries the instrument's zero.

Operator, ratio of this route to the route it replaces, median of three passes. The null arm
spans 0.989 to 1.003; † marks a width the predicate declines, where both arms run the same kernel:

shape	call site	T=1024	T=4096	T=8192
5120x17408	linear_add	0.895	0.936	0.931
5120x6144	linear_add	0.933	0.979	1.000 †
5120x17408	linear	0.877	0.915	0.889
5120x6144	linear	0.885	0.925	0.998 †
14336x5120	linear	0.989	0.950	0.952
16384x5120	linear	1.000 †	0.975	0.952
34816x5120	linear	1.000 †	0.937	0.912

End to end. Only qwen3.8-27b carries FP8_E4M3FN_ROW_BF16S weights — qwen3.6-27b is nvfp4 and
qwen3.6-35b-a3b groupwise-int end to end — so it is the only artifact this route runs on.
ninfer_bench -p 16384, five passes in mirrored order:

prefill chunk	base, tok/s	this branch, tok/s	change	null arm
1024	8595.5	8801.0	+2.40%	−0.03%
4096	9220.4	9677.9	+4.95%	−0.03%
8192	9184.0	9588.3	+4.38%	−0.04%

Each change is the median of the five paired ratios; past the first pass they spread under 0.1%.

Decode is not on this path, and that is a census. -n 512 with --cuda-graph-trace=node: 29
distinct kernels execute and fp8_a8_tma_kernel is not among them, 0 launches. Decode throughput
over four alternating repetitions ranges 0.9993 to 1.0002.

Suite and tools on the tree this pushes: ctest 114 of 114, and the same on clean b88c0f6f,
both built and run on one machine with one toolchain. compute-sanitizer: memcheck 0 errors at 1024, 1345,
4096 and 4288 on 14336x5120 and at 4001 on 5120x17408 — the last three reach this route only
since item 3; racecheck, synccheck and initcheck 0 at 1024.

Your b88c0f6f also closed the ninfer_attn_input_proj_test failure the previous revision reported
as not mine to fix. It passes on both arms now.

What this does not claim

• Not that the memory traffic is unchanged. Unchanged logical inputs do not establish unchanged
  physical traffic, and I did not measure traffic.
• No bandwidth derived from assumed bytes over time. Where a share of peak appears it is the column
  ninfer_linear_bench publishes, which is a model.
• Nothing 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.
• Not bitwise identical to the route it replaces — item 3.

The routing constant is fitted for this part and is not portable; what travels is the shape of the
comparison. The sweeps behind it, behind both width bounds, and behind the rasterisation choice are
in docs/maintainer/fp8-a8-tma-route.md.

🤖 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-09T19:07:31.573516Z	f0c5ed8	Manual request

ℹ️ 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.

[MichaelDementii]

Rebased onto b88c0f6f and re-measured there; the body is rewritten around the current numbers
rather than appended to, so it and the commit are the whole story now.

All four items are closed: 512-byte alignment on every tile base the producer writes — it was a
defect, not a missing sentence, and the measurement is in the body — calibration moved to
docs/maintainer/fp8-a8-tma-route.md, the multiple-of-tile condition removed together with the
sweep that justifies removing it, and cta_mbarrier_* reused. One thing is new since your review:
token-fastest rasterisation, matching ee9d5192 and the cp.async schedules this route falls back
to. It is not uniformly free and the body gives both numbers.

Prefill on qwen3.8-27b: +2.40% at chunk 1024, +4.95% at 4096, +4.38% at 8192, with a fourth arm —
the base under a second label — reading −0.03%. ctest 114 of 114 on this tree and on b88c0f6f.

Your b88c0f6f also closed the attn_input_proj graph-replay failure the previous revision
flagged as not mine to fix; it passes on both arms now.

Ready for another look whenever suits you.

[MichaelDementii]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. You're on a roll.

Reviewed commit: f0c5ed88b9

ℹ️ 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".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".