From d621ba1e87d09f066d304f260f0f684c674d208a Mon Sep 17 00:00:00 2001 From: Michael Dementii Date: Wed, 9 Sep 2026 01:45:24 +0200 Subject: [PATCH] perf(ops): choose the predicated W8 GEMM cache policy instead of inheriting it w8_rowsplit_gemm_mma_kernel loads activations, codes and scales one way on full tiles and another on predicated ones: if constexpr (Full) { cp_async<16, Cache::cg>(...); // chosen } else { cp_async_zfill<16>(...); // Cache::ca, the memory.cuh default } The predicated side never chose ca; it inherited cp_async_zfill's default. Which side runs is decided per launch, not per tile - full needs (w.n % BM) == 0 && (x.ne[1] % BN) == 0 && ... - so any token count that is not a multiple of BN sends the whole GEMM down the predicated path. This makes the policy explicit on the schedule and sets it where it has been measured. One parameter, defaulting to Cache::ca, so adding it changes no instantiation. One schedule opts in: MmaR64C128, the wide route of the linear family. WHY ONLY THAT ONE Because the sign is not a property of the schedule. Sweeping an arm with every predicated instantiation on cg, against master, on the shapes the dispatch actually sends to each: schedule shape cg flip null p95 mma_r128_c64 34816x5120 -10.59 % 0.025 % mma_r64x32_c64_k128_a1 248320x5120 -3.69 % 0.188 % mma_r64_c96 2048x16384 -3.22 % 0.000 % mma_r48_c96 / c64 / c128 2048x16384 -2.09..-1.21 % <=0.03 % mma_r32_c64 / c96 / c112 2048x16384 -0.94..-0.56 % <=0.10 % mma_r32_c128 2048x16384 -2.92 % 0.149 % mma_r32_c128 9216x2048 0.00 % 0.000 % unresolved mma_r64x16_c48_k128_a1 34816x5120 -3.27 % 0.000 % mma_r64x16_c48_k128_a1 248320x5120 +1.07 % 0.205 % LOSES Most schedules gain, several by a lot. But MmaR64x16C48K128A1 - a BM = 64 tile, not one of the BM = 16 ones - gains 3.3 % on one shape it serves and loses 1.1 % on another, 6 of its 8 token counts positive there. A per-schedule constant cannot express that, so it cannot be set from one shape's reading. The two BM = 16 schedules are worse still: forced onto cg they cost +16.8 % on the feature route and +21.0 % on the dflash2 attn route. MmaR64C128 is the one schedule measured across enough shapes to show a consistent sign, so it is the one that moves. The knob exists so the next schedule can follow with its own measurement. EFFECT One RTX 5090, one session, SM clock recorded per cell (2663 to 2872 MHz, mean 2770), base b88c0f6f, ninfer_linear_bench --execution graph --warmup 5 --repeat 40 --flush-mib 256, three passes with the arm order rotated; ratios paired inside a pass, median over passes per T, then median over T; CI from a paired bootstrap over the per-T medians; the null is a second build of master carried through the same rotation. No number below comes from any other capture. shape T range grid this patch CI95 null p95 feature 5120x25600 50..56 small_t +0.045 % [+0.00, +1.00] 0.034 % feature 5120x25600 57..64 BM=16 +0.007 % [+0.00, +0.38] 0.014 % feature 5120x25600 65..128 BM=32 +0.000 % [+0.00, +0.00] 0.013 % feature 5120x25600 129..256 160 CTAs -11.27 % [-11.59, -11.22] 0.045 % feature 5120x25600 257..384 240 CTAs -2.02 % [-2.26, -1.57] 0.163 % feature 5120x25600 385..400 320 CTAs -0.86 % 0.000 % generic 9216x2048 4090..4130 144 x 33 -2.76 % [-2.77, -2.76] 0.250 % mid 34816x5120 129..320 544 x 2..3 -3.77 % [-3.94, -3.52] 0.084 % lm head 248320x5120 129..320 3880 x 2..3 -3.79 % [-4.07, -3.45] 0.752 % sparse 2048x4096 896..1100 32 x 7..9 -1.39 % [-1.60, -0.08] 0.158 % The first three rows are the routes the patch deliberately does not touch, and they read zero to within their nulls. That is the check that it changes only what it claims. The -11.27 % is a partial-wave effect and is not the headline: 160 CTAs is 0.94 of a wave on a 170-SM part, and the gain falls to -2.02 % at 1.41 waves and -0.86 % at 1.88. Nothing past T=400 was measured on that shape. The number to carry is the -2.8 to -3.8 % on the three large shapes, whose grids run many waves. Against the machine, from the same rows: 8 to 25 % of the 1674.5 GB/s sustained read, so nothing here is bandwidth-bound, and 48 to 91 % of the 209.5 TFLOP/s dense bf16 tensor peak. The generic shape sits at 90.6 % of that peak on master and still gains 2.76 %, which is where the change is least expected and most worth having. This is not a roofline wall moving; it is L1 tag and allocation work removed from a path that gets nothing back from L1. A control the sweep already contains: full needs (T % BN) == 0, so T=4096 is the one point in the generic sweep the patch cannot act on. It reads +0.0000 % while T=4090 reads -2.71 %, and the spread over the 40 predicated T is 0.410 pp, so this is not an alignment-phase artefact. WHY ca IS RIGHT WHERE IT STAYS Counters from the same card (ncu; hit rates and byte counts only - no timing from that capture is used here, because ncu serialises launches and controls caches between replays, which removes the L2 reuse the policy depends on): capture L1 sector hit rate L2 bytes 1/(1-hit) BM=16 feature route T=57..64 35.61 % -> 0.64 % x1.549 1.553 BM=32 feature route T=65..128 0.30 % -> 1.07 % x0.992 1.003 MmaR64C128 n=5120 T=129..256 0.29 % -> 0.37 % x1.000 1.003 MmaR64C128 n=248320 T=129..256 1.59 % -> 1.53 % x1.017 1.016 Only the BM=16 schedule gets reuse out of L1, and bypassing it turns exactly those hits into L2 traffic: x1.549 measured against 1.553 predicted from the hit rate alone. At BM=16 the schedule is 128 threads and 22.8 KB of shared memory, so four CTAs sit on an SM and m/BM = 320 of them read the same 16 KB tile. MmaR64C128 is 46 KB and two CTAs, and at n=5120 its grid is 160 CTAs on 170 SMs, so no two share an SM. This does not separate "L1 has reuse" from "the tile is shared" - one captured kernel has a non-trivial hit rate and it is also the only BM=16 one - and the r64x16_c48 result above shows the rule is not simply about BM either. VERIFICATION A cache policy selects where a load is cached, not what it returns, so the output should be bit-identical. Witnessed rather than argued: the six op tests covering the eight translation units that include this header were run on master and on this commit with NINFER_OP_REPORT_STATS=1, which prints max_abs, max_rel and rel_l2 per case at %.17g. That is 4145 numeric records, byte-identical on both sides: ninfer_linear_w8_a16_test 600 records ninfer_linear_add_w8_a16_test 123 ninfer_linear_swiglu_w8_a16_test 192 ninfer_linear_pair_w8_a16_test 288 ninfer_attn_input_proj_test 2816 ninfer_gdn_input_proj_test 126 ctest is 114/114 on both sides. clang-format adds no replacement to either file. The two sides are separate builds from sources whose hashes differ, and the gate compares the resulting binaries and fails if they turn out equal. That check is there because the previous acceptance run for this branch did not have it: it rebuilt one tree twice and compared a build against itself. "Changes no instantiation but the one that opts in" is checked rather than asserted. Comparing the SASS of ninfer_ops between the two builds: 30 instantiations of the kernel and 454 LDGSTS on both sides, and exactly one instantiation differs - the predicated BM=64 BN=128 one, whose 15 LDGSTS all gain the BYPASS modifier (0 -> 15 of 15). The other 29 are unchanged. WHAT IS NOT HERE No end-to-end number. The one offered in the first version did not survive its own control: ragged pp8191 improved, but so did the unchanged aligned pp8192, leaving +0.102 % and -0.371 %. This claim is operator-level. Five of the nineteen dispatch cells that reach MmaR64C128 are measured, covering five of the fifteen distinct (n, k) shapes; the other fourteen cells are not. Every other predicated instantiation keeps the inherited ca, including the ones measured above to gain from cg. Moving them is a separate change with its own per-shape measurement, because r64x16_c48 shows a schedule can want opposite things on two shapes it serves. Co-Authored-By: Claude Opus 5 --- src/ops/linear/w8/w8_rowsplit_gemm_mma.cu | 15 +++++++- src/ops/linear/w8/w8_rowsplit_gemm_mma.cuh | 41 +++++++++++++++++++--- 2 files changed, 50 insertions(+), 6 deletions(-) diff --git a/src/ops/linear/w8/w8_rowsplit_gemm_mma.cu b/src/ops/linear/w8/w8_rowsplit_gemm_mma.cu index 0d41ca96a3..77f911cb40 100644 --- a/src/ops/linear/w8/w8_rowsplit_gemm_mma.cu +++ b/src/ops/linear/w8/w8_rowsplit_gemm_mma.cu @@ -30,6 +30,16 @@ template void launch_route(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream) { const bool full = (w.n % Schedule::BM) == 0 && (x.ne[1] % Schedule::BN) == 0 && w.k == w.padded_shape[1] && (w.k % Schedule::BK) == 0; + // cp.async.cg requires a 16 B-aligned source. On the predicated path the activation source + // steps by 8 bf16 and the scale source by kg = padded_k / 32 two-byte scales, so the encoding + // is legal only when both are multiples of 8. The full path inherits this from (w.k % BK) == 0 + // above; a schedule that asks for cg on the predicated path has to be checked here. + if constexpr (Schedule::kPredicatedCache == Cache::cg) { + if (!full && ((w.k % 8) != 0 || ((w.padded_shape[1] / 32) % 8) != 0)) { + throw std::invalid_argument( + "w8 rowsplit mma: predicated cg needs k and the group count to be multiples of 8"); + } + } for_each_token_slice(x.ne[1], Schedule::BN, [&](std::int32_t offset, std::int32_t count) { const Tensor x_slice = x.slice(1, offset, count); Tensor out_slice = out.slice(1, offset, count); @@ -77,7 +87,10 @@ using MmaR48C112 = W8RowSplitMmaGemmSchedule<48, 112, 48, 16, 2>; using MmaR48C128 = W8RowSplitMmaGemmSchedule<48, 128, 48, 16, 2>; using MmaR64C96 = W8RowSplitMmaGemmSchedule<64, 96, 64, 16, 2>; using MmaR64C112 = W8RowSplitMmaGemmSchedule<64, 112, 64, 16, 2>; -using MmaR64C128 = W8RowSplitMmaGemmSchedule<64, 128, 64, 16, 2, 2>; +// The wide route of the linear family is the one schedule measured to want cg on the predicated +// path; every other schedule keeps the inherited ca. +using MmaR64C128 = + W8RowSplitMmaGemmSchedule<64, 128, 64, 16, 2, 2>::with_predicated_cache; using MmaR96C96 = W8RowSplitMmaGemmSchedule<96, 96, 48, 16, 2>; using MmaR128C64 = W8RowSplitMmaGemmSchedule<128, 64, 64, 16, 2>; using MmaR128C80 = W8RowSplitMmaGemmSchedule<128, 80, 64, 16, 2>; diff --git a/src/ops/linear/w8/w8_rowsplit_gemm_mma.cuh b/src/ops/linear/w8/w8_rowsplit_gemm_mma.cuh index 13e555e1b1..154c275f6f 100644 --- a/src/ops/linear/w8/w8_rowsplit_gemm_mma.cuh +++ b/src/ops/linear/w8/w8_rowsplit_gemm_mma.cuh @@ -9,6 +9,7 @@ // m16n8k16 BF16 MMA with FP32 accumulation. #include "ops/common/mma.cuh" +#include "ops/common/memory.cuh" #include "ops/common/math.cuh" #include "ops/linear/w8/w8_rowsplit_output.cuh" @@ -26,8 +27,21 @@ union alignas(16) W8Bf16x8Bits { static_assert(sizeof(W8Bf16x8Bits) == 16); +// The predicated loads below inherited Cache::ca from cp_async_zfill's default, while the full path +// a few lines down spells cg. This parameter makes that a choice. It defaults to ca, so adding it +// changes no instantiation, and it governs the predicated branch only - the full branch keeps its +// own cg - which is why it is named for it. +// +// ca stays the default because whether cg pays is not a property of the schedule. On the two +// BM = 16 schedules in this tree - src/ops/linear/w8/w8_feature.cu and +// src/ops/attn_input_proj/w8/w8_dflash2_attn_input.cu - m / BM CTAs read the same activation tile +// (m is the weight-row count, the kernel's first size argument) and four of them share an SM; L1 +// serves 35.6 % of their sectors, and forcing cg there costs +16.8 % and +21.0 % of the operation. +// Most other schedules gain from cg, but not all, and not on every shape: MmaR64x16C48K128A1, a +// BM = 64 tile, gains 3.3 % on [34816, 5120] and loses 1.1 % on [248320, 5120]. So the policy is +// set per schedule and only where it has been measured across the shapes that reach it. template + int ACTIVATION_STAGES_ = STAGES_, Cache PredicatedCache_ = Cache::ca> struct W8RowSplitMmaGemmSchedule { static constexpr int BM = BM_; static constexpr int BN = BN_; @@ -48,6 +62,16 @@ struct W8RowSplitMmaGemmSchedule { static constexpr int SMEM_BYTES = BM * BK * 2 + ACTIVATION_STAGES * BN * BK * 2 + BM * BK + BM * SCALE_CACHE_BYTES; + // Cache policy for the predicated loads only; see the note above the template. + static constexpr Cache kPredicatedCache = PredicatedCache_; + + // Restate this schedule with a different predicated policy, leaving every other parameter where + // it is rather than respelling it - and its default with it - at the point of use. + template + using with_predicated_cache = + W8RowSplitMmaGemmSchedule; + static_assert(BM % WM == 0 && BN % WN == 0); static_assert(WM % 16 == 0 && WN % 8 == 0); static_assert(THREADS <= 1024); @@ -132,7 +156,11 @@ __global__ __launch_bounds__(Cfg::THREADS, Cfg::MIN_BLOCKS) void w8_rowsplit_gem cp_async<16, Cache::cg>(dst, &x[static_cast(nn) * k + kk]); } else { const int valid = (nn < n && kk < k) ? min(8, k - kk) * 2 : 0; - ninfer::ops::cp_async_zfill<16>( + // cp.async.cg needs the source naturally aligned to 16 B. kk steps by 8 and x is + // bf16, so that holds exactly when k is a multiple of 8. The full path gets this + // from its own (k % BK) == 0 test; the predicated path has no such guarantee, so + // launch_route checks it for any schedule that asks for cg. + ninfer::ops::cp_async_zfill<16, Cfg::kPredicatedCache>( dst, &x[static_cast(nn < n ? nn : 0) * k + (kk < k ? kk : 0)], valid); } @@ -156,8 +184,8 @@ __global__ __launch_bounds__(Cfg::THREADS, Cfg::MIN_BLOCKS) void w8_rowsplit_gem } else { const bool valid_row = output_tile.valid(grow, m); const std::int64_t gi = static_cast(valid_row ? grow : 0) * kg + g0; - ninfer::ops::cp_async_zfill<16>(dst, &codes[gi * 32 + chunk * 16], - valid_row ? 16 : 0); + ninfer::ops::cp_async_zfill<16, Cfg::kPredicatedCache>( + dst, &codes[gi * 32 + chunk * 16], valid_row ? 16 : 0); } } if ((kt % SCALE_CACHE_TILES) == 0) { @@ -173,7 +201,10 @@ __global__ __launch_bounds__(Cfg::THREADS, Cfg::MIN_BLOCKS) void w8_rowsplit_gem const int valid_scales = valid_row && g0 < kg ? min(8, kg - g0) : 0; const std::int64_t gi = static_cast(valid_row ? grow : 0) * kg + min(g0, kg - 1); - ninfer::ops::cp_async_zfill<16>(dst, &scales[gi * 2], valid_scales * 2); + // gi steps by kg, and scales are 2 B, so 16 B alignment here needs kg to be + // a multiple of 8. launch_route checks that alongside k. + ninfer::ops::cp_async_zfill<16, Cfg::kPredicatedCache>(dst, &scales[gi * 2], + valid_scales * 2); } } }