perf(ArcKernels): blockwise-FP8 GEMM on tensor cores, scale promoted per 128-K

[heydryft]

Feeds the #194 release gate. Branched off master 3c22cb5b5.

🔴 UNVERIFIED ON HARDWARE — never run

No GPU was available for this wave. Nothing in this PR has executed. Every performance number below is a derivation from published machine limits, not a measurement. Do not quote any of it as a result. Rule 2 normally bans this; it was explicitly overridden for this wave.

The target

fp8_matmul_tiled is 66% of a B=256 decode step — 524 ms of 794. Its inner loop is

acc += s_input[ty][k] * s_weight[tx][k];

two shared-memory float loads to feed one FMA, no register blocking, no tensor-core instruction anywhere in it. Shared memory delivers 32 floats/clk/SM against 128 FP32 lanes, so that loop cannot exceed ~1/4 of the FP32 CUDA-core rate however it is tuned. It is a structural ceiling, not a tuning problem.

What this adds

fp8_matmul_wmma, alongside the old kernel: a tensor-core GEMM that keeps weights in FP8 and applies the [N/128, K/128] block scale by promoting an FP32 accumulator at every scale-block boundary along K — DeepGEMM's structure. K_BLK == block_size_x, so one K-tile is exactly one scale block and the scale is a single scalar per tile.

