Module w8a8_prefill

Module w8a8_prefill 

Source
Expand description

W8A8 block-scaled prefill branch (#917/#928). A CHILD module, not a sibling: it adds impl DenseFfnLayer methods that read this layer’s private kernel handles, and this file is already at the CI size cap. W8A8 block-scaled dense-FFN PREFILL — the gate/up/down GEMM arm that the native-FP8 dispatch in dense_ffn.rs reaches ahead of its W8A16 branches.

§Reachability — read this before measuring anything here

This whole module is opt-in and unreached by default. It runs only when the dense FFN holds block-scaled FP8 weights, and qwen35_dense.rs::load_layers installs those only when

dense_fp8_enabled() && tp_world_size == 1
    && variant == Nvfp4Variant::Fp8Dequanted
    && proj_is_native_fp8(gate_proj)

and dense_fp8_enabled() is ATLAS_DENSE_FP8 == "1". git grep ATLAS_DENSE_FP8 returns ONE hit — its own reader. No CI job, no BENCH.toml entry and no gate sets it, so no certification record on any branch has ever exercised this code. Without it self.fp8_weights is None, forward_prefill_inner never reaches the selection below, and NEITHER of the two route log lines is emitted. Three served A/B attempts were spent discovering that, each reading as “the lever did not arm”.

It is off by default because on GB10 native dense FP8 LOSES to the NVFP4 autoquant fallback, and not narrowly. Same box, same binary, same byte-identical 949-token prompt, Qwen3.6-27B-FP8, TTFT median of 5 (spark-256a, 2026-09-11):

default (NVFP4)                       1437.2 ms   1.00x
ATLAS_DENSE_FP8=1, W8A16 (capped)     2555.8 ms   1.78x slower
ATLAS_DENSE_FP8=1, W8A8  (no cap)     3343.3 ms   2.33x slower

So the ceiling below is worth 23.4% within the dense-FP8 path, and the dense-FP8 path is still the slower choice on this arch. GB10’s FP8 W8A16 kernels are simply less tuned than its NVFP4 W4A16 ones (unfused per-projection GEMV, non-transposed prefill GEMM); closing that is kernel work, not loader wiring. On H100 the trade is the other way round, which is why the ceiling is per-arch data in kernels/<hw>/HARDWARE.toml and not a constant here.

WHY (#917 / #928). On a native-FP8 checkpoint the dense FFN’s prefill GEMMs ran w8a16_gemm_pipelined: BF16 activations against E4M3 weights, so the MMA is the BF16 tensor-core path and the FP8 bytes are pure memory savings. Measured on H100 (2026-09-11, 1193-token prompt): TTFT 1075 ms against vLLM’s 287 ms for the same model and prompt, with the pipelined kernel turning ~12 TFLOP/s on the FFN shapes. The attention Q/K/V/O projections had already moved to the W8A8 block-scaled path (paged_qkv.rs / paged_oproj.rs), and the MoE shared expert with them (moe/forward_prefill_fp8.rs) — the dense FFN was the one large prefill consumer still on W8A16, and on a dense model it is most of the FLOPs.

This gives it the same arithmetic vLLM uses: per-token 1x128 FP32 activation scales (per_token_group_quant_fp8) multiplied against the checkpoint’s 128x128 FP32 weight scales in an FP32 epilogue, with the product accumulated by mma.sync.m16n8k32.e4m3 — native on sm_90a (H100) and sm_121. Two GEMM implementations sit behind one selector:

  • cuBLASLt fp8_gemm_act_weight_t_blkscaled (weight as A with BLK128x128 scales, activation as B with VEC128 scales — the DeepSeek block-FP8 scheme), when ATLAS_CUBLAS_GEMM=1. The Hopper fast path. Its VEC128 scales go through fp8_act_scale_to_kmajor first: cuBLASLt documents that operand’s scales with the TOKEN index contiguous, which is the transpose of what the quantizer writes (see spark_runtime::cublaslt::scale_layout).
  • ops::fp8_gemm_t_blockscaled, the in-tree kernel, otherwise.

Both consume the SAME quantized activation and the same FP32 epilogue, so they are expected to agree to a BF16 ULP or two; the microtest (examples/native_fp8_ffn_w8a8_microtest.rs) pins that.

ACCURACY. W8A8 is lossier than W8A16 by construction — the activation is quantized to E4M3 per 128-element group instead of kept in BF16. That is vLLM’s dynamic W8A8 arithmetic and a deliberate precision trade, not a bug: the microtest gates it at cosine >= 0.999 / relative RMS <= 2% against the W8A16 reference, and the serve logs the selected path once at INFO so which arithmetic ran is visible in any TTFT report. ATLAS_FFN_W8A16_ONLY=1 restores the old path byte-for-byte.

Functions§

ffn_w8a16_only
ATLAS_FFN_W8A16_ONLY kill switch: PRESENCE (any value, including empty) keeps the dense-FFN prefill on today’s W8A16 kernels. Presence rather than =1 because this is an escape hatch an operator reaches for while a serve is misbehaving, and ATLAS_FFN_W8A16_ONLY=0 meaning “on” is a trap.