From 3c2449065371c3d9749b283b79a02dbfcd227dd2 Mon Sep 17 00:00:00 2001 From: Andrewxu313 Date: Tue, 9 Jun 2026 12:19:49 +0000 Subject: [PATCH] fix: use disable_ue8m0_cast for FP8 GEMM on Blackwell (sm_100) w8a16_gemm had disable_ue8m0_cast dropped, so deep_gemm.fp8_gemm_nt defaulted to the UE8M0/MXFP8 scale path on sm_100, which needs power-of-2 scales we don't supply -> all-NaN (GLM-5.1-FP8 emitted '!!!!' on B200). Pass disable_ue8m0_cast=True (DeepGEMM's regular blockwise-FP8 path; matches w8a8_deepgemm; no-op on Hopper). Also fix w8a16_gemm_dequant's 3D reshape to use the explicit N dim instead of -1 (ambiguous for empty/0-token sub-batches). Verified: coherent GLM-5.1-FP8 generation on 8xB200. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Andrewxu313 --- batchgen/attention/mla/fa3_backend.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/batchgen/attention/mla/fa3_backend.py b/batchgen/attention/mla/fa3_backend.py index f1ec03420..2c4f9cfc3 100644 --- a/batchgen/attention/mla/fa3_backend.py +++ b/batchgen/attention/mla/fa3_backend.py @@ -618,11 +618,12 @@ def w8a16_gemm( # act_quant now routes bf16 inputs to the validated batchgen_kernels quant; # non-bf16 / empty sub-batches fall back to the legacy in-file Triton path. x_fp8 = act_quant(x) - # disable_ue8m0_cast removed — on Hopper (SM90) the flag is a no-op - # (layout.hpp:22 early-exits for arch_major==9 regardless), and omitting - # lets DeepGEMM's default handling apply (same as SGLang). This also - # ensures Blackwell upgrade path uses UE8M0 natively when appropriate. - deep_gemm.fp8_gemm_nt(x_fp8, y_fp8, out) + # disable_ue8m0_cast=True: use DeepGEMM's regular blockwise-FP8 scale path. + # The default (False) selects the UE8M0/MXFP8 path on Blackwell (sm_100), which + # expects UE8M0 (power-of-2) scales that act_quant / checkpoint scale_inv do NOT + # provide -> all-NaN output (verified on B200). True is a no-op on Hopper. + # Matches w8a8_deepgemm.py. See batchgen_design/blackwell/numerical_debug_plan.md. + deep_gemm.fp8_gemm_nt(x_fp8, y_fp8, out, disable_ue8m0_cast=True) if activation_bf16.dim() == 3: out = out.view(n_group, l, n) else: @@ -654,7 +655,9 @@ def w8a16_gemm_dequant( x = activation_bf16.view(-1, activation_bf16.size(-1)) out = torch.mm(x, weight_bf16.T) if is_3d: - out = out.view(n_group, l, -1) + # explicit N (not -1): -1 is ambiguous for 0-element / empty sub-batches + # (ranks with 0 tokens), which the deep_gemm path avoided via explicit n. + out = out.view(n_group, l, out.size(-1)) return out