B300 tcgen05 gemm rs / ag gemm - #35
Open
monopodium wants to merge 26 commits into
Open
Conversation
Collaborator
Author
|
8× B300, single node, BF16 gemm_rs — M×K×N with K = M/8, N = M; FLOPs = 2·M·N·K
ag_gemm — M×K×N with K = M, N = M/8; FLOPs = 2·M·K·N
|
The Hopper wgmma path cannot assemble on sm_103a (ptxas rejects wgmma.fence
and wgmma.mma_async). Port gemm_rs's MMA to Blackwell tensor cores, behind
GPU=blackwell / MKERNEL_TCGEN05 so the sm_90a path is untouched.
Vendored from ThunderKittens, following the existing tk_-prefixed layout:
tensor-memory types (tt, tensor_allocator), thread- and group-scope tcgen05
MMA, tmem->register loads, and the tcgen05 sync primitives that the original
subset left as an empty section. Also restores the packed int/fp8e8m0/fp4
aliases and cuda_fp4.h those headers name.
Kernel changes, all in the compute path:
- A_tile becomes a single 128-row tile, since tcgen05 issues one M=128 MMA
for the whole row block instead of two 64-row wgmma halves. Same bytes,
one TMA load instead of two.
- A dedicated MMA warp (producer warpgroup, warp 2) issues the tcgen05
MMAs into a tensor-memory accumulator; consumers no longer do MMA.
- Consumers wait on the accumulator, pull their 16-row slice out of tmem
using TK's group<8> row mapping, release tmem, then stage to shared for
the existing epilogue TMA warp. Shared/global layouts are unchanged.
- Two new semaphores: mma_done and tmem_free.
Bench: allow single-node runs. With n_send == n_reduce == 0 the kernel
already takes its intra_only_debug path and never dereferences the session,
so skip create_session entirely for NUM_NODES=1, and undo the chunk-major
staging view before the correctness check (the reduce CTAs that normally
publish the output buffer do not launch).
Verified on 8x B300 (sm_103a): SASS carries UTCHMMA and no wgmma;
correctness PASSes at M=2048/4096/8192/16384; 1073 TFLOP/s at M=32768.
Same pattern as the gemm_rs port, behind the same GPU=blackwell /
MKERNEL_TCGEN05 switch; the sm_90a wgmma path is untouched. No new
ThunderKittens vendoring was needed - the tcgen05 headers brought in for
gemm_rs cover this too.
Compute path:
- A_tile becomes a single 128-row tile (A_comm_tile, the ring all-gather
granularity, is unchanged), so both the local and the A_recv load
coordinates drop their *2+i indexing.
- Producer warpgroup warp 2 becomes the MMA issuer. Unlike gemm_rs this is
a persistent kernel, so the warp replays the loader's task loop verbatim,
including the identical remote-skip, to stay in lockstep with the input
pipeline.
- Consumers wait on the accumulator, pull their 16-row slice out of tmem
with TK's group<8> row mapping, release tmem, then stage to shared for
the existing epilogue TMA warp. Shared and global layouts are unchanged.
- Two new semaphores: mma_done and tmem_free.
Two single-node fixes, both independent of the MMA port:
- The ring receive area is empty when there are no peers, and TMA
descriptors cannot be built over a zero-extent tensor. Clamp A_recv to
one A_comm_tile of rows on both the host and bench sides. The kernel
never reads it: with num_nodes == 1 every task decodes to the local
shard, so is_remote is never true.
- Skip the phase-0 prologue when num_nodes == 1. It posts RDMA work
requests to the d2h FIFO, and a caller with no session has null FIFO
handles, so it faulted with an illegal access. There are no peers to
post to either way.
Bench: skip create_session for NUM_NODES=1, as for gemm_rs.
Verified on 8x B300 (sm_103a): SASS carries UTCHMMA and no wgmma;
correctness PASSes exactly (max_abs 0) at M=4096/8192/16384/32768;
725 TFLOP/s at 32768x32768x4096.
Tensor memory is 512 columns and a 128x256 fp32 accumulator costs 256, so
carry two and alternate. The MMA warp can then start the next tile's MMAs
while the consumers are still draining the previous accumulator, instead of
blocking on tmem_free. mma_done/tmem_free become rings of the same depth,
each ring slot tracking its own phase bit.
ag_gemm also goes from 3 to 4 pipeline stages. At 48 KB per stage plus 64 KB
of outputs aliased onto the last one, 3 stages left 48 KB of the 227 KB
shared budget unused; 4 stages lands at 208 KB, the same as gemm_rs.
Measured on 8x B300, best of default and 148 comp SMs:
gemm_rs 2048 +7.5% 4096 +3.3% 8192 +4.7% 16384 +9.0% 32768 +1.4%
ag_gemm 4096 -1.8% 8192 -0.2% 16384 +1.5% 32768 +4.5%
Modest, and the reason is that tmem was not the binding constraint. Two
experiments pinned that down:
- Giving outputs their own shared allocation removes the loader's
per-tile wait on outputs_finished, which is the edge that actually
starves the MMA warp. But at 4 stages that needs 256 KB, over the 227 KB
limit, so it only fits at 3 stages - and 3 stages costs ~18% at M=32768.
- Running 3 stages with outputs still aliased reproduces the same ~18%,
confirming the regression is pipeline depth, not the aliasing.
So input pipeline depth dominates, and the remaining gap to ThunderKittens'
b200 kernels is most likely arithmetic intensity: per CTA tile this kernel
reads 3 MB to do 268 MFLOP (85 FLOP/byte), while TK's 2-CTA cluster covers
CLUSTER_M=512 x 256 and reaches ~171 FLOP/byte. Closing that means either a
second accumulator per consumer warpgroup (two row blocks per CTA, sharing
one B tile) or real 2-CTA clusters - both of which want the same 512 tmem
columns this ring now uses.
Correctness unchanged: gemm_rs PASSes at M=4096/16384/32768, ag_gemm PASSes
exactly (max_abs 0) at M=4096/16384.
Replaces the tensor-memory ring. Tensor memory holds exactly two 128x256 fp32 accumulators; spending them on two row blocks of the same task instead of on double-buffering one lets a single B tile feed both MMAs, which is what actually moves the needle. Arithmetic intensity per CTA task at K=4096: before out 128x256 read 3.1 MB 268 MFLOP 85 FLOP/byte now out 256x256 read 4.2 MB 537 MFLOP 128 FLOP/byte TK b200 out 512x256 read 6.3 MB 1074 MFLOP 171 FLOP/byte Measured, best of default and 148 comp SMs, vs the ring it replaces: 2048 -7.0% 4096 -0.2% 8192 -0.9% 16384 +0.5% 32768 +22.7% At 32768x4096x32768 that is 6.586 ms / 1336 TFLOP/s, up from 8.081 ms / 1089, and 0.80x of ThunderKittens' b200 kernel (was 0.65x). Small shapes give back a little because a task now spans twice the work and there are half as many to spread over the CTAs. The pairing falls out of the existing decode: tasks are laid out as (row block within slice, slice, column) with column varying fastest, so the partner of task_id is task_id + tiles_per_round - same slice, same column, next row block. Same slice also means the same owner device, so the staging math is untouched. Row blocks per device slice must now be even; checked on the host. Pipeline stages drop 4 -> 3 because a stage now carries two A tiles plus the shared B (64 KB), landing at 192 KB. Each stage feeds twice the MMA work, so the shallower pipeline still covers more latency than the old 4. The shared output tile stays single-buffered, so the two row blocks drain through it in sequence: outputs_free hands it back to the consumers after the first, outputs_finished releases the whole task to the loader after the second. Both store-warp rounds wait on the same outputs_arrived semaphore and so share one phase bit, and outputs_free needs a low-half bit because it must start unsignalled - getting either wrong deadlocks. Correctness PASSes at M=2048/4096/8192/16384/32768.
tma::cluster::load_async for tiles is the one cluster primitive mKernel's
ThunderKittens subset was missing; tma::cluster::{wait,expect,arrive,sync},
cluster_ctarank(), everyone, tensor_allocator<1,2> and the mm2_*/mma2_*
wrappers are all already present. Retyped onto mKernel's
detail::tma_layout_for descriptor concept, since mKernel passes
dist::local_tensor where upstream passes gl.
Unused for now - the 2-CTA conversion itself is not in this commit. Recorded
here so the next step starts from a tree that has every primitive it needs.
Also pins down the ncta=2 shape semantics empirically, which the source alone
reads ambiguously. With mm2_ABt and A=st_bf<128,64>, B=st_bf<128,64>,
D=tt<float,128,256> the instruction type-checks, so:
M = A::rows * ncta = 256 - A is split by rows across the CTA pair,
each CTA keeps its own 128 accumulator rows
N = B::rows * ncta = 256 - B is stored N-major as st<N/2, K> and split
by columns across the pair
So a cluster computes 512x256 from A read once (512 rows) and B read once
(256 columns): 6 MB per cluster task at K=4096, 171 FLOP/byte, versus 128 for
the current single-CTA paired-row-block kernel. That is the whole of the
remaining gap to upstream's b200 kernel.
The catch this exposes: mm2_ABt wants B as (N, K). mKernel takes B as (K, N)
everywhere - globals, host entrypoint, bench allocation and the correctness
reference - so the conversion is not confined to the compute path.
Step one of the 2-CTA conversion, kept separate so the layout change is validated on its own. mm2_ABt derives its N extent from B::rows * ncta, so the CTA pair can split B by columns and read it once per cluster - which requires B stored (N, K). Switch the Blackwell path to that layout now, still single-CTA with mm_ABt and a full-width st_bf<COL_BLOCK, RED_BLOCK> tile. B.cols() meant N on the K-major layout; it is B.rows() here. The bench keeps one logical B and hands the kernel whichever layout the build wants, so the correctness reference is unchanged. Hopper is untouched. Same bytes per stage, and no measurable cost: M=16384 1.328 ms (was 1.326), M=32768 6.692 ms (was 6.586). Correctness PASSes at M=4096/16384.
Completes the conversion. A cluster of two CTAs computes 512x256 per task:
mm2_ABt splits A by rows across the pair and B by columns, so each is read
once per cluster instead of once per CTA - 171 FLOP/byte at K=4096, up from
128.
Structure follows upstream's b200 kernel:
- Both CTAs load their own slice (cluster_mask = 1 << ctarank) but direct
the transaction count at CTA 0, so the leader's single wait covers the
whole cluster's inputs. Needed one more vendored TMA variant for that.
- The leader's warps 2 and 3 issue one mm2_ABt each, one per accumulator.
Each issues its own cluster::expect, which is where inputs_arrived's two
arrivals come from and how the six redirected tile loads are accounted.
- commit<2> multicasts, so both CTAs' inputs_finished and mma_done fire and
the follower's loader and consumers stay in step without a leader-side
broadcast.
- Consumers and the epilogue are unchanged apart from the row-block mapping,
which becomes row_idx + 2*ctarank + h.
- __cluster_dims__ on the kernel, grid rounded down to whole clusters, and a
cluster-wide sync before the tensor allocator so no peer redirects a
transaction at a barrier that does not exist yet.
Row blocks per device slice must now be a multiple of 4, so M=2048 no longer
builds a valid launch; checked on the host.
Measured, best of default and 148 comp SMs, against the single-CTA version:
4096 +5.6% 8192 -8.9% 16384 -7.8% 32768 +3.5%
So it wins only at the largest shape - 6.365 ms / 1382 TFLOP/s, 0.83x of
upstream (was 0.80x) - and loses in the middle, where having the leader issue
every MMA while the follower's issue warps idle costs more than the halved
operand traffic saves. The intensity model keeps over-predicting: 128 -> 171
FLOP/byte is 1.33x on paper and 1.035x in practice, the same way 85 -> 128
predicted 1.5x and delivered 1.23x. L2 is evidently already capturing much of
the operand reuse that the model charges to HBM.
Correctness PASSes at M=4096/8192/16384/32768.
The 2-CTA conversion regressed the middle shapes. A roofline built from known bytes and measured time says why, and it is not what the last two rounds assumed. At M=32768, per GPU: working set A+B 537 MB -> 0.08 TB/s of HBM. Not remotely bound. 1-CTA paired 68.7 GB requested -> 10.43 TB/s L2->SM, 1336 TFLOP/s 2-CTA cluster 51.5 GB requested -> 8.10 TB/s L2->SM, 1382 TFLOP/s upstream b200 51.5 GB requested -> 9.75 TB/s L2->SM, 1664 TFLOP/s This kernel and upstream now run the identical 512x256 cluster tile, so they request identical bytes for identical FLOPs, and upstream is still 1.2x faster. The remaining gap therefore has nothing to do with arithmetic intensity - which is also why the last two intensity improvements both came in far under their predicted speedup. It is prefetch depth: keeping the tensor cores fed, not moving fewer bytes. Upstream runs 4 pipeline stages. This kernel dropped to 3 back when B was full width and a stage cost 64 KB. The 2-CTA split halved the B tile, so a stage is 48 KB again and 4 stages fit in 208 KB - the same budget and depth upstream uses. Restoring it is a one-line change. Best of default and 148 comp SMs, against the 3-stage cluster build and the single-CTA build it replaces: 4096 0.399 -> 0.378 -> 0.396 ms 8192 0.574 -> 0.630 -> 0.568 ms 16384 1.326 -> 1.438 -> 1.286 ms 32768 6.586 -> 6.365 -> 6.180 ms So the regressions are gone and M=32768 lands at 6.18 ms / 1423 TFLOP/s, 0.86x of upstream, up from 0.80x before the cluster work. M=16384 measured 1.286/1.412/1.307 ms across three runs, so the earlier apparent 3.5% loss there was inside run-to-run spread; M=32768 measured 6.194/6.185/6.175, so its gain is not. Also adds MKERNEL_PROFILE_NO_BARRIER, which drops the end-of-kernel cross-GPU barrier. Nsight replays the kernel on one rank while the others run once, so that barrier deadlocks under the profiler. Profiling builds only; results from such a build are meaningless. Correctness PASSes at M=4096/8192/16384/32768.
The roofline said the kernel moves the same bytes as upstream but more slowly,
so the remaining gap is keeping the tensor cores fed. The edge that starves
them is the loader's per-task wait on outputs_finished, which exists only
because the output tile aliases the last input stage.
Two things had blocked removing it. At 64 KB of outputs it only fit by
dropping to 3 stages, and that cost more than the gate. And the group<8> tmem
layout spreads each warpgroup across all 128 rows, so the two 64-row halves
of a row block are written by warps {0,1,4,5} and {2,3,6,7} - neither a
warpgroup - which is why the epilogue staged both halves at once.
Both are answered by staging one 64-row half at a time. That halves the
buffer to 32 KB, which fits alongside four 48 KB stages at 224 KB, and the
half-selection is just a predicate on the warp's tmem row with every warp
still riding the shared barrier. The epilogue becomes four sub-rounds per
task (two row blocks x two halves); outputs_free hands the tile between them
and outputs_finished closes the task.
Also releases each accumulator's tensor memory as soon as its registers are
read rather than at the end of the epilogue, matching upstream. Measured on
its own that changed nothing, which is consistent with the diagnosis: the MMA
warp was never blocked on tmem, it was blocked behind the loader.
Best of default and 148 comp SMs, M=32768 over five runs
(5.956/5.923/5.939/6.098/5.956, median 5.956):
4096 0.396 -> 0.404 ms
8192 0.568 -> 0.573 ms
16384 1.286 -> 1.280 ms
32768 6.180 -> 5.956 ms 1477 TFLOP/s, 0.89x of upstream
The small shapes move within their run-to-run spread; M=32768 is the real
gain. Correctness PASSes at M=4096/8192/16384/32768.
The epilogue wrote every output tile twice: once into workspace and once as the cross-GPU store_add into the owner's staging. workspace is only ever read back by fused_comm_tile_impl, which belongs to the inter-node reduce path, and that path does not launch when there is no session. So on a single node half the epilogue's TMA stores and half its global write traffic were dead - 128 KB per task. Gated on G.rt != nullptr, which is exactly the condition under which the reduce CTAs exist. Multi-node behaviour is unchanged. Best of default and 148 comp SMs; M=32768 over five runs (5.774/5.767/5.764/5.753/5.766, median 5.766): 4096 0.404 -> 0.392 ms 8192 0.568 -> 0.555 ms 16384 1.280 -> 1.285 ms 32768 5.956 -> 5.766 ms 1525 TFLOP/s, 0.92x of upstream This came out of looking for a way to let the follower CTA share the MMA issue, which turned out to be a dead end twice over. tcgen05.mma.cta_group::2 must be issued by the leader of the pair - upstream guards it the same way - so the follower structurally cannot. And it would not matter: a K step takes about 1574 cycles of tensor core work and carries two MMA instructions, so issue is on the order of 0.1% of the time and the leader's issue warps are idle for the rest of it. Correctness PASSes at M=4096/8192/16384/32768.
…line Brings across the two changes that paid off on gemm_rs. One task now covers two adjacent row blocks sharing a single B tile, and the epilogue stages one 64-row half at a time so outputs can have their own allocation instead of aliasing the last input stage and gating the loader. The pairing falls out of decode_comp_task: it varies rb fastest within a band, so flat and flat+1 are adjacent row blocks at the same column. SUPER_M is 12 and final_rows is 4 or 8 for every shape here, all even, so flat = 2k never straddles a band or the super/tail boundary and rb always comes out even - which also puts the pair inside one 256-row intra block, leaving the per-K-strip barrier waits untouched. Shared memory: a stage carries two A tiles plus the shared B at 64 KB, and 3 stages plus a 32 KB staging tile is 224 KB. Staying at 4 stages would need 256 KB, but each stage now feeds twice the MMA work so the shallower pipeline still covers more. Two runs, essentially identical: 4096 0.167 -> 0.170 ms 8192 0.598 -> 0.544 ms +9.9% 16384 2.379 -> 2.233 ms +6.5% 32768 11.610 -> 10.553 ms +10.0% M=32768 reaches 833 TFLOP/s, 0.53x of upstream, up from 0.48x. Correctness PASSes exactly (max_abs 0) at M=4096/8192/16384/32768.
…comm-SM table on Blackwell Two problems, both in the harness rather than the kernel, and together worth more than every kernel change made to ag_gemm so far. The per-shape num_comm_sms override assigned into args inside the shape loop and never restored it, so a shape absent from the table inherited whatever the previous shape had set. M=32768 is not in the table, so in any multi-shape sweep it ran with 8 comm CTAs instead of the 64 default: 10.6 ms rather than 8.6. Every ag_gemm number reported from a multi-shape run was measured that way. Reset before the lookup. The table itself is Hopper tuning - it exists to cut the comm-CTA budget at small M, where the fused path otherwise loses to NCCL. On Blackwell the tcgen05 compute path is fast enough that the intra-node all-gather becomes the critical path instead, and 8 comm CTAs starve it badly. Sweeping the count at M=32768: 8 -> 10.58 ms, 16 -> 8.67, 32 -> 8.67, 64 -> 8.66. Same story at the shapes the table does cover: at 8192, 0.551 -> 0.370 -> 0.325 -> 0.326; at 16384, 2.238 -> 1.437 -> 1.349 -> 1.348. Flat past 32, so keep the 64 default on that path and leave Hopper alone. Medians of three runs of 30 iterations, against the previously reported numbers: 4096 0.170 -> 0.170 ms 8192 0.547 -> 0.326 ms +67.8% 16384 2.227 -> 1.349 ms +65.1% 32768 10.595 -> 8.637 ms +22.7% ag_gemm now sits at 0.54x / 0.77x / 0.67x / 0.65x of upstream, from 0.54x / 0.46x / 0.41x / 0.53x. Correctness unchanged - still exact at every shape.
Adds a per-task activity trace (AG_GEMM_TRACE=1) and fixes what it found.
Each record carries %globaltimer at task start/end -- device-wide, so tasks on
different SMs land on one axis -- plus clock64() deltas accumulated inside the
task for stall accounting, which is nearly free and correct within one SM.
busy_clk lets each record self-calibrate ns per tick.
At M=32768 the trace showed compute CTAs resident and looping from t=0 but
spending 86-92% of the first 3.3 ms stalled, doing 0-10% useful work, while all
8 gather CTAs ran flat out. The overlap existed on paper and delivered nothing.
The cause is a task-order mismatch. A GPU gathers only its own 1/8 shard --
intra rows [d*L, d*L+L) in row order -- and all 8 run concurrently, so readiness
sweeps the row space as {0,L,2L,...}, then {1,L+1,...}. Compute's first wave of
140 CTAs asked for rows 0..139, of which only 8 could be ready.
Relabelling row r as (r % NUM_DEVICES) * L + (r / NUM_DEVICES) makes the k-th
row consumed the k-th row produced. It is a bijection, and independent of
dev_idx, so the lockstep tile order the multicast read sharing depends on is
preserved -- unlike the device-rotation experiment, which regressed.
Compute's useful-work share goes 30.8% -> 47.4%, readiness stall 42.7% -> 24.3%:
M=16384 1.321 -> 1.051 ms 0.69x -> 0.87x of TK
M=32768 8.697 -> 6.540 ms 0.64x -> 0.86x of TK (1011 -> 1345 TFLOP/s)
Bit-exact at every shape (max_abs = 0).
Also retires a wrong conclusion: the "gather-CTA cliff" (8 -> 16 CTAs costing
2.5x) was an artifact of the bad ordering and is gone -- 6.612 vs 6.607 ms.
Two follow-ups in the same area:
AG_GEMM_FASTPOLL (on for GPU=blackwell) skips the per-K-strip readiness poll
once phase 1 is globally complete, since every plane-0 slot is then set.
256 global loads per task become one. Worth 0.9%.
AG_GEMM_OWNSHARD (off) skips the wait for rows the device owns, which are
already in place. Correct but slower -- it desynchronises the cross-device
tile order. Kept disabled per the repository's convention for failed flags.
…lock
Takes ag_gemm's compute path to the same 2-CTA shape gemm_rs uses: A split by
rows across the CTA pair, B taken N-major and split by columns, one MMA warp
per accumulator on the leader, cluster-scope TMA loads redirected at the
leader's barrier, commit<2> multicasting the stage release back to both CTAs.
M=16384 1.051 -> 0.996 ms 0.87x -> 0.91x of TK
M=32768 6.540 -> 6.518 ms 0.86x of TK (1350 TFLOP/s)
M=8192 0.325 -> 0.320 ms
Bit-exact at every shape (max_abs = 0).
An earlier attempt at this hung nondeterministically and was reverted without a
diagnosis. The hang is real and it is one line of placement:
~tensor_allocator() { if (warpid()==0) { if (ncta==2) barrier.cluster...; } }
At ncta 2 the destructor runs a cluster barrier, so every CTA of every cluster
has to construct and destroy the allocator. gemm_rs does that -- its allocator
lives at fused_kernel scope, above the role dispatch, so its comm CTAs build one
too. ag_gemm's sat inside fused_comp_sm, so gather CTAs never reached the
barrier and their cluster partners waited forever. Moved to fused_kernel scope.
Finding it needed a tool, since a deadlocked kernel never returns its trace:
AG_GEMM_TRACE now also keeps a live progress array, written in place as work
advances and read back with cudaMemcpyFromSymbolAsync on a non-blocking stream.
The copy engines run independently of SM execution, so it reports where every
warp role is stuck while every CTA is spinning. AG_GEMM_HANG_PROBE=1 launches
one iteration without synchronising and polls it. That walked the diagnosis from
"it hangs" to "gather completed and passed its gate", to "all four compute warp
roles exited their loops", to "no compute CTA reaches the statement after
fused_comp_sm returns" -- which leaves only the destructor.
Also required, and enforced rather than assumed:
- num_intra_comm is rounded up to a multiple of CLUSTER_SIZE. Clusters are
consecutive blockIdx.x and gather CTAs hold the low indices, so an odd
count would put a gather CTA and a compute CTA in one cluster and deadlock
the compute half on the first cluster-scope barrier.
- num_comp_sms and the launch grid are truncated to whole clusters.
- SUPER_M % ROW_BLOCKS_PER_CLUSTER == 0 (static_assert) and row blocks per
shard likewise, so a cluster's four row blocks never straddle a super-tile
band -- which is what lets one decode_comp_task call serve both CTAs.
The fourth pipeline stage that the halved B tile pays for is left off: measured
over 3 runs it is 1.5% better at M=16384 and 1.7% worse at M=32768. The default
follows the larger shape; AG_GEMM_STAGES=4 overrides.
…t on Follow-up to the 2-CTA cluster work. Re-tracing the cluster build at M=32768 moved the picture again: readiness stall is down to 1.12 ms wall-equivalent (0.83 of it before gather ends) while pipeline backpressure is now 1.84 ms, so the loader is no longer the pole -- the MMA and epilogue are. All three attempts below aimed at that, and the measurements say none of them pays. AG_GEMM_EPILOGUE_READ_WAIT: store_async_read_wait instead of the global-commit store_async_wait in the epilogue, which is what gemm_rs does. Medians of 3: 0.6% at M=32768, noise elsewhere. Not worth it against the race the original comment warns about on the inter-node path, which single-node cannot test. AG_GEMM_DOUBLE_OUTPUT: two epilogue staging tiles instead of one, so the consumers can stage the next 64-row half while the store warp drains the previous. Fits (3 stages x 48 KB + 2 x 32 KB = 208 KB). Measured neutral: M=16384 0.996 -> 0.986, M=32768 6.518 -> 6.563. AG_GEMM_ROWPOLL: a per-row completion counter on the otherwise unused barrier plane 1, so a task can drop its 256 per-K-strip polls as soon as its own row is complete -- much earlier than the global flag, which waits for the slowest GPU. Clearly worse at every shape (M=32768 6.495 -> 6.748, M=4096 0.173 -> 0.228). The signal is a cross-GPU multimem atomic and all 256 tasks of a row hit one address; the contention costs gather more than the saved polls buy compute, and gather is on the critical path early. All three are left in behind default-off flags, per the repository's convention for failed experiments, with the numbers in the comments. The default build is unchanged in behaviour, and the numbers stand: M×K×N start now TK of TK 4096x4096x512 0.173 0.171 0.092 0.54x 8192x8192x1024 0.327 0.318 0.252 0.79x 16384x16384x2048 1.321 0.986 0.910 0.92x 32768x32768x4096 8.697 6.508 5.592 0.86x Bit-exact at every shape. gemm_rs unaffected (6.245 ms at M=32768).
Generalises the ag_gemm tracer into common/mkernel_activity_trace.cuh (records
plus the live progress array, no kernel-specific assumptions) and wires gemm_rs
up to it behind GEMM_RS_TRACE=1, with one record per task per warp role so the
three can be compared directly.
gemm_rs asks a different question than ag_gemm: there is no gather on the
single-node path, and the reduce-scatter is fused into the epilogue as a
cross-GPU store_add. So the roles carry:
LOADER a = ticks blocked on inputs_finished (pipeline backpressure)
MMA a = ticks blocked on inputs_arrived (waiting for operands)
b = ticks blocked on tmem_free (waiting for the epilogue)
STORE a = ticks blocked on outputs_arrived (waiting for consumers)
b = ticks in store_add + read_wait (the cross-GPU push)
At M=32768 (span 5.83 ms, trace overhead below noise: 6.239 vs 6.245 ms):
MMA 54.3% waiting for operands, 4.5% waiting for the epilogue
STORE 82.5% waiting for consumers, 11.2% actually storing (0.63 ms wall-eq)
LOADER 78.8% blocked on the pipeline
Two things follow. The cross-GPU reduce-scatter is not the problem -- it is
0.63 ms of wall-equivalent time and the store warp is idle most of its life.
And the MMA issue warp is data-bound, not epilogue-bound.
That last point corrects a claim in README_B300.md. It said the gap to upstream
"is prefetch depth keeping the tensor cores fed". Depth is now measurable
(GEMM_RS_STAGES), and at M=32768 it is past its knee:
2 stages 8.484 ms 3 stages 6.444 ms (-24%) 4 stages 6.240 ms (-3.2%)
A 5th stage would be worth ~1% on that curve, and does not fit anyway (4 x 48 KB
+ 32 KB = 224 of 226 KB). So depth explains the earlier gains but not the
remaining 1.2x, and the operand wait is a bandwidth/locality problem rather than
a latency one: 51.5 GB of operand reads per GPU in 5.83 ms is 8.84 TB/s against
a 537 MB cold footprint, i.e. 96x reuse that L2 has to absorb.
Which points at task order. The decode walks columns fastest, so the ~60
concurrent clusters share A row blocks and stream 60 distinct B tiles. That is
the better of the two pure orders (A traffic is 34 GB against B's 17 GB), but a
super-tile swizzle that shares both -- what ag_gemm already does with SUPER_M --
is the untested lever. Left for next time; the trace can now measure it.
No behaviour change: the default build is bit-identical in configuration and
gemm_rs still measures 6.240 ms at M=32768 with correctness passing.
The comment asserted the remaining gap to upstream was prefetch depth. Depth is now measurable and past its knee at 4 stages (2: 8.484 ms, 3: 6.444, 4: 6.240), and the trace shows the MMA warp still blocked on operands 54% of the time -- a bandwidth/locality problem, not a latency one.
Net -179 lines, no behaviour change: ag_gemm measures 0.171 / 0.322 / 0.985 /
6.520 ms and gemm_rs 1.290 / 6.225 ms, both matching the previous commit, and
ag_gemm stays bit-exact.
- ag_gemm_trace.cuh 132 -> 37 lines. It duplicated the machinery that
common/mkernel_activity_trace.cuh now holds; it is a shim over it, like
gemm_rs_trace.cuh. Roles and progress slots are the only kernel-specific
part. Callers use the shared MKERNEL_TRACE_* macro spelling, and
session.cuh refers to mkernel::trace:: directly.
- Trace markers in ag_gemm.cu, 29 -> 12. Most were sub-round and phase
encodings added to localise the 2-CTA cluster deadlock; that is fixed and
documented, so what remains is one mark per role per task plus a DONE
sentinel and the kernel-tail slots -- enough to say which role is stuck.
Finer marks are cheap to re-add if another hang turns up.
- Dropped AG_GEMM_OWNSHARD. Failed experiments normally stay behind disabled
#ifdefs here, but this one made the hottest loop's condition a four-way
preprocessor chain for a variant measured slower; the finding is in
README_B300.md, which is the part worth keeping.
- The ag_gemm flag block was eight lines of boilerplate per switch. It is now
two lists and a foreach, with one comment block documenting what each flag
does and what the rejected ones measured. Verified across four
configurations that the resulting -D set is unchanged.
monopodium
force-pushed
the
b300-tcgen05-gemm-rs
branch
from
August 23, 2026 03:49
f2906a2 to
7d58e75
Compare
Net -117 lines on top of the previous tidy. gemm_rs_trace.cuh and ag_gemm_trace.cuh were 41 and 37 lines that existed only to name a handful of constants, roughly 40% of it comment. Both now live in the operator headers they belong to, next to the other per-kernel constants, and the files are gone. That leaves this branch adding six files to main, five of which are vendored ThunderKittens headers mirroring upstream's own layout; the sixth is common/mkernel_activity_trace.cuh, which both kernels share. Comments trimmed where I had been writing paragraphs: the trace header's block comment 28 -> 14 lines, the row-permutation rationale 17 -> 7, the ag_gemm and gemm_rs pipeline-depth notes 9 -> 5 and 16 -> 9, the Makefile flag table 26 -> 8. The measured numbers and the non-obvious "why" stay; the prose around them moved to README_B300.md, which is where it belongs. Verification is partial and worth stating plainly. gemm_rs's device code comes out byte-identical. ag_gemm's differs by 27 bytes -- but a control build of identical source differs by 54, so nvcc is not deterministic here and the test cannot resolve a difference that small either way. All five build configurations compile. A benchmark run is still owed: the box is busy with another user's job.
Net -873 lines. types.cuh states the rule this branch had been ignoring: "This
release only exposes the ThunderKittens types used by the five kernels." The
tcgen05 port vendored ~2100 lines of TK to reach ten entry points.
What the two kernels call, in full: mm_ABt, mma_ABt, mm2_ABt, mma2_ABt,
warp::load_async(rt, tt), tensor_load_wait, tensor_commit, tensor_allocator, tt,
and subtile. Everything else was dead.
tk_ops_thread_mma_tcgen05.cuh 985 -> 605 the _AB / _AtB / _AtBt
orientations, the fp8
scale-factor overloads,
and mma*_chunk
tk_ops_group_mma_tcgen05.cuh 236 -> 70 same set, group scope
tk_ops_group_memory_tile_tensor_to_reg.cuh 571 -> 321 store_async (reg -> tmem
is never used)
tk_types_tensor_tt.cuh 135 -> 105 30 tt_* aliases, none
referenced -- the
kernels write
tt<float, R, C> directly
tk_ops_thread_memory_tile_shared_to_tensor.cuh deleted -- it held only
load_mxnv_scale_async,
an fp8 scale loader
nothing calls
tk_common_base_types.cuh 8 of 12 added aliases dropped, now
unreferenced once the fp8
and fp4 paths are gone
Verification. All five Blackwell configurations compile, which for pruned
templates is close to proof: a definition that is needed fails to resolve.
Overload resolution could in principle shift silently, so the device code was
diffed too -- gemm_rs comes out byte-identical, ag_gemm differs by 54 bytes,
which is exactly nvcc's measured noise floor here (two builds of identical
source differ by the same amount).
gemm_ar and ring_attention do not build, but they do not build on main either:
upstream's tk_ops_group_mma_tcgen05_bf16.cuh, added with dispatch_gemm on
Blackwell, does not compile for sm_90a. Pre-existing, and not this branch's to
fix.
A benchmark run is still owed; the box has had two other users' jobs on it.
31% at M=4096 (0.171 -> 0.119 ms, 0.54x -> 0.77x of TK), nothing elsewhere. Bit-exact, three clean runs. Tracing M=4096 -- which nobody had done, all the earlier work looked at M=32768 -- shows a different kernel from the large shapes: compute does 11.6% useful work and waits 69.5% gather runs 0 -> 0.121 ms of a 0.208 ms span, on 4 CTAs of 148 CTAs launched, 4 gather, 32 compute, 112 do nothing (there are only 16 cluster tasks at this shape) So there is not enough compute to hide gather behind. Shortening gather is what matters and handing CTAs back to compute buys nothing -- the opposite of what the adaptive table assumed, and it was capping gather at 4 CTAs. Sweep: gather CTAs 4 8 16 32 48 M=4096 0.170 0.136 0.121 0.129 - M=8192 - 0.373 0.324 0.316 0.353 M=16384 - 1.265 0.985 1.071 1.182 M=16384 and M=32768 were already at their optima; only the two small-M caps move, and only on the MKERNEL_TCGEN05 path. The Hopper values are untouched with their original comments, since the reasoning there was about a different kernel shape and I cannot test it. Also adds a trace record for the kernel tail (role 2, AG_GEMM_TRACE builds), which is what showed the tail costs 6.9% of an iteration at M=4096 and 3.1% at M=32768 -- and that the cost is grid_sync waiting on straggler CTAs, not the barrier clearing itself, whose median is 0.5 us.
…tens
The two benches were timing differently and their flags meant opposite things
(gemm_rs_bench.py:368 vs ag_gemm_bench.py:393), so no ratio between them, or
against TK, was sound. MKERNEL_BENCH_TIMING now selects the mode and means the
same thing in both:
batch (default) warm up, then N back-to-back launches inside one event pair,
divided by N. No per-iteration host work, no barrier, no
sync -- exactly TK's common.py:benchmark_no_l2_clear.
periter each launch bracketed and reset between. Leaves idle gaps
that let the GPU boost, so it is NOT comparable to TK.
The root problem was that timing and correctness were coupled: the correctness
check read whatever state the timed loop happened to leave. gemm_rs cannot
survive that, because its epilogue store_adds into staging and needs the buffer
zeroed each iteration, so a back-to-back loop necessarily accumulates. Both
benches now run one clean iteration after timing, purely for the check. That is
also what made MKERNEL_BENCH_NO_SYNC=1 "fail" correctness at M=4096 -- the
kernel was fine, the harness was reading a racing loop's leftovers. It passes
now, at every shape, in both modes.
The correction to the numbers is large, and it goes the opposite way at the two
ends. Medians of 3, all runs guarded against other users' jobs:
gemm_rs was now TK of TK (previously)
4096x512x4096 0.397 0.077 0.078 1.02x 0.21x
8192x1024x8192 0.574 0.247 0.238 0.97x 0.42x
16384x2048x16384 1.264 0.969 0.985 1.02x 0.78x
32768x4096x32768 5.766 7.083 5.290 0.75x 0.92x
So gemm_rs is at parity with upstream at three of four shapes; the small-M gap
was almost entirely the harness launching a 0.08 ms kernel onto an idle GPU at
120 MHz, thirty times. At M=32768 the number gets worse, not better, because
sustained back-to-back running throttles the SM clock 2032 -> 1522 MHz (SW power
cap, 502 W) -- the per-iteration sync had been hiding that behind idle gaps.
ag_gemm moves less, since it was already batch-timed; it gains at small M from
the warmup the loop had been missing (M=4096 0.171 -> 0.113, 0.54x -> 0.81x).
Repeatability improves sharply: three runs now land within 0.5% of each other,
against 15% spread between min and median on the old path.
The trace hook captured its iteration after a sleep(0.1), so the traced launch started from an idle GPU at 120 MHz and the numbers described the clock ramp as much as the kernel. Eight warm iterations first. It mattered: M=8192's apparent gather bandwidth was 336 GB/s cold and is 626 GB/s warm, so half of what looked like a serious inefficiency at that shape was the clock. The other shapes move less (M=4096 779 -> 755, M=32768 unchanged at ~600) because their gather windows are long enough to average over the ramp.
AG_GEMM_RBPT selects how many row blocks a CTA computes per task. The default stays 2. Measured, medians of 3: M=4096 RBPT 1: 0.108 RBPT 2: 0.117 1 wins by 8% (0.84x -> 0.88x of TK) M=8192 0.332 0.313 M=16384 1.215 0.981 M=32768 8.670 6.576 1 loses by 32% Only M=4096 is task-starved enough to pay for the lost B-tile reuse -- at that shape just 32 of 148 CTAs have a tile at all, and doubling the task count is worth more than halving the operand bytes per FLOP. Everywhere else the reuse dominates. Per-shape selection would mean a runtime value where the tile types, semaphore counts and epilogue round count all read a compile-time constant, so the knob is left at 2 and the numbers are recorded next to it. Trying it surfaced a real latent bug. Both CTAs of a cluster point their TMA transactions at the leader's barrier, so the expects have to add up to CLUSTER_SIZE x (RBPT A tiles + B). With RBPT=2 the two MMA warps each expecting (A0, A1, B) happened to sum to exactly that; with RBPT=1 there is one MMA warp and it expected half the bytes, so every working compute CTA deadlocked on the first task. The coupling was invisible while RBPT and the MMA warp count were both 2. NUM_MMA_WARPS now names that relationship and the expect covers the cluster explicitly. The hang probe found it in one run: 64 working compute CTAs stuck at task 1 while gather had completed and the 68 CTAs with no work had exited cleanly. Default build unchanged and re-verified: 0.117 / 0.313 / 0.981 / 6.576, PASS.
7.6% at M=32768 (7.081 -> 6.552 ms, 0.86x -> 0.94x of TK) and 6.6% at M=24576. Bit-identical results; every shape PASS. Everything structural was already the same as upstream -- tile 128x256x64, cluster 512x256 at CLUSTER_SIZE 2, pipeline depth 4, 51.5 GB of operand traffic at 171 FLOP/byte. The steady-state trace put the MMA issue warp 53.0% blocked on inputs_arrived against 4.0% on the epilogue and 9.1% of the store warp inside the cross-GPU store_add, so it was starved of operands rather than held up by the reduce-scatter. Same bytes requested by both kernels means the difference is how many of them L2 serves, which is the task order. TK's get_task_idx bands the cluster rows (SUPER_M=8) and varies the row fastest inside a band. mKernel's decode varied the column fastest across the whole range, so the ~60 concurrent clusters pinned one A tile against 60 distinct B tiles. gemm_rs_decode_cluster_task does the banding in cluster units and keeps the device rotation that carries reduce-scatter ownership, expressed globally (device d starts total/NUM_DEVICES tasks in) rather than by slice. It is not a uniform win, so the order is chosen at runtime -- SUPER_M is pure index arithmetic, nothing about the tile types or semaphore counts depends on it. Cluster rows against the supertile's delta: M=4096 8 rows / 1 band +2% M=8192 16 / 2 ~0% M=16384 32 rows / 4 bands +6% M=24576 48 / 6 -6% M=32768 64 rows / 8 bands -8% Below ~6 bands the column-fastest order already wraps into the next row and covers a compact block by itself. The crossover sits between 4 and 6 bands and was not localised further, so the threshold is stated as the measurement it is. A first attempt keyed on columns-vs-clusters looked mechanistic but picked wrong at M=16384 (64 columns against 60 clusters), which is how the band count won. Also retracts an earlier claim in README_B300.md that sustained clocks cost gemm_rs 13% that TK does not pay: measured side by side, TK throttles more (6% of samples against 2%) and draws more peak power (666 W against 518 W).
Member
|
@monopodium Thanks! let me know if it is ready for review. |
…l site Cleanup only; no behaviour change intended. bench: the batch timing block had been copied into both benches with slight variation, which is how they drifted apart in the first place. It now lives in common.py as benchmark_batch() alongside the other shared helpers, with resolve_timing_mode() owning the env-var precedence (including the two legacy flags whose senses were opposite). ag_gemm_bench.py loses 18 lines, gemm_rs_bench.py 22, common.py gains 47 with the docstrings that explain why the timed loop leaves state dirty. The point is not the line count -- it is that the two benches can no longer drift. gemm_rs: the cluster loop computed round_pair/within/task_id on every task and discarded them whenever the super-tile order was selected. gemm_rs_decode_cluster now owns the choice, so the remap only runs on the path that needs it and there is one decode call site instead of two under an #ifdef. Verified: all seven Blackwell build configurations compile, and both timing modes run and pass correctness. gemm_rs on the Hopper path still fails, exactly as it does on origin -- upstream's tk_ops_group_mma_tcgen05_bf16.cuh does not compile for sm_90a. Timing was re-checked only under contention from another user's job, so the perf confirmation is owed.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.