Numerically this is better than what it replaces, not a trade. The scalar kernel computes f32(act) * f32(w * scale) and rounds every product. Here the FP8 → fp16/bf16 weight conversion is exact (e4m3's 3 mantissa bits and 2^-9..2^8 exponent range are representable in both), the tensor core forms each product at full width into f32, and the scale is applied once per 128 K-elements instead of once per element.

Derived cost — a derivation, NOT a measurement

V4 at B=256: 7 sites/layer × 43 layers = 4.238 G params ⇒ FLOP = 2 × 256 × 4.238e9 = 2.170 TFLOP.

bound	rate	time
H200 FP32 non-tensor	67 TFLOP/s	32.4 ms — the scalar kernel's own bound; it achieves 6.8% of it
H200 BF16/FP16 tensor	989.5 TFLOP/s	2.19 ms ← this kernel's bound
H200 FP8 tensor	1979 TFLOP/s	1.10 ms
HBM over 4.238 GB	4.8 TB/s	0.88 ms

Compute-bound at a derived 2.19 ms floor against a measured 524 ms today. A WMMA GEMM of this shape typically realises 50–70% of peak, which would put it at a derived 3.1–4.4 ms. Arithmetic only; none of it has run.

The remaining 2× to 1.10 ms is not reachable by tuning this kernel — mma with e4m3 operands needs both operands in FP8, hence FP8 activations with their own per-token 128 scales. Named, not built.

Cheap checks done first, as instructed

cuBLASLt block scaling: supported on our target, unavailable at our toolkit pin. CUBLASLT_MATMUL_DESC_B_SCALE_MODE = CUBLASLT_MATMUL_MATRIX_SCALE_BLK128x128_32F is documented at compute capability 9.0 — Hopper, exactly our target (and notably not on Blackwell, which uses the MX VEC32_UE8M0 modes). But the enum does not exist before CUDA 12.9: absent in 12.4.1, absent in 12.8.0, present in 12.9.0 — and .github/workflows/cuda_compile_check.yaml pins 12.4.1. It also requires FP8 activations with VEC128_32F per-token scales and TN layout, so it is not the one-attribute change it first looks like. Worth revisiting if the toolkit pin moves.

wgmma is reachable but deliberately not used. cudaforge auto-suffixes sm_90 → sm_90a (cudaforge-0.1.5 compute_cap.rs::auto_suffix), which is the target wgmma requires, and it would be worth maybe another 1.3–1.5×. Not used because this file cannot be run before it is committed, and a hand-rolled wgmma descriptor that is subtly wrong returns plausible logits rather than an error. nvcuda::wmma fixes the fragment layout in the compiler, compiles for both sm_80 and sm_90a, and has a working precedent in-tree (kernels/mxfp4/mxfp4_gemm_wmma.cu). The wgmma rung belongs on a box that can run the A/B this PR sets up.

Kill switch — fast-path-default

(default)         -> tensor-core WMMA GEMM
ARC_NO_FP8_WMMA=1 -> the scalar fp8_matmul_tiled, i.e. exactly the
                     behaviour of the commit before this one

One binary, both arms, no rebuild. Run the control arm first.

No M threshold, deliberately

This codebase has frozen a dispatch threshold from a single measured point twice — ARC_FP8_CUBLAS_MIN_M = 512 (512 was the only M ever measured, so M = 5..511 fell through to the scalar kernel) and qtip::gather_policy's n >= 683. A threshold is a claim about every value it excludes. I have zero measured points, so inventing one would be strictly worse than either of those. Sweep on hardware, then add one if the sweep shows one.

Correctness guards worth a reviewer's eye

• Eligibility is read from the kernel, not duplicated in Rust. fp8_matmul_wmma_tile_dims exports N_BLK/K_BLK so fp8_wmma_eligible cannot drift from the tiling the kernel was compiled with. The failure mode of that drift is a numerically wrong GEMM, not a crash.
• Shared-memory pad is +8 and must stay a multiple of 8. wmma::load_matrix_sync requires a 16-byte-multiple leading dimension for 16-bit types; (128+8)*2 = 272 B = 17*16 satisfies it. This is the opposite of the right answer for the scalar kernel next door, where +1 breaks a bank conflict. A +1 pad here would be undefined. Called out in the source so nobody "harmonises" them.
• --use_fast_math, which side we're on: unaffected. No transcendentals, no division, no sqrt. -fmad=true only contracts the promotion into an FMA (more accurate). --ftz=true flushes f32 denormals and nothing here goes near 1e-38. Tensor-core MMA is not governed by these flags.

What was actually verified

• cargo check -p mistralrs-quant (no cuda) — green.
• cargo test -p mistralrs-quant --test cuda_kernel_build_guard — 5/5, including expected_kernel_count_matches_disk.
• EXPECTED_KERNEL_COUNT 41 → 43 (kernel + its cc<8.0 link stub); the dummy-stub count in that file's own note updated 5 → 6.
• nvcc for sm_80 and sm_90 plus cargo check --workspace --features cuda --tests, via the free no-GPU Actions gate on this PR.

Nothing else. In particular: no kernel executed, no output compared against a reference, no timing taken.

🤖 Generated with Claude Code

https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7

[github-actions]

Code Metrics Report

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Language              Files        Lines         Code     Comments       Blanks
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 C Header                  5          305          210           52           43
 CSS                       2         1181         1036           34          111
 CUDA                     78        27867        19381         5568         2918
 Dockerfile                1           39           22            8            9
 JavaScript               16         3546         2676          482          388
 Jinja2                    7          694          656            5           33
 JSON                     74         4600         4597            0            3
 Makefile                  1            6            5            0            1
 Metal Shading Lan|       33        12224         9431         1142         1651
 PowerShell                1          300          227           30           43
 Python                  147        15371        12677          824         1870
 Shell                    42        10258         6849         2734          675
 Plain Text                4         3801            0         2479         1322
 TOML                     33         1498         1294           54          150
 YAML                      3           25           23            2            0
─────────────────────────────────────────────────────────────────────────────────
 HTML                      4         2687         2604           43           40
 |- CSS                    2          543          479           37           27
 |- JavaScript             1         1233         1215           12            6
 (Total)                             4463         4298           92           73
─────────────────────────────────────────────────────────────────────────────────
 Jupyter Notebooks         4          122           83           23           16
 |- Markdown               1           60           30           22            8
 |- Python                 1          122          113            1            8
 (Total)                              304          226           46           32
─────────────────────────────────────────────────────────────────────────────────
 Markdown                204        45330            0        35183        10147
 |- BASH                  72         1655         1203          331          121
 |- C                      3           17           17            0            0
 |- CUDA                   2           84           56           16           12
 |- JSON                  18          708          708            0            0
 |- PowerShell             1            1            1            0            0
 |- Python                23         1008          787          113          108
 |- Rust                  66         2051         1716           77          258
 |- TOML                   6          207          164            0           43
 |- YAML                   5           41           36            5            0
 (Total)                            51102         4688        35725        10689
─────────────────────────────────────────────────────────────────────────────────
 Rust                    678       335044       288629        17172        29243
 |- Markdown             497        30361          471        26183         3707
 (Total)                           365405       289100        43355        32950
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Total                  1337       502989       357396        92632        52961
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

[heydryft]

Route-2 handover: the cuBLASLt block-scaled FP8 path is more reachable than either of us thought

Recorded here so it isn't rediscovered. Three findings, each verified, that change route 2's scope.

1. 🔑 The Rust bindings already have the enums — no toolkit pin bump is needed to compile

I originally reported route 2 as blocked because CUBLASLT_MATMUL_MATRIX_SCALE_BLK128x128_32F is absent from the CUDA 12.4.1 C header, which is what cuda_compile_check.yaml pins. That reasoning was about the C header and does not apply to our actual call path.

We reach cuBLASLt through cudarc, and candle-core 0.9.2 pins cudarc 0.19.4, whose src/cublaslt/sys/mod.rs contains:

CUBLASLT_MATMUL_DESC_A_SCALE_MODE          = 31
CUBLASLT_MATMUL_DESC_B_SCALE_MODE          = 32
CUBLASLT_MATMUL_MATRIX_SCALE_VEC128_32F    = 4
CUBLASLT_MATMUL_MATRIX_SCALE_BLK128x128_32F = 5

cudarc 0.19.4 has no bindgen in its build.rs — the bindings are static Rust constants, not generated from the local toolkit. So the descriptor wiring compiles today, at the 12.4.1 pin, through the existing gate. The ≥12.9 requirement is a runtime cuBLAS-library requirement on the serving box, not a build-time one.

This also independently corroborates the peer's header grep from a real 13.0 box — 31, 32, and BLK128x128_32F = 5 match exactly, from a completely different source. Those magic numbers are no longer unverified.

Also present in the 12.4-era cudarc 0.17.8, so this is not new.

2. 🔑 The activation quantiser already exists — it does not need writing, and EXPECTED_KERNEL_COUNT does not move for route 2

quant_fp8_blockwise_kernel (kernels/blockwise_fp8/blockwise_fp8.cu) is fully general in weight_block_size_y/_x. Called as fp8_blockwise_quantize(x_2d, vec![1, 128]) it produces scales of shape [M, ceil(K/128)], M-major — which is exactly the M x L layout VEC128_32F requires. The Rust op fp8_blockwise_quantize is already pub and already exported.

3. ⚠️ But it is not fit for the decode hot path as written — this is the real work item

With block_size_y = 1 the launcher still uses dim3 blockDim(32, 32) = 1024 threads, and the tile loop is local_y = threadIdx.y; local_y < 1; local_y += 32. So:

• 32 of 1024 threads do work; 31/32 of every block idles.
• The absmax reduction is atomicMaxFloat onto one shared address — 128 serialised atomics per 128-element block.
• At M=256, K=7168 that is 14,336 blocks x 1024 threads to quantise 1.8 M elements.

It is correct, so it unblocks a first end-to-end A/B immediately. It is not something to leave in the steady-state decode path. The fix is conventional — one warp per 128-block, __shfl_down_sync reduction, vectorised loads — and it is a genuine new kernel (or a block_size_y == 1 specialisation), so that is what would move EXPECTED_KERNEL_COUNT, not the wiring.

Remaining scope for route 2, with the unknowns removed

1. Warp-per-block specialisation of the quantiser for block_size_y == 1 (new kernel + stub, count 43 -> 45).
2. cublaslt/api.rs: lift the a_scale/b_scale f32-scalar requirement (:70-76), pass device scale tensors, set A_SCALE_MODE = VEC128_32F and B_SCALE_MODE = BLK128x128_32F. weight_scale_inv is already the correct B layout.
3. Wire into BlockwiseFP8Linear::forward, dropping dequantize_w().
4. Constraints to honour: TN on Hopper, M and N multiples of 4, scale base addresses 16 B aligned, BLK128x128 scales K-major with column stride a multiple of 4, D not FP8.
5. Runtime gate: the box needs cuBLAS >= 12.9. Below that cublasLtMatmulDescSetAttribute will reject attribute 31/32 — loudly, per the docs ("If the corresponding scale is set for a non-narrow precisions tensor, cuBLAS will return an error"), so it fails visibly rather than silently. Needs a runtime version check and a fallback, since our boxes have run 12.4.

How this PR relates

It does not compete with route 2 — it is the control arm for it, and the only arm that works below cuBLAS 12.9. Route 2's ceiling is better (FP8 tensor cores, 1.10 ms derived, vs 2.19 ms here) because it quantises activations to FP8; this kernel deliberately does not touch activation numerics. Both delete the 12.7 GB/step dequantize_w().

Whoever lands route 2 should keep ARC_NO_FP8_WMMA and this kernel as the fallback, and run the three-arm sweep documented in ops.rs, not a two-arm one.

None of this has run. Every figure above is a derivation or a code reading, not a measurement.