From f0c5ed88b9e2d2f148f94bfbb3d3472c5ca9bd3e Mon Sep 17 00:00:00 2001 From: MichaelDementii Date: Thu, 3 Sep 2026 00:24:40 +0000 Subject: [PATCH] perf(fp8): stage the A8 GEMM operands through TMA Add a second implementation of the FP8 row-scaled A8 GEMM and route to it where a cost model finds it cheaper. The arithmetic is unchanged: 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. What differs is the CTA's shape and how it is fed. The output tile grows from 64 to 256 tokens and the consumer warp tile from 32x32 to 64x64, so the MMA issued per ldmatrix goes from 8 to 32. Feeding that tile from one CTA needs a deeper, narrower pipeline - four stages of K=64 in place of two of K=128, the buffered K extent unchanged at 256 - and a dedicated producer warp issuing cp.async.bulk.tensor under an mbarrier pair, in place of the CTA-wide barrier the old K loop takes on every stage. Occupancy falls from two CTAs of 256 threads to one of 288, which is why narrow shapes are declined. The 64-byte TMA swizzle is a function of the shared-memory address, not of an offset relative to the destination, so a tile base must sit on a 512-byte swizzle atom. Measured: the same tile loaded 128 bytes past an atom boundary de-swizzles wrongly in all 1024 of its bytes. That alignment is declared on the storage, the union, the outer storage and the allocation, with static asserts that both stage strides and a paired block's second branch are whole atoms. CTAs are rasterised token-fastest, so the blocks that share a weight tile run at the same time. Each CTA computes the same tile either way, so the output does not change; the operator gains 1.5 to 3.3% on 34816x5120, whose weights do not fit in L2, and loses 0.2 to 0.5% on 5120x17408 through linear_add. Routing is a comparison of two quantised costs - waves needed times the work one SM carries through a wave - with one empirical constant for this part, plus a width floor and one per-geometry ceiling the model cannot express. The sweeps behind all three, the alignment measurement, the rasterisation choice, and what each bound costs are in docs/maintainer/fp8-a8-tma-route.md. Measured on an RTX 5090 (sm_120a, 170 SMs, 525 W) against the route it replaces, the arms alternating inside each pass and the base repeated under a second label to carry the instrument's zero: operator, linear: 0.877 to 0.989 where the route is taken operator, linear_add: 0.895 to 0.979 where the route is taken end to end, qwen3.8-27b prefill: +2.40% at chunk 1024, +4.95% at 4096, +4.38% at 8192 decode: 0.9993 to 1.0002 over four repetitions, and the route launches zero times in a decode-only kernel census taken with CUDA graphs expanded Output is not bitwise identical to the route it replaces. Where the width is a whole number of the old kernel's token tiles the two match byte for byte; at other widths they round differently by construction, and the numerical tests check those against a host reference. Co-Authored-By: Claude Opus 5 --- docs/README.md | 1 + docs/maintainer/fp8-a8-tma-route.md | 252 +++++++++ src/CMakeLists.txt | 2 +- .../attn_input_proj/fp8/fp8_attn_input_a8.cu | 23 + .../gdn_input_proj/fp8/fp8_gdn_input_a8.cu | 12 + src/ops/linear/fp8/fp8_a8.cu | 24 +- src/ops/linear/fp8/fp8_a8_schedule.cuh | 284 +++++++++++ src/ops/linear/fp8/fp8_a8_tma.cuh | 478 ++++++++++++++++++ src/ops/linear_add/fp8/fp8_linear_add_a8.cu | 14 +- .../linear_swiglu/fp8/fp8_linear_swiglu_a8.cu | 13 + tests/ops/linear/test_fp8_a8.cpp | 26 + tests/ops/linear_add/test_fp8.cpp | 9 +- tests/ops/linear_swiglu/test_fp8.cpp | 8 +- tests/ops/test_attn_input_proj.cpp | 5 +- tests/ops/test_gdn_input_proj.cpp | 6 +- 15 files changed, 1150 insertions(+), 7 deletions(-) create mode 100644 docs/maintainer/fp8-a8-tma-route.md create mode 100644 src/ops/linear/fp8/fp8_a8_tma.cuh diff --git a/docs/README.md b/docs/README.md index 66942624f8..949a7e6b8b 100644 --- a/docs/README.md +++ b/docs/README.md @@ -53,6 +53,7 @@ Runtime and Op references: - [Op admission, contracts, ownership, qualification, and performance rules](maintainer/op-development.md) - [ReplaySSM GDN technical reference](maintainer/replayssm-gdn.md) - [Linear benchmark contract and registered suites](maintainer/linear-benchmark.md) +- [FP8 A8 GEMM TMA route: measurements behind its routing constants](maintainer/fp8-a8-tma-route.md) `engine-architecture.md` is the sole top-level Engine architecture reference. `resource-scheduling-and-context-cache.md` is its narrower authority for resource selection, diff --git a/docs/maintainer/fp8-a8-tma-route.md b/docs/maintainer/fp8-a8-tma-route.md new file mode 100644 index 0000000000..b7aadd5ca7 --- /dev/null +++ b/docs/maintainer/fp8-a8-tma-route.md @@ -0,0 +1,252 @@ +# FP8 A8 GEMM: the TMA-staged route + +Performance reference for the second implementation of the FP8 row-scaled A8 GEMM. The routing +header states invariants; the measurements that chose the constants live here, where they can be +re-taken and amended without touching code. + +Nothing here is a semantic authority. The Op contract, its represented inputs and its numerical +qualification are unchanged by this route and remain where they were. + +## 1. What the route is + +`fp8_a8_tma_kernel` is a second implementation of the GEMM `fp8_a8_mma.cuh` computes: E4M3 codes +with a per-row BF16 weight scale and a per-token FP32 activation scale, accumulated in FP32 through +`m16n8k32`, with the Op's epilogue and output policy applied unchanged. It is reached from five +call sites, all prefill-width: + +| Op | geometry `` | epilogue / output | +|---|---|---| +| `linear` | all five | identity, contiguous | +| `linear_add` | `<5120,6144>`, `<5120,17408>` | residual add folded into the store | +| `linear_swiglu` | `<34816,5120>` | paired rows, SwiGLU store | +| `attn_input_proj` | `<14336,5120>` | four-output split store | +| `gdn_input_proj` | `<16384,5120>` | identity, contiguous | + +Of the published artifacts only **qwen3.8-27b** carries `FP8_E4M3FN_ROW_BF16S` weights — attention +and GDN projections, MLP gate/up and down on layers 56..63, the output head and the embedding. +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 it there returns 1.000 by construction. + +## 2. What actually changed + +The route is 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 (`cuobjdump -res-usage`) | 94 | 166 | +| dynamic shared per CTA | 49152 B + 1024 B static | 98816 B + 0 B static | +| 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. That is the mechanism behind +the width floor in section 5, 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, before allocation granularity. + +## 3. Shared-memory alignment: a precondition, not a preference + +`fp8_tma_shared_byte()` derives the sixteen-byte segment of a swizzled row from the row index taken +**relative to the tile base**. The 64-byte swizzle XORs the segment index with `(row / 2) % 4` and +so closes after eight rows: its repeating unit is 512 bytes. + +The hardware's swizzle is a function of the shared-memory address, not of an offset relative to the +destination. Measured directly — the same tile, the same descriptor, two destinations: + +```text +dynamic shared base : 0x600 (mod 512 = 0) +tile at atom-aligned base : 0 / 1024 bytes wrong +tile at base + 128 : 1024 / 1024 bytes wrong +``` + +So the tile base must be 512-byte aligned. That is stronger than the 128 bytes the TMA store itself +requires, and the gap is reachable: `__align__` on the `extern __shared__` array is what places the +base, and with 16 bytes declared it follows whatever static shared the kernel carries. + +| declared alignment | dynamic base, by preceding static shared (16 / 64 / 128 / 1024 / 1152 B) | +|---|---| +| `__align__(16)` | 0x410, 0x440, 0x480, 0x800, 0x880 — 16, 64, 128, 0, **128** mod 512 | +| `__align__(128)` | 0x400 in all five — 0 mod 512 | +| `__align__(512)` | 0x400 in all five — 0 mod 512 | + +`__align__(128)` happens to produce an atom-aligned base on CUDA 13.1, which is why the route +computed correct results before this was stated. That is a property of one toolchain's layout, not +of the source. `kFp8A8TmaSwizzleAtomBytes = 512` now sits on the tensor storage, the union around +it, the outer storage and the allocation, with static asserts that the stage stride (16384 B), the +weight stage stride (8192 B) and a paired block's second branch (4096 B) are all whole atoms. The +declaration costs 384 bytes of padding — `sizeof` goes from 98432 to 98816 against a 101376 cap — +and no time: the base is 0x400 either way. + +## 4. The routing model, and what it does not represent + +`fp8_a8_tma_cheaper` compares two quantised costs: waves needed, times the work one SM carries +through a wave. Both kernels leave part of a wave idle and they quantise differently — one CTA of +256 tokens against two CTAs of 64 — so the multiprocessor count enters on both sides rather than +being frozen into the constant. + +`kFp8A8TmaWorkRatio` is the TMA route's time per token of work relative to the route it replaces. +It is an empirical constant for one part, solved from the widest measured point where wave +quantisation is mildest. It is **not** portable: on another part the two kernels' intrinsic speeds +differ and it must be re-solved. What travels is the shape of the comparison. + +The model has no term for pipeline fill, none for K, and none for the cost of an epilogue. Both +width bounds below exist because of shapes it therefore cannot tell apart. `kFp8A8TmaMargin` is a +margin on modelled cost, not on measured time. + +## 5. Width bounds + +Both bounds are measurements, and both were re-taken on `b88c0f6f`. Every table below is the ratio +of the route's time to the time of the route it replaces, on one RTX 5090 in one session, arms +alternating in mirrored order. Clock locking is not available on this machine and the card idles at +180 MHz, so alternation and repetition are the whole defence against drift. + +Widths the shipped predicate declines were taken with a scratch build whose floor is lowered to the +tile and whose ceiling is removed. Where the cost model then declines a width on its own, both arms +run the same kernel and the cell reads 1.000 — that is the model's verdict, not the bound's. + +### 5.1 Floor, `kFp8A8TmaMinTokens = 1024` + +Below the floor the ratio is a sawtooth and it swings both ways. The route's time steps every 256 +tokens and the route it replaces steps every 64, so within a band the ratio falls as the older +kernel climbs, then jumps when the route needs another tile. With the floor removed and the cost +model deciding alone: + +| T | 256 | 320 | 384 | 448 | 512 | 576 | 640 | 704 | 768 | 832 | 896 | 960 | +|---|---|---|---|---|---|---|---|---|---|---|---|---| +| `14336x5120` | **1.108** | 1.000 | 1.000 | 1.000 | 1.000 | 1.000 | 1.017 | **0.910** | **0.903** | 0.993 | 1.000 | 1.000 | +| `16384x5120` | **0.900** | 1.014 | 1.000 | 1.000 | 1.026 | 1.009 | 1.000 | 1.000 | 1.000 | 1.000 | 1.000 | 1.000 | +| `5120x6144`, `linear_add` | 1.000 | 1.000 | 1.000 | 1.000 | 0.979 | **1.182** | **1.164** | **1.162** | **1.159** | 0.920 | 0.921 | 0.921 | + +Seventeen cells where the model takes the width: **eight gains, 0.7 to 10.0 percent, against nine +losses, 0.9 to 18.2 percent.** The worst loss is nearly twice the best gain, and both extremes are +in the same 256-token band on different shapes, which is what a constant fitted above 1024 looks +like when it is asked below it. + +That is what the floor buys. It is a bound on the constant, not on the model: `kFp8A8TmaWorkRatio` +was solved at and above the widths the product runs and has no validation below them. Whether it +should move is a question for a wider sweep than this one; on this sample, moving it down costs more +than it returns. + +### 5.2 Ceiling, `kFp8A8TmaMaxTokens = 4096` + +The model counts blocks along output rows and token tiles and has no term for K, so it returns one +verdict for the two 5120-row residual shapes. Measured through `linear_add`, the call site both are +reached from in production, they diverge above 4096: + +| T | 4096 | 6144 | 8192 | 10240 | 12288 | 14336 | +|---|---|---|---|---|---|---| +| `5120x6144` | **0.979** | 1.011 | 1.006 | 1.000 | 1.024 | 1.021 | +| `5120x17408` | **0.937** | 0.952 | 0.932 | 1.000 | **0.846** | **0.850** | + +`5120x6144` loses 0.6 to 2.4 percent everywhere above the bound while `5120x17408` gains 4.8 to 15.4 percent; +10240 reads 1.000 on both because the model declines it unaided. `5120x6144` is the smallest GEMM of +the five — 62.9 MFLOP per token against 146.8, 167.8, 178.3 and 356.5 — so it has the least work to +amortise the pipeline over. That is the direction the numbers point in, not a mechanism this +measurement establishes. + +**The bound is not free, and the two benches disagree about it.** Through plain `linear` the same +geometry gains at the same widths — 0.928, 0.961, 0.955, 1.001, 0.964, 0.972 — so the ceiling costs +that path 2.8 to 7.2 percent. It is kept because `linear_add` is the production call site for +`5120x6144` and `linear` at that shape is a development surface. Anyone who makes plain `linear` a +production path at this geometry should re-take this table before trusting the bound. + +### 5.3 Rasterisation + +The CTAs are rasterised token-fastest, so the blocks that share a weight tile run at the same time +and the weight matrix is read once instead of once per token tile. Each CTA computes the same tile +either way, so this cannot change the output; it changes only the order the work distributor hands +blocks to SMs. Ratio of token-fastest to the stock grid, same branch otherwise: + +| shape | weights | T=2048 | T=4096 | T=8192 | +|---|---|---|---|---| +| `34816x5120` | 178 MB | **0.983** | **0.985** | **0.967** | +| `5120x17408` | 89 MB | 1.000 | 0.999 | 1.000 | +| `16384x5120` | 84 MB | 1.000 | 0.996 | 0.997 | +| `14336x5120` | 73 MB | 1.000 | 0.998 | 1.000 | +| `5120x6144` | 31 MB | 1.007 | 1.000 | 0.998 | +| `5120x17408`, `linear_add` | 89 MB | 1.002 | 1.005 | 1.004 | + +The one geometry that gains is the one whose weights do not fit in this part's 96 MB of L2. Where +they fit, L2 already supplies the reuse and the order is worth nothing; through `linear_add` at +`5120x17408` it is worth −0.2 to −0.5 percent, because there the activation is the larger stream — +142 MB at T=8192 against 89 MB of weights — and token-fastest shares the smaller one. + +End to end the gain wins: −0.12 percent at prefill chunk 1024, +0.71 at 4096, +0.97 at 8192. + +A per-geometry choice would recover the loss. The quantity to choose on is computable rather than +fitted: the bytes the concurrently resident CTAs touch under each order, which follows from the tile +sizes, K, the grid and the multiprocessor count. On this sample it predicts the sign of every cell +above. It is not taken here because it is a decision for three routes — this one, the cp.async +schedules, which are all registered `TokenFast`, and the NVFP4 TMA route — and it should be settled +for all three at once, on a wider sweep than the six geometries here. + +## 6. The multiple-of-tile condition, and why it is gone + +An earlier form of the predicate also required `tokens % MmaSchedule::kBlockTokens == 0`. It bought +bit-identity with the route being replaced: at those widths the old kernel takes its `FullTokens` +branch, whose expression this kernel matches exactly, so correctness could be settled by `memcmp` +rather than by a tolerance argument. + +It was removed because it paid for that with speed and with coverage, and rescued nothing. + +Swept over 1024..1520 in steps of 16 on `14336x5120`, with three arms on one card — the route it +replaces, 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 the condition declined | widths | model would have | route measured | so the condition | +|---|---|---|---|---| +| 1040..1328, not whole tiles | 15 | declined anyway | 1.199 down to 1.039 | repeated a decision | +| 1360..1520, not whole tiles | 9 | **taken** | **0.975 down to 0.915** | 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 could not. The model's only +width-dependent inputs are `ceil(T / 256)` and `ceil(T / 64)`, both constant between adjacent +multiples of 64, so its verdict is fixed across each band and the condition can only subtract from +a decision already made. + +Confirmed against the product rather than inferred: with the condition removed and nothing else +changed, the widths where the model declines measure 0.999 to 1.001 — removing it changes nothing +there — and 1360, 1440 and 1520 measure 0.975, 0.915 and 0.916. + +Removing it admits between 1134 and 5292 further widths per geometry over 1024..8192. Those widths +are not byte-comparable with the previous kernel by construction, so the numerical tests cover them +against a host reference instead: `1345` on `14336x5120`, `1153` on `16384x5120` and `34816x5120`, +`4001` on both residual shapes. + +## 7. Reproducing + +One RTX 5090, driver 580.159.03, CUDA 13.1, 525 W limit, 170 SMs. Absolute times from this machine +do not travel; ratios inside one run do. + +```bash +# operator, one shape +./build/bench/ninfer_linear_bench --qtype FP8 --policy a8 --n 14336 --k 5120 \ + --sweep 1024:8192:1024 --warmup 3 --repeat 12 --csv-out out.csv + +# operator, the residual call site +./build/bench/ninfer_fp8_linear_add_bench --k 6144 --policy a8 \ + --t-sweep 1024,2048,4096,8192 --warmup 3 --repeat 12 --csv-out out.csv + +# end to end, on the only artifact this route runs on +./build/bench/ninfer_bench --weights qwen3_8_27b_nvfp4.ninfer -p 16384 -r 3 --warmup 1 \ + --prefill-chunk 4096 --max-ctx 32768 -o csv --output-file out.csv +``` + +The swizzle probe of section 3 is out of tree: it issues `cp.async.bulk.tensor` directly into two +chosen shared offsets, which has no in-tree home. diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 3e99d8c7e4..a7f00182c4 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -307,7 +307,7 @@ add_library(ninfer_ops STATIC ops/wrapper/vision_pos_embed.cpp) ninfer_internal_includes(ninfer_ops) ninfer_cuda_archive(ninfer_ops) -target_link_libraries(ninfer_ops PUBLIC ninfer_core PRIVATE ninfer_nvfp4_non_rdc) +target_link_libraries(ninfer_ops PUBLIC ninfer_core PRIVATE ninfer_nvfp4_non_rdc CUDA::cuda_driver) # Product-independent text and media mechanisms. add_library(ninfer_text STATIC diff --git a/src/ops/attn_input_proj/fp8/fp8_attn_input_a8.cu b/src/ops/attn_input_proj/fp8/fp8_attn_input_a8.cu index 3062c21220..611301be27 100644 --- a/src/ops/attn_input_proj/fp8/fp8_attn_input_a8.cu +++ b/src/ops/attn_input_proj/fp8/fp8_attn_input_a8.cu @@ -3,6 +3,8 @@ #include "core/device.h" #include "ops/attn_input_proj/fp8/fp8_attn_input_output.cuh" #include "ops/linear/fp8/fp8_a8_mma.cuh" +#include "ops/linear/fp8/fp8_a8_schedule.cuh" +#include "ops/linear/fp8/fp8_a8_tma.cuh" #include "ops/linear/fp8/fp8_config.h" #include "ops/linear/fp8/fp8_output.cuh" @@ -39,6 +41,24 @@ void launch_mma(const Weight& weight, Tensor& q, Tensor& gate, Tensor& k, Tensor CUDA_CHECK(cudaGetLastError()); } +// The TMA-staged route. Placed in the wide arm of the ladder below, next to the schedule it +// falls back to: the predicate compares the two, so the schedule it names has to be the one that +// would otherwise run at this width. +void run_tma(const Weight& weight, Tensor& q, Tensor& gate, Tensor& k, Tensor& v, + Fp8A8Workspace workspace, int tokens, cudaStream_t stream) { + const Fp8AttentionInputOutput output{ + static_cast<__nv_bfloat16*>(q.data), + static_cast<__nv_bfloat16*>(k.data), + static_cast<__nv_bfloat16*>(gate.data), + static_cast<__nv_bfloat16*>(v.data), + }; + fp8_a8_tma_launch::Type>( + workspace.codes, workspace.scales, static_cast(weight.qdata), + static_cast(weight.scales), tokens, Fp8IdentityEpilogue{}, output, + stream); + CUDA_CHECK(cudaGetLastError()); +} + template void run(const Weight& weight, Tensor& q, Tensor& gate, Tensor& k, Tensor& v, Fp8A8Workspace workspace, int tokens, cudaStream_t stream) { @@ -76,6 +96,9 @@ void fp8_attn_input_a8_launch(const Tensor& x, const Weight& weight, Tensor& q, run(weight, q, gate, k, v, workspace, x.ne[1], stream); else if (x.ne[1] <= 144) run(weight, q, gate, k, v, workspace, x.ne[1], stream); + else if (fp8_a8_tma_applies::Type, Prefill>( + x.ne[1], workspace.codes, weight.qdata)) + run_tma(weight, q, gate, k, v, workspace, x.ne[1], stream); else run(weight, q, gate, k, v, workspace, x.ne[1], stream); } diff --git a/src/ops/gdn_input_proj/fp8/fp8_gdn_input_a8.cu b/src/ops/gdn_input_proj/fp8/fp8_gdn_input_a8.cu index 0f5cf8a98b..67be291f33 100644 --- a/src/ops/gdn_input_proj/fp8/fp8_gdn_input_a8.cu +++ b/src/ops/gdn_input_proj/fp8/fp8_gdn_input_a8.cu @@ -47,6 +47,18 @@ void launch_mma(const Weight& weight, Tensor& qkv, Tensor& z, Fp8A8Workspace wor void fp8_gdn_input_a8_launch(const Tensor& x, const Weight& weight, Tensor& qkv, Tensor& z, Fp8A8Workspace workspace, cudaStream_t stream) { launch_fp8_a8_quantize(x, weight, workspace, stream); + using TmaSchedule = typename Fp8LinearA8TmaSchedule::Type; + if (fp8_a8_tma_applies(x.ne[1], workspace.codes, + weight.qdata)) { + const Fp8GdnInputOutput output{static_cast<__nv_bfloat16*>(qkv.data), + static_cast<__nv_bfloat16*>(z.data)}; + fp8_a8_tma_launch(workspace.codes, workspace.scales, + static_cast(weight.qdata), + static_cast(weight.scales), + x.ne[1], Fp8IdentityEpilogue{}, output, stream); + CUDA_CHECK(cudaGetLastError()); + return; + } if ((x.ne[1] % Schedule::kBlockTokens) == 0) { launch_mma(weight, qkv, z, workspace, x.ne[1], stream); } else { diff --git a/src/ops/linear/fp8/fp8_a8.cu b/src/ops/linear/fp8/fp8_a8.cu index a937254fbb..34bd31d2c7 100644 --- a/src/ops/linear/fp8/fp8_a8.cu +++ b/src/ops/linear/fp8/fp8_a8.cu @@ -4,6 +4,7 @@ #include "ops/common/math.cuh" #include "ops/common/warp.cuh" #include "ops/linear/fp8/fp8_a8_schedule.cuh" +#include "ops/linear/fp8/fp8_a8_tma.cuh" #include "ops/linear/fp8/fp8_config.h" #include "ops/linear/fp8/fp8_output.cuh" @@ -100,10 +101,31 @@ void launch_quantize_exact(const Tensor& x, Fp8A8Workspace workspace, cudaStream CUDA_CHECK(cudaGetLastError()); } +// The TMA route is chosen by fp8_a8_tma_applies, not by this function: it needs at least one full +// 256-token tile, a width inside this geometry's measured bounds, and a modelled cost below the +// cp.async route's. A partial trailing tile is allowed: TMA zero-fills the rows past the descriptor +// extent, and the epilogue neither reads their activation scale nor stores them. Everything else +// keeps the cp.async route. +template +void launch_tma(const Weight& weight, Tensor& out, Fp8A8Workspace workspace, std::int32_t tokens, + cudaStream_t stream) { + fp8_a8_tma_launch( + workspace.codes, workspace.scales, static_cast(weight.qdata), + static_cast(weight.scales), tokens, Fp8IdentityEpilogue{}, + Fp8ContiguousOutput{static_cast<__nv_bfloat16*>(out.data), Geometry::kOutputRows}, stream); + CUDA_CHECK(cudaGetLastError()); +} + template void launch_problem(const Weight& weight, Tensor& out, Fp8A8Workspace workspace, std::int32_t tokens, cudaStream_t stream) { - using Schedule = typename Fp8LinearA8ProductionSchedule::Type; + using Schedule = typename Fp8LinearA8ProductionSchedule::Type; + using TmaSchedule = typename Fp8LinearA8TmaSchedule::Type; + if (fp8_a8_tma_applies(tokens, workspace.codes, + weight.qdata)) { + launch_tma(weight, out, workspace, tokens, stream); + return; + } if ((tokens % Schedule::kBlockTokens) == 0) { launch_mma(weight, out, workspace, tokens, stream); } else { diff --git a/src/ops/linear/fp8/fp8_a8_schedule.cuh b/src/ops/linear/fp8/fp8_a8_schedule.cuh index 613452197f..e2dbd99b40 100644 --- a/src/ops/linear/fp8/fp8_a8_schedule.cuh +++ b/src/ops/linear/fp8/fp8_a8_schedule.cuh @@ -1,8 +1,15 @@ #pragma once #include "ops/linear/fp8/fp8_a8_mma.cuh" +#include "ops/linear/fp8/fp8_a8_tma.cuh" #include "ops/linear/fp8/fp8_config.h" +#include +#include +#include +#include +#include + namespace ninfer::ops::detail { template @@ -38,4 +45,281 @@ struct Fp8LinearA8ProductionSchedule { Fp8MmaFragmentPipeline::PingPong, Fp8MmaRaster::TokenFast>; }; +// TMA route. One CTA per SM with a 256-token tile. Four stages put the tensor pipeline at +// 4 * (256 + 128) * 64 = 96 KiB, which is the largest depth that fits under the 99 KiB per-CTA +// cap; five would need 120 KiB and fails to build. The epilogue shares this storage through a +// union, so it adds nothing to the budget. +// Declared only, like Fp8LinearA8ProductionSchedule above it: a geometry that has never been +// measured on this route must be a build error, not a silent opt-in to another problem's schedule. +template +struct Fp8LinearA8TmaSchedule; + +// 256-token tile, four stages. Four is the deepest that fits: 4 * (256 + 128) * 64 = 96 KiB against +// the 99 KiB per-CTA cap, and five would need 120 KiB. One CTA per SM follows from that budget. +using Fp8A8TmaMeasuredSchedule = Fp8A8TmaSchedule<256, 4, 1>; + +template <> +struct Fp8LinearA8TmaSchedule { + using Type = Fp8A8TmaMeasuredSchedule; +}; + +template <> +struct Fp8LinearA8TmaSchedule { + using Type = Fp8A8TmaMeasuredSchedule; +}; + +template <> +struct Fp8LinearA8TmaSchedule { + using Type = Fp8A8TmaMeasuredSchedule; +}; + +template <> +struct Fp8LinearA8TmaSchedule { + using Type = Fp8A8TmaMeasuredSchedule; +}; + +template <> +struct Fp8LinearA8TmaSchedule { + using Type = Fp8A8TmaMeasuredSchedule; +}; + +// Multiprocessor count of the device this thread will launch on - the current device, which is +// where the launch that follows goes. Cached per device ordinal: the count steers which kernel +// runs, and a process that touches two different GPUs must not steer the second one with the +// first one's number. Returning zero declines the route, which is always a safe answer. +inline std::int32_t fp8_a8_multiprocessor_count() { + static std::array, kFp8A8MaxDevices> cache{}; + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess || device < 0 || device >= kFp8A8MaxDevices) { + return 0; + } + const std::int32_t cached = cache[device].load(std::memory_order_acquire); + if (cached != 0) { return cached; } + int value = 0; + if (cudaDeviceGetAttribute(&value, cudaDevAttrMultiProcessorCount, device) != cudaSuccess) { + return 0; + } + cache[device].store(value, std::memory_order_release); + return value; +} + +// Which of the two routes is cheaper at this width. +// +// Both kernels tile the same problem and both leave part of a wave idle at the end, but they +// quantise differently: the TMA route runs one CTA per SM over a 256-token tile, the cp.async route +// two CTAs per SM over a 64-token tile. So the decision is a comparison of two quantised costs, not +// a score for one of them - which also means the multiprocessor count enters the model on both +// sides instead of being frozen into a fitted constant. +// +// Cost of a route is (waves it needs) x (work one CTA does). Work per CTA is proportional to its +// token tile, so the tile widths carry it and cancel into kFp8A8TmaWorkRatio below. + +// Time the TMA route takes per token of work, relative to the route it replaces. An empirical +// constant for one part, solved from the widest measured point on an RTX 5090 (sm_120a, 170 SMs); +// it is not portable and not a claim about either kernel in general. What the model represents is +// wave quantisation and nothing else - it has no term for pipeline fill, for K, or for the cost of +// an epilogue, and the two width bounds below exist because of shapes it therefore cannot tell +// apart. The sweeps behind the constant and both bounds are in +// docs/maintainer/fp8-a8-tma-route.md. +inline constexpr double kFp8A8TmaWorkRatio = 0.936; + +// How much cheaper the model must find the TMA route before the route is taken. This is a margin on +// modelled cost, not on measured time: the model is a wave count times a tile width, so it is +// coarse, and a decision it calls within two percent is a decision it has not really made. Widths +// whose measured gain is smaller than this are still admitted - the margin buys confidence in the +// comparison, not a floor on the payoff. +inline constexpr double kFp8A8TmaMargin = 0.02; + +// Which of the two routes the model calls cheaper, as a pure function of shape, width and part. +// +// This is the whole decision and it is written once. The runtime guard below calls it with the +// multiprocessor count it reads; the coverage checks call it with the count of the device the +// sweeps were taken on. Both therefore ask the same arithmetic about the same instantiated tiles, +// which a second copy written out in terms of loose integers would not. +template +constexpr bool fp8_a8_tma_cheaper(std::int32_t tokens, std::int64_t multiprocessors) { + const std::int64_t tma_blocks = fp8_a8_tma_blocks(tokens); + const std::int64_t mma_tiles = + (static_cast(tokens) + MmaSchedule::kBlockTokens - 1) / + MmaSchedule::kBlockTokens; + const std::int64_t mma_blocks = + static_cast(Geometry::kOutputRows / MmaSchedule::kBlockRows) * mma_tiles; + const std::int64_t mma_slots = + static_cast(MmaSchedule::kMinBlocksPerSm) * multiprocessors; + const std::int64_t tma_slots = + static_cast(TmaSchedule::kMinBlocksPerSm) * multiprocessors; + const std::int64_t tma_waves = (tma_blocks + tma_slots - 1) / tma_slots; + const std::int64_t mma_waves = (mma_blocks + mma_slots - 1) / mma_slots; + // Cost is waves times the work one SM carries through a wave - not waves alone. The two routes + // put different amounts of work on an SM at once, and each route's occupancy comes from its own + // schedule rather than from a number written here: the shipped TMA schedule places one CTA of + // BlockM tokens, the cp.async schedule two CTAs of BlockTokens each. Comparing wave counts + // without that weight makes the wider tile look free, which is exactly backwards - and reading + // one route's occupancy from its schedule while fixing the other's in the model is how a future + // schedule change silently stops being modelled. + const double tma = + static_cast(tma_waves * TmaSchedule::kMinBlocksPerSm * TmaSchedule::kBlockM) * + kFp8A8TmaWorkRatio; + const double mma = + static_cast(mma_waves * MmaSchedule::kMinBlocksPerSm * MmaSchedule::kBlockTokens); + return tma < mma * (1.0 - kFp8A8TmaMargin); +} + +// The widest token count the route is allowed to take, per geometry. +// +// The model counts blocks along output rows and token tiles and has no term for K, so it gives one +// verdict for the two 5120-row residual shapes. Measured, they diverge above 4096. The bound is a +// measurement, stated per geometry so that an unswept shape cannot inherit it; the sweep is in the +// document above. +template +inline constexpr std::int32_t kFp8A8TmaMaxTokens = std::numeric_limits::max(); +template <> +inline constexpr std::int32_t kFp8A8TmaMaxTokens = 4096; + +// The narrowest width the route is offered at, for all geometries. +// +// This is a bound on the constant, not on the model. kFp8A8TmaWorkRatio was solved at and above the +// widths the product runs and has no validation below them, so the floor goes where the calibration +// starts. The model itself does better here than that suggests - at six sampled points below 1024 +// it agrees with the measurement at five - and the floor accordingly gives up measured gains of 8 +// to 12 percent to avoid one measured loss of 2. Whether it should move is a question for a wider +// sweep than the one behind it; both are in the document above. +inline constexpr std::int32_t kFp8A8TmaMinTokens = 1024; + +// The same question asked of a geometry, with both of its registered schedules filled in, and the +// measured width bound applied on top of the modelled cost. +template +constexpr bool fp8_a8_tma_admits(std::int32_t tokens, std::int64_t multiprocessors) { + if (tokens < kFp8A8TmaMinTokens) { return false; } + if (tokens > kFp8A8TmaMaxTokens) { return false; } + return fp8_a8_tma_cheaper::Type, + typename Fp8LinearA8ProductionSchedule::Type>( + tokens, multiprocessors); +} + +// The device the width bounds and the ratio were measured on. It is here so the coverage checks +// below ask their question against a fixed machine rather than against whatever built the tree. +inline constexpr std::int64_t kFp8A8CalibrationSms = 170; + +// A width is covered if the runtime would take it. The runtime adds representability, the address +// test and the grid.y limit on top of this, and none of those depend on the width. +template +constexpr bool fp8_a8_tma_covers(std::int32_t tokens) { + return fp8_a8_tma_admits(tokens, kFp8A8CalibrationSms); +} + +// The widths the numerical tests run at, pinned so they cannot quietly stop covering the route. +// These assert coverage and nothing else: that the predicate still sends these widths down the TMA +// path, so the tests that run them are testing the kernel they were extended for. They are not a +// statement that the route is faster - a predicate answer cannot carry that - and re-tuning the +// constant is expected to trip one of them. The fix then is to re-pick the test width from the new +// frontier and re-run the tests on it, not to widen the constant until the assertion passes. +// +// tests/ops/linear/test_fp8_a8.cpp and tests/ops/linear_add/test_fp8.cpp check values against a +// host reference and cannot see which kernel produced them, so without these the suite would stay +// green over a route that had disappeared. One aligned width and one leaving a partial trailing +// tile per geometry, the trailing tile being the majority of the admitted set. +static_assert(fp8_a8_tma_covers(4096) && + fp8_a8_tma_covers(4288) && + fp8_a8_tma_covers(1345), + "test_fp8_a8.cpp covers attn_input at 4096, 4288 and 1345"); +static_assert(fp8_a8_tma_covers(4096) && + fp8_a8_tma_covers(4160) && + fp8_a8_tma_covers(1153), + "test_fp8_a8.cpp covers gdn_input at 4096, 4160 and 1153"); +static_assert(fp8_a8_tma_covers(4096) && + fp8_a8_tma_covers(4288) && + fp8_a8_tma_covers(1153), + "test_fp8_a8.cpp covers mlp_gate_up at 4096, 4288 and 1153"); +static_assert(fp8_a8_tma_covers(1664) && + fp8_a8_tma_covers(4001), + "test_fp8_a8.cpp covers residual 5120x6144 at 1664, both fp8 tests at 4001"); +// attn_input reaches the route at the shipped default chunk, which is also a width its test +// runs. Every other route-taking test width is pinned below; without this one a retune could +// drop 1024 off the route and leave that test green over the kernel it was added to exercise. +static_assert(fp8_a8_tma_covers(1024), + "attn_input T=1024 is a route-taking test width"); +static_assert(fp8_a8_tma_covers(4160) && + fp8_a8_tma_covers(4001), + "both fp8 tests cover residual 5120x17408 at 4160 and 4001"); + +// The two residual geometries differ only in K and the model has no K term, so the pair of +// coverage checks above is one decision written twice. They are written out separately anyway: the +// day a schedule is registered per geometry is the day they stop agreeing. + +// Guard for the TMA route. +// +// Seven conditions, in the order fp8_a8_tma_applies runs them, which is the order of what they +// protect: first what the hardware cannot describe, then what the calibration cannot speak +// for, then what the launch geometry cannot carry, and only last a question of speed. +// +// TMA descriptors carry a global address, and cuTensorMapEncodeTiled rejects one it cannot +// describe. The activation codes are a workspace this op allocated, so they are aligned by +// construction. The weight codes are a byte offset into the loaded artifact, so their alignment is +// a property of the layout and not something this patch establishes. Declining is the whole +// remedy - the route being replaced reads the same bytes with no such requirement - and it keeps a +// new failure mode out of a patch that is supposed to change only speed. +// +// This is a runtime address test and lives only here. fp8_a8_tma_admits and fp8_a8_tma_covers stay +// as they are: they are compile-time statements about widths, and an address is not a width. +inline constexpr std::uintptr_t kFp8A8TmaAddressAlignment = 16; + +inline bool fp8_a8_tma_addresses_admit(const void* activation_codes, const void* weight_codes) { + const auto a = reinterpret_cast(activation_codes); + const auto b = reinterpret_cast(weight_codes); + if (a == 0 || b == 0) { return false; } + return (a % kFp8A8TmaAddressAlignment) == 0 && (b % kFp8A8TmaAddressAlignment) == 0; +} + +// Representability is a property of the geometry. The branch below is documentation, not a +// fallback: the call sites test this predicate at runtime, so the kernel template is instantiated +// whatever it answers, and a geometry the kernel cannot tile fails its own static_assert first. A +// geometry is admitted to this route by being registered in Fp8LinearA8TmaSchedule above, and that +// list is declared-only for the same reason. +// +// There is no condition on the width being a whole number of cp.async token tiles, and there used +// to be: it bought bit-identity with the route being replaced, and it was removed because it +// blocked widths the cost model takes and the route measures faster at. It could not do otherwise. +// The model's only width-dependent inputs are the two tile counts, both constant between adjacent +// multiples of the smaller tile, so its verdict is already fixed across a band and a condition on +// the width within that band can only subtract from a decision already made. The sweep is in +// docs/maintainer/fp8-a8-tma-route.md. +// +// So it is a question of speed, and of nothing else. +template ::Type> +bool fp8_a8_tma_applies(std::int32_t tokens, const void* activation_codes, + const void* weight_codes) { + // The cost model below compares this route against the one that would otherwise run, so the + // schedule it names has to be that one. Pinning it here means a future per-op schedule + // override cannot silently make the comparison quote a tile nobody falls back to. + static_assert( + std::is_same_v::Type>, + "the cost model must name the schedule this route actually falls back to"); + if constexpr (!kFp8A8TmaRepresentable) { + return false; + } else { + // Not a kernel requirement - the copy zero-fills past the extent and the store drops + // those rows, so one partial tile is fine. It is a statement about the calibration: the + // ratio was solved at and above kBlockM-sized widths, and a floor below the tile would put + // the model somewhere it was never fitted. + static_assert(kFp8A8TmaMinTokens >= TmaSchedule::kBlockM, + "the floor must stay inside the range the ratio was solved in"); + // Addresses before widths: a pointer the descriptor cannot describe is a hard failure in + // the launcher, where a width that does not suit is only a slower route. + if (!fp8_a8_tma_addresses_admit(activation_codes, weight_codes)) { return false; } + if (tokens < kFp8A8TmaMinTokens) { return false; } + if (tokens > kFp8A8TmaMaxTokens) { return false; } + // The launcher puts token tiles on grid.y, which tops out at 65535 where the cp.async route + // linearises into grid.x and does not. No real width comes near - it would take 16.7 M + // tokens in one chunk - but the limit is the new route's alone, and exceeding it is a + // launch failure rather than a slow answer, so it is cheaper to decline than to explain + // later. + if (fp8_a8_tma_token_tiles(tokens) > 65535) { return false; } + const std::int32_t multiprocessors = fp8_a8_multiprocessor_count(); + if (multiprocessors <= 0) { return false; } + return fp8_a8_tma_cheaper(tokens, multiprocessors); + } +} + } // namespace ninfer::ops::detail diff --git a/src/ops/linear/fp8/fp8_a8_tma.cuh b/src/ops/linear/fp8/fp8_a8_tma.cuh new file mode 100644 index 0000000000..54198d12c0 --- /dev/null +++ b/src/ops/linear/fp8/fp8_a8_tma.cuh @@ -0,0 +1,478 @@ +#pragma once + +// Row-scaled E4M3 weight x row-scaled E4M3 activation GEMM, staged by TMA with warp +// specialization. Same arithmetic as fp8_a8_mma.cuh: the accumulator, the m16n8k32 instruction, +// the scale application and the epilogue are unchanged. Only the staging differs - dedicated +// producer threads issue cp.async.bulk.tensor under an mbarrier ring, so one CTA per SM with a +// large token tile replaces two CTAs with a small one. + +#include "ops/common/math.cuh" +#include "ops/common/mbarrier.cuh" +#include "ops/common/memory.cuh" +#include "ops/common/mma.cuh" +#include "ops/linear/fp8/fp8_a8_mma.cuh" // Fp8MmaIdentityRows +#include "ops/linear/fp8/fp8_output.cuh" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace ninfer::ops::detail { + +// Highest device ordinal this route keeps per-device state for. Two things are cached by ordinal - +// the shared-memory opt-in in the launcher below, and the multiprocessor count the routing guard +// reads - and both must agree on the bound, or one of them silently stops caching where the other +// still does. +inline constexpr int kFp8A8MaxDevices = 16; + +struct alignas(128) Fp8A8TmaDescriptors { + CUtensorMap a_codes; + CUtensorMap b_codes; +}; + +inline void fp8_check_driver(CUresult status, const char* operation) { + if (status == CUDA_SUCCESS) { return; } + const char* name = nullptr; + (void)cuGetErrorName(status, &name); + throw std::runtime_error(std::string(operation) + ": " + + (name != nullptr ? name : "CUDA error")); +} + +inline CUtensorMap fp8_make_tma_2d(void* address, std::uint64_t columns, std::uint64_t rows, + std::uint64_t row_stride_bytes, std::uint32_t box_columns, + std::uint32_t box_rows, const char* operation) { + CUtensorMap map{}; + const std::uint64_t global_dim[] = {columns, rows}; + const std::uint64_t global_stride[] = {row_stride_bytes}; + const std::uint32_t box_dim[] = {box_columns, box_rows}; + const std::uint32_t element_stride[] = {1, 1}; + fp8_check_driver( + cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, address, global_dim, + global_stride, box_dim, element_stride, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_64B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + operation); + return map; +} + +// A is [tokens, K] and B is [output rows, K], both one byte per element and K-contiguous. The +// innermost box is the 64-byte K tile that 64B swizzle expects. +template +Fp8A8TmaDescriptors make_fp8_a8_tma_descriptors(const std::uint8_t* activation_codes, + const std::uint8_t* weight_codes, + std::int32_t tokens) { + constexpr std::uint32_t kRowBytes = 64; + // A paired block computes both branches, so its weight box is one branch tall and the + // producer issues one request per branch. + constexpr std::uint32_t kWeightBoxRows = PairRows ? BlockN / 2 : BlockN; + Fp8A8TmaDescriptors descriptors{}; + descriptors.a_codes = + fp8_make_tma_2d(const_cast(activation_codes), Geometry::kInputRows, + static_cast(tokens), Geometry::kInputRows, kRowBytes, BlockM, + "encode fp8 activation codes TMA"); + descriptors.b_codes = fp8_make_tma_2d( + const_cast(weight_codes), Geometry::kInputRows, Geometry::kOutputRows, + Geometry::kInputRows, kRowBytes, kWeightBoxRows, "encode fp8 weight codes TMA"); + return descriptors; +} + +// The route's shape, exported so the host can reason about a launch without restating it. The +// admission model needs the same grid the launcher builds; when those were written out twice, in +// two files, nothing would have caught them drifting apart. +template +constexpr std::int64_t fp8_a8_tma_token_tiles(std::int32_t tokens) { + return (static_cast(tokens) + Schedule::kBlockM - 1) / Schedule::kBlockM; +} + +template +constexpr std::int64_t fp8_a8_tma_blocks(std::int32_t tokens) { + return static_cast(Geometry::kOutputRows / Schedule::kBlockN) * + fp8_a8_tma_token_tiles(tokens); +} + +// Shapes the kernel can express at all. This is a property of the geometry, not a tuning decision, +// and the kernel asserts the same conditions - so it is stated once and consumed by both. +template +constexpr bool kFp8A8TmaRepresentable = (Geometry::kOutputRows % Schedule::kBlockN) == 0 && + (Geometry::kInputRows % Schedule::kBlockK) == 0; + +template +struct Fp8A8TmaSchedule { + static_assert(BlockM == 64 || BlockM == 128 || BlockM == 256); + static_assert(Stages >= 2 && Stages <= 8); + // Not a free parameter. fp8_a8_tma_cheaper multiplies this by the multiprocessor count to get + // the slots a wave fills, so a value above 1 divides the predicted wave count while the + // per-wave work stays whole, and the route starts winning its own comparison at every partial + // wave. Raising it means checking the shared-memory budget against it and re-fitting the ratio. + static_assert(MinBlocksPerSm == 1, + "the cost model reads this as achieved occupancy; see fp8_a8_tma_cheaper"); + + static constexpr int kBlockM = BlockM; + static constexpr int kBlockN = 128; + static constexpr int kBlockK = 64; + static constexpr int kRowBytes = kBlockK; + static constexpr int kStages = Stages; + static constexpr int kWarpsM = 4; + static constexpr int kWarpsN = 2; + static constexpr int kConsumerWarps = kWarpsM * kWarpsN; + static constexpr int kConsumerThreads = kConsumerWarps * 32; + // One warp produces. Only lane 0 issues the TMA descriptors; the rest of a warp is the + // smallest unit the block can be built from. An earlier form padded this to a full 128-thread + // warpgroup so that setmaxnreg could hand the producer's registers to the consumers, but the + // kernel allocates 166 registers with no spills, so there was nothing for the donation to buy - + // and relocatable device code discards the transfer anyway, which is why the build file keeps + // the nvfp4 warp-specialized kernels out of the RDC archive. Dropping to one warp measured + // identical (1.0004 / 0.9997 / 0.9996 against the padded form) and removes that conflict. + static constexpr int kProducerThreads = 32; + // Raising this to a warpgroup would re-arm the setmaxnreg pair below, and this kernel is + // compiled into a relocatable archive where that transfer is discarded. Keep it a build error. + static_assert( + kProducerThreads == 32, + "a warpgroup producer re-arms setmaxnreg, which relocatable device code discards"); + static constexpr int kThreads = kConsumerThreads + kProducerThreads; + static constexpr int kWarpM = kBlockM / kWarpsM; + static constexpr int kWarpN = kBlockN / kWarpsN; + static constexpr int kMmaM = kWarpM / 16; + static constexpr int kMmaN = kWarpN / 8; + static constexpr int kMmaKPerStage = kBlockK / 32; + static constexpr int kMinBlocksPerSm = MinBlocksPerSm; + + static_assert(kWarpM % 16 == 0 && kWarpN % 8 == 0); +}; + +// The 64B swizzle pattern repeats over eight rows of 64 bytes. fp8_tma_shared_byte reads the +// segment out of a row index taken relative to the tile base, which is the hardware's mapping only +// where that base is itself at the start of an atom - so the requirement is 512 bytes, and 128, +// which is all the TMA store itself needs, is not enough to state it. Every tile base has to hold +// it, not just the first: the stage stride and, for a paired block, the second branch's offset +// both land inside this storage. +inline constexpr std::size_t kFp8A8TmaSwizzleAtomBytes = 512; + +template +struct alignas(kFp8A8TmaSwizzleAtomBytes) Fp8A8TmaTensorStorage { + alignas(kFp8A8TmaSwizzleAtomBytes) + std::uint8_t a_codes[Schedule::kStages][Schedule::kBlockM * Schedule::kRowBytes]; + alignas(kFp8A8TmaSwizzleAtomBytes) + std::uint8_t b_codes[Schedule::kStages][Schedule::kBlockN * Schedule::kRowBytes]; + + static_assert((Schedule::kBlockM * Schedule::kRowBytes) % kFp8A8TmaSwizzleAtomBytes == 0, + "each activation stage must begin on a swizzle atom"); + static_assert((Schedule::kBlockN * Schedule::kRowBytes) % kFp8A8TmaSwizzleAtomBytes == 0, + "each weight stage must begin on a swizzle atom"); + // A paired block loads its second branch at half the weight stage. That offset is a tile base + // as much as the stage is, so it carries the same requirement. + static_assert(((Schedule::kBlockN / 2) * Schedule::kRowBytes) % kFp8A8TmaSwizzleAtomBytes == 0, + "the second branch of a paired block must begin on a swizzle atom"); +}; + +template +union alignas(kFp8A8TmaSwizzleAtomBytes) Fp8A8TmaScratch { + Fp8A8TmaTensorStorage tensors; + __nv_bfloat16 output[Schedule::kBlockM * (Schedule::kBlockN + 8)]; +}; + +template +struct alignas(kFp8A8TmaSwizzleAtomBytes) Fp8A8TmaSharedStorage { + Fp8A8TmaScratch scratch; + alignas(8) std::uint64_t full[Schedule::kStages]; + alignas(8) std::uint64_t empty[Schedule::kStages]; +}; + +__device__ __forceinline__ void fp8_tma_load_2d(void* destination, const CUtensorMap* descriptor, + std::int32_t coordinate0, std::int32_t coordinate1, + std::uint64_t* barrier) { + asm volatile("cp.async.bulk.tensor.2d.shared::cta.global.tile.mbarrier::complete_tx::bytes " + "[%0], [%1, {%2, %3}], [%4];" + : + : "r"(smem_addr(destination)), "l"(descriptor), "r"(coordinate0), "r"(coordinate1), + "r"(smem_addr(barrier)) + : "memory"); +} + +// 64B swizzle places the sixteen-byte segment of a row at segment XOR ((row / 2) % 4). The +// pattern closes after eight rows, so the unit it repeats over is kFp8A8TmaSwizzleAtomBytes, and +// this mapping is the hardware's only for a tile whose base is at the start of one. The row here +// is relative to that base; Fp8A8TmaTensorStorage is what makes the base sit where that is true. +__device__ __forceinline__ int fp8_tma_shared_byte(int row, int logical_byte) { + return ((logical_byte >> 4) ^ ((row >> 1) & 3)) * 16 + (logical_byte & 15); +} + +// The work distributor hands CTAs to SMs in linear order with blockIdx.x fastest, so the stock +// grid - x over weight-row tiles, y over token tiles - puts a different weight tile in every CTA +// that runs at the same time, and the whole weight matrix is re-read once per token tile. Walking +// the token index fastest instead makes the CTAs that share a weight tile run together, and the +// matrix is read once. Every other route here already makes that choice: the cp.async schedules +// this one falls back to are registered with Fp8MmaRaster::TokenFast, and nvfp4_w4a4_tma.cuh +// carries the same arithmetic. A shared home for it is a follow-up; putting one here would edit +// NVFP4 files this change has no other reason to touch. +__device__ __forceinline__ void fp8_tma_raster_blocks(int& block_x, int& block_y) { + const int rows = static_cast(gridDim.y); + const int linear = + static_cast(blockIdx.y) * static_cast(gridDim.x) + static_cast(blockIdx.x); + block_y = linear % rows; + block_x = linear / rows; +} + +template +__global__ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void fp8_a8_tma_kernel( + const __grid_constant__ Fp8A8TmaDescriptors descriptors, + const float* __restrict__ activation_scales, const __nv_bfloat16* __restrict__ weight_scales, + std::int32_t tokens, const __grid_constant__ Epilogue epilogue, + const __grid_constant__ Output output, const __grid_constant__ RowPolicy row_policy) { + static_assert(kFp8A8TmaRepresentable, + "the geometry must tile evenly in both K and output rows"); + // The weight scales are fetched one 32-bit word at a time at parent_row0 and applied to + // parent_row0 and parent_row1, which is right only while the row policy maps adjacent local + // rows to adjacent parent rows. Fp8SwiGluRows does not, at its branch boundary: there it jumps + // by the intermediate size. local_row0 is always even, so a pair can only straddle that + // boundary if the branch width is odd, and the branch width is kBlockN / 2 - hence % 4, not + // % 2. It cannot fire while kBlockN is fixed at 128; it is here so the invariant is written + // down where the pair is loaded rather than rediscovered later. + // The unpaired case needs the identity policy for an unrelated reason: the producer loads + // B and the store writes at the raw row while the epilogue would apply the mapped one, and + // nothing reconciles the two. + static_assert(PairRows || std::is_same_v, + "a non-identity row policy is only wired up for the paired store"); + static_assert(!PairRows || (Schedule::kBlockN % 4) == 0, + "paired rows need an even branch width so a scale pair cannot straddle it"); + + // The alignment the swizzle indexing needs, declared on the allocation that backs it. The + // TMA store alone would be satisfied by 128. + extern __shared__ __align__(kFp8A8TmaSwizzleAtomBytes) unsigned char shared_bytes[]; + auto& shared = *reinterpret_cast*>(shared_bytes); + constexpr int kBranchRows = PairRows ? Schedule::kBlockN / 2 : Schedule::kBlockN; + int block_x = 0; + int block_y = 0; + fp8_tma_raster_blocks(block_x, block_y); + const int token_begin = block_y * Schedule::kBlockM; + const int row_begin = block_x * kBranchRows; + + if (threadIdx.x == 0) { +#pragma unroll + for (int stage = 0; stage < Schedule::kStages; ++stage) { + cta_mbarrier_init(&shared.full[stage], 1); + cta_mbarrier_init(&shared.empty[stage], Schedule::kConsumerWarps); + } + cta_mbarrier_fence_init(); + } + __syncthreads(); + + constexpr int kKTiles = Geometry::kInputRows / Schedule::kBlockK; + + if (threadIdx.x < Schedule::kProducerThreads) { + if (threadIdx.x == 0) { +#pragma unroll 1 + for (int k_tile = 0; k_tile < kKTiles; ++k_tile) { + const int stage = k_tile % Schedule::kStages; + const std::uint32_t empty_phase = 1U ^ ((k_tile / Schedule::kStages) & 1U); + cta_mbarrier_wait(&shared.empty[stage], empty_phase); + constexpr std::uint32_t kTransactionBytes = + (Schedule::kBlockM + Schedule::kBlockN) * Schedule::kRowBytes; + cta_mbarrier_arrive_expect_tx(&shared.full[stage], kTransactionBytes); + + auto& tensors = shared.scratch.tensors; + fp8_tma_load_2d(tensors.a_codes[stage], &descriptors.a_codes, + k_tile * Schedule::kRowBytes, token_begin, &shared.full[stage]); + if constexpr (PairRows) { + fp8_tma_load_2d(tensors.b_codes[stage], &descriptors.b_codes, + k_tile * Schedule::kRowBytes, + row_policy.weight_row(row_begin, 0), &shared.full[stage]); + fp8_tma_load_2d(tensors.b_codes[stage] + kBranchRows * Schedule::kRowBytes, + &descriptors.b_codes, k_tile * Schedule::kRowBytes, + row_policy.weight_row(row_begin, kBranchRows), + &shared.full[stage]); + } else { + fp8_tma_load_2d(tensors.b_codes[stage], &descriptors.b_codes, + k_tile * Schedule::kRowBytes, row_begin, &shared.full[stage]); + } + } + } + return; + } + + auto& tensors = shared.scratch.tensors; + const int consumer_thread = static_cast(threadIdx.x) - Schedule::kProducerThreads; + const int lane = consumer_thread & 31; + const int warp = consumer_thread >> 5; + const int warp_m = warp / Schedule::kWarpsN; + const int warp_n = warp - warp_m * Schedule::kWarpsN; + + const int a_matrix = lane >> 3; + const int a_row_offset = (lane & 7) + ((a_matrix & 1) << 3); + const int a_column_byte = (a_matrix >> 1) * 16; + const int b_row_offset = lane & 7; + const int b_column_byte = ((lane >> 3) & 1) * 16; + + float accumulators[Schedule::kMmaM][Schedule::kMmaN][4] = {}; +#pragma unroll 1 + for (int k_tile = 0; k_tile < kKTiles; ++k_tile) { + const int stage = k_tile % Schedule::kStages; + const std::uint32_t full_phase = (k_tile / Schedule::kStages) & 1U; + cta_mbarrier_wait(&shared.full[stage], full_phase); + +#pragma unroll + for (int k_step = 0; k_step < Schedule::kMmaKPerStage; ++k_step) { + unsigned a_fragments[Schedule::kMmaM][4]; + unsigned b_fragments[Schedule::kMmaN][2]; +#pragma unroll + for (int mma_m = 0; mma_m < Schedule::kMmaM; ++mma_m) { + const int row = warp_m * Schedule::kWarpM + mma_m * 16 + a_row_offset; + const auto* address = tensors.a_codes[stage] + row * Schedule::kRowBytes + + fp8_tma_shared_byte(row, k_step * 32 + a_column_byte); + ldmatrix_x4(a_fragments[mma_m][0], a_fragments[mma_m][1], a_fragments[mma_m][2], + a_fragments[mma_m][3], smem_addr(address)); + } +#pragma unroll + for (int mma_n = 0; mma_n < Schedule::kMmaN; ++mma_n) { + const int row = warp_n * Schedule::kWarpN + mma_n * 8 + b_row_offset; + const auto* address = tensors.b_codes[stage] + row * Schedule::kRowBytes + + fp8_tma_shared_byte(row, k_step * 32 + b_column_byte); + ldmatrix_x2(b_fragments[mma_n][0], b_fragments[mma_n][1], smem_addr(address)); + } +#pragma unroll + for (int mma_m = 0; mma_m < Schedule::kMmaM; ++mma_m) { +#pragma unroll + for (int mma_n = 0; mma_n < Schedule::kMmaN; ++mma_n) { + mma_fp8_e4m3(accumulators[mma_m][mma_n][0], accumulators[mma_m][mma_n][1], + accumulators[mma_m][mma_n][2], accumulators[mma_m][mma_n][3], + a_fragments[mma_m][0], a_fragments[mma_m][1], + a_fragments[mma_m][2], a_fragments[mma_m][3], + b_fragments[mma_n][0], b_fragments[mma_n][1]); + } + } + } + if (lane == 0) { cta_mbarrier_arrive(&shared.empty[stage]); } + } + + // The epilogue reuses the tensor pipeline's storage, so every consumer warp must finish its + // last tensor read before any warp starts overwriting it. + asm volatile("bar.sync 1, %0;" : : "r"(Schedule::kConsumerThreads) : "memory"); + + constexpr int kOutputStride = Schedule::kBlockN + 8; + auto* shared_output = shared.scratch.output; + const int accumulator_token = lane >> 2; + const int accumulator_row = 2 * (lane & 3); +#pragma unroll + for (int mma_m = 0; mma_m < Schedule::kMmaM; ++mma_m) { + const int local_token0 = warp_m * Schedule::kWarpM + mma_m * 16 + accumulator_token; + const int local_token1 = local_token0 + 8; + const int token0 = token_begin + local_token0; + const int token1 = token_begin + local_token1; + // TMA zero-fills the rows past the end, so their accumulators are harmless; they are + // simply never stored. Nothing may be read on their behalf either: the activation scale is + // indexed by token, and so is the residual an epilogue may add, so both are gated on the + // row existing. Rows that do exist take exactly the arithmetic they took before. + const bool token0_valid = token0 < tokens; + const bool token1_valid = token1 < tokens; + const float activation_scale0 = token0_valid ? activation_scales[token0] : 0.0F; + const float activation_scale1 = token1_valid ? activation_scales[token1] : 0.0F; +#pragma unroll + for (int mma_n = 0; mma_n < Schedule::kMmaN; ++mma_n) { + const int local_row0 = warp_n * Schedule::kWarpN + mma_n * 8 + accumulator_row; + const int parent_row0 = row_policy.weight_row(row_begin, local_row0); + const int parent_row1 = row_policy.weight_row(row_begin, local_row0 + 1); + const std::uint32_t scale_bits = load_vec(weight_scales + parent_row0); + const float2 weight_scale = bf16x2_bits_to_float2(scale_bits); + const float value00 = token0_valid + ? epilogue.apply(parent_row0, token0, + accumulators[mma_m][mma_n][0] * + activation_scale0 * weight_scale.x) + : 0.0F; + const float value01 = token0_valid + ? epilogue.apply(parent_row1, token0, + accumulators[mma_m][mma_n][1] * + activation_scale0 * weight_scale.y) + : 0.0F; + const float value10 = token1_valid + ? epilogue.apply(parent_row0, token1, + accumulators[mma_m][mma_n][2] * + activation_scale1 * weight_scale.x) + : 0.0F; + const float value11 = token1_valid + ? epilogue.apply(parent_row1, token1, + accumulators[mma_m][mma_n][3] * + activation_scale1 * weight_scale.y) + : 0.0F; + auto* destination0 = reinterpret_cast<__nv_bfloat162*>( + shared_output + local_token0 * kOutputStride + local_row0); + auto* destination1 = reinterpret_cast<__nv_bfloat162*>( + shared_output + local_token1 * kOutputStride + local_row0); + *destination0 = __floats2bfloat162_rn(value00, value01); + *destination1 = __floats2bfloat162_rn(value10, value11); + } + } + + asm volatile("bar.sync 1, %0;" : : "r"(Schedule::kConsumerThreads) : "memory"); + constexpr int kVectorsPerRow = kBranchRows / 8; + constexpr int kOutputVectors = Schedule::kBlockM * kVectorsPerRow; + for (int task = consumer_thread; task < kOutputVectors; task += Schedule::kConsumerThreads) { + const int local_token = task / kVectorsPerRow; + if (token_begin + local_token >= tokens) { continue; } + const int row_vector = task - local_token * kVectorsPerRow; + const auto* row_base = shared_output + local_token * kOutputStride + row_vector * 8; + const uint4 values = load_vec(row_base); + if constexpr (PairRows) { + output.store_pair_vector(row_begin + row_vector * 8, token_begin + local_token, values, + load_vec(row_base + kBranchRows)); + } else { + output.store_vector(row_begin + row_vector * 8, token_begin + local_token, values); + } + } +} + +// One launch path for every Op family: they differ only in the epilogue and output policies. +template +void fp8_a8_tma_launch(const std::uint8_t* activation_codes, const float* activation_scales, + const std::uint8_t* weight_codes, const __nv_bfloat16* weight_scales, + std::int32_t tokens, Epilogue epilogue, Output output, cudaStream_t stream, + RowPolicy row_policy = {}) { + constexpr std::size_t kSharedBytes = sizeof(Fp8A8TmaSharedStorage); + static_assert(kSharedBytes <= 99 * 1024); + // Per device, not per process: the opt-in raises the dynamic shared limit on the current + // device only, so a process that reaches a second GPU would otherwise launch there without it + // and fail at 96 KiB. Keyed by ordinal, set once per device. + static std::array, kFp8A8MaxDevices> raised{}; + int attribute_device = 0; + // Throwing here, where the routing guard merely declines: the guard can fall back to the + // cp.async route and produce the right answer, but by this point the launch is committed and a + // kernel that needs 96 KiB would fail on the device we could not name. + if (cudaGetDevice(&attribute_device) != cudaSuccess || attribute_device < 0 || + attribute_device >= kFp8A8MaxDevices) { + throw std::runtime_error("fp8 TMA: cannot identify the current device"); + } + // Acquire/release rather than relaxed for consistency with the multiprocessor-count cache, + // which is the same shape of problem. Neither orders the driver-side effect itself - C++ + // atomics cannot - and neither needs to: setting the attribute twice is idempotent, and a + // thread that reads false when another has already set it simply sets it again. + if (!raised[attribute_device].load(std::memory_order_acquire)) { + const cudaError_t attribute = cudaFuncSetAttribute( + fp8_a8_tma_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(kSharedBytes)); + if (attribute != cudaSuccess) { + throw std::runtime_error(std::string("fp8 TMA shared memory attribute: ") + + cudaGetErrorName(attribute)); + } + raised[attribute_device].store(true, std::memory_order_release); + } + + const Fp8A8TmaDescriptors descriptors = + make_fp8_a8_tma_descriptors( + activation_codes, weight_codes, tokens); + const auto token_tiles = fp8_a8_tma_token_tiles(tokens); + const dim3 blocks(Geometry::kOutputRows / Schedule::kBlockN, + static_cast(token_tiles)); + fp8_a8_tma_kernel + <<>>( + descriptors, activation_scales, weight_scales, tokens, epilogue, output, row_policy); +} + +} // namespace ninfer::ops::detail diff --git a/src/ops/linear_add/fp8/fp8_linear_add_a8.cu b/src/ops/linear_add/fp8/fp8_linear_add_a8.cu index db246d9361..a281107464 100644 --- a/src/ops/linear_add/fp8/fp8_linear_add_a8.cu +++ b/src/ops/linear_add/fp8/fp8_linear_add_a8.cu @@ -44,7 +44,19 @@ void launch_mma(const Weight& weight, Tensor& residual, Fp8A8Workspace workspace template void launch_problem(const Weight& weight, Tensor& residual, Fp8A8Workspace workspace, std::int32_t tokens, cudaStream_t stream) { - using Schedule = typename Fp8LinearA8ProductionSchedule::Type; + using Schedule = typename Fp8LinearA8ProductionSchedule::Type; + using TmaSchedule = typename Fp8LinearA8TmaSchedule::Type; + if (fp8_a8_tma_applies(tokens, workspace.codes, + weight.qdata)) { + auto* output = static_cast<__nv_bfloat16*>(residual.data); + fp8_a8_tma_launch( + workspace.codes, workspace.scales, static_cast(weight.qdata), + static_cast(weight.scales), tokens, + Fp8AddResidualEpilogue{output, Geometry::kOutputRows}, + Fp8ContiguousOutput{output, Geometry::kOutputRows}, stream); + CUDA_CHECK(cudaGetLastError()); + return; + } if ((tokens % Schedule::kBlockTokens) == 0) { launch_mma(weight, residual, workspace, tokens, stream); } else { diff --git a/src/ops/linear_swiglu/fp8/fp8_linear_swiglu_a8.cu b/src/ops/linear_swiglu/fp8/fp8_linear_swiglu_a8.cu index c2d1e03db4..e9819ae4c8 100644 --- a/src/ops/linear_swiglu/fp8/fp8_linear_swiglu_a8.cu +++ b/src/ops/linear_swiglu/fp8/fp8_linear_swiglu_a8.cu @@ -4,6 +4,7 @@ #include "ops/linear/fp8/fp8_a8_mma.cuh" #include "ops/linear/fp8/fp8_a8_plan.h" #include "ops/linear/fp8/fp8_a8_schedule.cuh" +#include "ops/linear/fp8/fp8_a8_tma.cuh" #include "ops/linear/fp8/fp8_config.h" #include "ops/linear/fp8/fp8_output.cuh" #include "ops/linear_swiglu/fp8/fp8_linear_swiglu_output.cuh" @@ -54,6 +55,18 @@ void fp8_linear_swiglu_a8_launch(const Tensor& x, const Weight& weight, Tensor& const Fp8A8Workspace scratch = allocate_fp8_a8_workspace(workspace, x.ne[1], Geometry::kInputRows); launch_fp8_a8_quantize(x, weight, scratch, stream); + using TmaSchedule = typename Fp8LinearA8TmaSchedule::Type; + using TmaRows = Fp8SwiGluRows; + if (fp8_a8_tma_applies(x.ne[1], scratch.codes, weight.qdata)) { + fp8_a8_tma_launch( + scratch.codes, scratch.scales, static_cast(weight.qdata), + static_cast(weight.scales), x.ne[1], Fp8IdentityEpilogue{}, + Fp8SwiGluOutput{static_cast<__nv_bfloat16*>(out.data), kIntermediate}, stream, + TmaRows{}); + CUDA_CHECK(cudaGetLastError()); + return; + } if ((x.ne[1] % Schedule::kBlockTokens) == 0) { launch_mma(weight, out, scratch, x.ne[1], stream); } else { diff --git a/tests/ops/linear/test_fp8_a8.cpp b/tests/ops/linear/test_fp8_a8.cpp index 77144dbb84..b39c6eb032 100644 --- a/tests/ops/linear/test_fp8_a8.cpp +++ b/tests/ops/linear/test_fp8_a8.cpp @@ -10,6 +10,17 @@ using namespace ninfer; using namespace ninfer::test::linear; int run_fp8_a8() { + // Widths for the TMA-staged route, one aligned and one with a partial trailing tile per shape, + // read off the predicate rather than assumed: admission is a step function of two wave counts, + // so it differs per geometry and is not monotonic in token count. 4160 is declined on + // attn_input and mlp_gate_up, so those two take 4288; 5120x6144 is bounded at 4096, so + // its partial width has to come from below that, hence 1664. Partial trailing tiles are the + // common case among admitted widths, not an edge. + // + // The last width in each list is not a whole cp.async token tile. Those widths reach the route + // only because the multiple-of-tile condition was removed, and they are the ones whose output + // the previous kernel cannot be compared against byte for byte, so a host reference is the + // only thing that checks them. constexpr std::array attn_invocations{ Invocation{12, CallForm::Policy, ops::LinearPolicy::AllowA8}, Invocation{48, CallForm::Policy, ops::LinearPolicy::AllowA8}, @@ -17,6 +28,9 @@ int run_fp8_a8() { Invocation{65, CallForm::Policy, ops::LinearPolicy::AllowA8}, Invocation{1023, CallForm::Policy, ops::LinearPolicy::AllowA8}, Invocation{1024, CallForm::Policy, ops::LinearPolicy::AllowA8}, + Invocation{4096, CallForm::Policy, ops::LinearPolicy::AllowA8}, + Invocation{1345, CallForm::Policy, ops::LinearPolicy::AllowA8}, + Invocation{4288, CallForm::Policy, ops::LinearPolicy::AllowA8}, }; int failures = run_shape("FP8_A8", ActivationCompute::A8, make_fp8_weight, {14336, 5120, 829U, Comparison::Sampled, true, attn_invocations}); @@ -25,6 +39,9 @@ int run_fp8_a8() { Invocation{48, CallForm::Policy, ops::LinearPolicy::AllowA8}, Invocation{65, CallForm::Policy, ops::LinearPolicy::AllowA8}, Invocation{1024, CallForm::Policy, ops::LinearPolicy::AllowA8}, + Invocation{4096, CallForm::Policy, ops::LinearPolicy::AllowA8}, + Invocation{1153, CallForm::Policy, ops::LinearPolicy::AllowA8}, + Invocation{4160, CallForm::Policy, ops::LinearPolicy::AllowA8}, }; failures += run_shape("FP8_A8", ActivationCompute::A8, make_fp8_weight, {16384, 5120, 839U, Comparison::Sampled, true, gdn_invocations}); @@ -34,6 +51,9 @@ int run_fp8_a8() { Invocation{48, CallForm::Policy, ops::LinearPolicy::AllowA8}, Invocation{65, CallForm::Policy, ops::LinearPolicy::AllowA8}, Invocation{1024, CallForm::Policy, ops::LinearPolicy::AllowA8}, + Invocation{4096, CallForm::Policy, ops::LinearPolicy::AllowA8}, + Invocation{1153, CallForm::Policy, ops::LinearPolicy::AllowA8}, + Invocation{4288, CallForm::Policy, ops::LinearPolicy::AllowA8}, }; failures += run_shape("FP8_A8", ActivationCompute::A8, make_fp8_weight, {34816, 5120, 853U, Comparison::Sampled, true, mlp_invocations}); @@ -42,7 +62,10 @@ int run_fp8_a8() { Invocation{25, CallForm::Policy, ops::LinearPolicy::AllowA8}, Invocation{48, CallForm::Policy, ops::LinearPolicy::AllowA8}, Invocation{65, CallForm::Policy, ops::LinearPolicy::AllowA8}, + Invocation{1664, CallForm::Policy, ops::LinearPolicy::AllowA8}, Invocation{1024, CallForm::Policy, ops::LinearPolicy::AllowA8}, + Invocation{4001, CallForm::Policy, ops::LinearPolicy::AllowA8}, + Invocation{4096, CallForm::Policy, ops::LinearPolicy::AllowA8}, }; failures += run_shape("FP8_A8", ActivationCompute::A8, make_fp8_weight, {5120, 6144, 857U, Comparison::Sampled, true, residual6144_invocations}); @@ -51,6 +74,9 @@ int run_fp8_a8() { Invocation{48, CallForm::Policy, ops::LinearPolicy::AllowA8}, Invocation{65, CallForm::Policy, ops::LinearPolicy::AllowA8}, Invocation{1024, CallForm::Policy, ops::LinearPolicy::AllowA8}, + Invocation{4096, CallForm::Policy, ops::LinearPolicy::AllowA8}, + Invocation{4001, CallForm::Policy, ops::LinearPolicy::AllowA8}, + Invocation{4160, CallForm::Policy, ops::LinearPolicy::AllowA8}, }; failures += run_shape("FP8_A8", ActivationCompute::A8, make_fp8_weight, diff --git a/tests/ops/linear_add/test_fp8.cpp b/tests/ops/linear_add/test_fp8.cpp index b9f907f15f..466be1618a 100644 --- a/tests/ops/linear_add/test_fp8.cpp +++ b/tests/ops/linear_add/test_fp8.cpp @@ -106,11 +106,18 @@ int run_shape(std::int32_t n, std::int32_t k, std::int32_t first_a8, std::uint32 Invocation{96, ops::LinearPolicy::AllowA8}, Invocation{128, ops::LinearPolicy::AllowA8}, Invocation{129, ops::LinearPolicy::AllowA8}, + // The residual add is the one epilogue that folds into the store, so it is the one that + // has to be checked where the store drops rows. 4001 and 4160 leave the route a partial + // trailing tile, and 4001 is not a whole cp.async token tile - a width the route reaches + // only since that condition was removed. 5120x6144 is bounded at 4096 and declines 4160. + Invocation{4001, ops::LinearPolicy::AllowA8}, + Invocation{4096, ops::LinearPolicy::AllowA8}, + Invocation{4160, ops::LinearPolicy::AllowA8}, }; for (int columns = 2; columns <= 24; ++columns) { invocations.push_back({columns, ops::LinearPolicy::A16Only}); } - constexpr std::int32_t kMaximumTokens = 1024; + constexpr std::int32_t kMaximumTokens = 4160; quantized_weight::PackedWeight host_weight = quantized_weight::make_patterned_weight(QType::FP8_E4M3FN_ROW_BF16S, n, k, seed); const std::vector rows = sampled_indices(n); diff --git a/tests/ops/linear_swiglu/test_fp8.cpp b/tests/ops/linear_swiglu/test_fp8.cpp index 2741f65f5d..101fe438d9 100644 --- a/tests/ops/linear_swiglu/test_fp8.cpp +++ b/tests/ops/linear_swiglu/test_fp8.cpp @@ -10,7 +10,13 @@ int main() { try { constexpr std::array kA16Cases{1, 2, 4, 16, 128}; - constexpr std::array kA8Cases{1, 2, 3, 8, 16, 48, 64, 65, 96, 128, 1024}; + // The last three entries are the point of this list: below them the TMA-staged route + // declines, so without them the paired-rows instantiation this op is the only user of - + // two TMA loads per stage, a non-identity row policy - is executed by no test. 1153 is + // not a whole cp.async token tile, which is a case the route reaches only since the + // multiple-of-tile condition was removed. + constexpr std::array kA8Cases{1, 2, 3, 8, 16, 48, 64, + 65, 96, 128, 1024, 1153, 4096, 4288}; int failures = 0; failures += run_profile( "LinearSwiGLU FP8_A16", diff --git a/tests/ops/test_attn_input_proj.cpp b/tests/ops/test_attn_input_proj.cpp index 1d7f71598c..e8fa7b1564 100644 --- a/tests/ops/test_attn_input_proj.cpp +++ b/tests/ops/test_attn_input_proj.cpp @@ -382,7 +382,10 @@ int run_fp8_target() { std::cerr << "FP8 attention projection workspace interval mismatch\n"; ++failures; } - for (int t : {129, 144, 145, 160, 161, 192, 193, 256, 257, 1024}) + // 1024 and 4096 reach the TMA-staged route; 4288 leaves it a partial trailing tile, + // which this op's four-output store otherwise never sees; 1345 is not a whole cp.async + // token tile, a case the route reaches only since that condition was removed. + for (int t : {129, 144, 145, 160, 161, 192, 193, 256, 257, 1024, 1345, 4096, 4288}) failures += run_target_projection_case(parent, nullptr, t, policy); for (int t : {1, 4, 5, 6, 8, 9, 16, 24, 25, 32, 33, 34, 64, 65, 80, 81, 96, 97, 128, 129, 144, 145, 160, 161}) diff --git a/tests/ops/test_gdn_input_proj.cpp b/tests/ops/test_gdn_input_proj.cpp index b2a78aeb4e..b28148dc1d 100644 --- a/tests/ops/test_gdn_input_proj.cpp +++ b/tests/ops/test_gdn_input_proj.cpp @@ -309,7 +309,11 @@ int run_fp8() { failures += run_fp8_case(parent, 1, ops::LinearPolicy::A16Only, true); failures += run_fp8_case(parent, 2, ops::LinearPolicy::A16Only); - for (const std::int32_t tokens : {1, 2, 7, 8, 48, 65, 1024}) { + // 1153, 4096 and 4160 reach the TMA-staged route; below them it declines, so without them + // this op output type never runs there. 1153 and 4160 leave a partial trailing tile, and + // 1153 is not a whole cp.async token tile - a width the route reaches only since that + // condition was removed. + for (const std::int32_t tokens : {1, 2, 7, 8, 48, 65, 1024, 1153, 4096, 4160}) { failures += run_fp8_case(parent, tokens, ops::LinearPolicy::AllowA8); } return failures;