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;