Module w8a8_prefill

Module w8a8_prefill 

Source
Expand description

2026-09-25: The W8A8 block-scaled FP8 prefill arm. 2026-09-25: The W8A8 block-scaled arm of the native-FP8 dense FFN: which projections take it, and the GEMM it runs.

Owner: model-layers (dense FFN). Invariants:

  • w8a8_prefill_selected holds only for 4 < m <= max_m, Fp8BlockScaled weights, k and n multiples of 128, the quantizer and GEMM handles loaded, METRALE_FFN_W8A16_ONLY unset and dispatch.fp8_blockscaled_prefill on.

The arm is reachable only on a layer that holds native-FP8 dense-FFN weights. The one loader that installs them, qwen35_dense.rs, does so only when ffn_fp8_arm_selected holds: METRALE_DENSE_FP8 is exactly 1, config.tp_world_size.max(1) == 1, the variant is Nvfp4Variant::Fp8Dequanted, and the layer’s mlp.gate_proj is block-scaled FP8 on disk (proj_is_native_fp8). Without it neither route log in this file is emitted.

The activation is quantized to E4M3 with one FP32 scale per token per 128-wide K group (per_token_group_quant_fp8). w8a8_gemm then runs cuBLASLt fp8_gemm_act_weight_t_blkscaled when dispatch.cublas.ffn holds (METRALE_CUBLAS_GEMM lists ffn, or is 1, true or all) and its other clauses pass, and the in-tree ops::fp8_gemm_t_blockscaled otherwise. examples/native_fp8_ffn_w8a8_microtest.rs grades W8A8 against the W8A16 reference at cosine >= 0.999 and relative RMS <= 2%.

The upper M bound is per target: kernels/<hw>/HARDWARE.toml [defaults] w8a8_prefill_max_m_widening / _narrowing. gb10 declares 64 and 384 beside its measurement; every other target has no cap (u32::MAX).

Functions§

ffn_w8a16_only
2026-09-25: METRALE_FFN_W8A16_ONLY, by presence (any value, including empty or 0): no dense-FFN projection takes the W8A8 arm. Read once per process (OnceLock).