[Bugfix] Handle E8M0 block scales in CUTLASS and Triton FP8 linear kernels

[waynehacking8]

Purpose

Fixes #47818.

DeepSeek-V4-style FP8 block checkpoints declare "scale_fmt": "ue8m0", so Fp8LinearMethod loads weight_scale_inv as exponent-only float8_e8m0fnu (fp8.py is_scale_e8m0). The DeepGEMM kernel handles that dtype in its own weight post-processing, but on any setup where DeepGEMM is unavailable or disabled (which today includes SM120, see #47436) the layer routes to the CUTLASS or Triton block kernels, and both crash on the first block-scaled linear during profile_run:

• CUTLASS: dispatch_scaled_mm requires fp32 scales (STD_TORCH_CHECK(b_scales.scalar_type() == Float) in scaled_mm_helper.hpp) — the exact error in [Bug]: Cutlass C3X dispatch_scaled_mm crashes on SM120 Blackwell (RTX PRO 6000) with FP8 block-scaled model (DeepSeek-V4-Flash) #47818. The issue's guess of a missing SM120 dispatch entry is not what happens; the SM120 blockwise kernel exists and works, the dispatch dies on the scale dtype check before reaching it.
• Triton: the kernel launch cannot bind an E8M0 tensor (KeyError: 'float8_e8m0fnu') — the reporter's --linear-backend triton traceback.

E8M0 holds pure powers of two, so upcasting to fp32 is lossless. This PR:

1. CutlassFp8BlockScaledMMKernel.process_weights_after_loading: upcast E8M0 weight scales to fp32 once at load time, reusing the existing _upcast_e8m0_to_fp32 helper (same treatment deepgemm_post_process_fp8_weight_block already applies on the DeepGEMM path).
2. w8a8_triton_block_scaled_mm: the E8M0→fp32 upcast already existed but was gated to ROCm/XPU; the Triton limitation is platform-independent, so drop the gate.
3. CutlassFp8BlockScaledMMKernel.can_implement: on SM12x, decline layers whose weight output dim is not a multiple of 128 so selection falls through to Triton. I swept the SM120 blockwise kernel on an RTX PRO 6000 across M ∈ {1..512} × N ∈ {512, 576, 1536, 2112, 7168} × K ∈ {1024, 7168} with fp32 scales: every N % 128 == 0 case passes (max rel err ~2-3e-3 vs fp32 dequant reference, bf16 rounding level), every N % 128 != 0 case fails CUTLASS can_implement ("Invalid status" from cutlass_gemm_caller.cuh) on all three tile configs. That kernel-level gap also explains why test_w8a8_block_fp8_cutlass_matmul (N=576, direct op call) fails on SM120 today — pre-existing, not addressed here, filed as [Bug]: SM120 CUTLASS blockwise FP8 GEMM rejects N % 128 != 0 (Invalid status; kv_a_proj N=576 in DeepSeek shapes) #47990.

Why this is not duplicating an existing PR

No open PR references #47818. #41834 (DSV4 SM12x enablement branch) bundles an equivalent E8M0 upcast among its ~116 files, but it routes SM12x away from CUTLASS entirely; my SM120 sweep shows the blockwise kernel is correct for 128-aligned N, so this PR keeps CUTLASS for those shapes and only falls through for the shapes the kernel genuinely rejects. This is a standalone fix for the crash on current main, independent of the model-enablement branch.

Test Plan

Unit (RTX PRO 6000, SM120, torch 2.11.0+cu130, precompiled nightly kernels):

python -m pytest tests/kernels/quantization/test_block_fp8.py -k "e8m0 or sm12x"

• test_w8a8_block_fp8_matmul_e8m0_scales — Triton path with E8M0 As/Bs must bit-match the fp32-scale run.
• test_cutlass_block_fp8_e8m0_weight_scale_upcast — kernel-class level: after process_weights_after_loading the scale param is fp32 with exactly round-tripped values, and apply_weights output bit-matches an fp32-scale layer.
• test_cutlass_block_fp8_sm12x_declines_unaligned_n — can_implement falls through for N=576 on family-120, accepts elsewhere (platform mocked, runs on any host).

Runtime A/B on a 2-layer dummy-weight DeepSeek-V4-Flash config (real attention dims and quantization_config, load_format=dummy, VLLM_USE_DEEP_GEMM=0, single RTX PRO 6000):

• stock main (f7fc0ca): Selected CutlassFp8BlockScaledMMKernel for Fp8LinearMethod then RuntimeError: dispatch_scaled_mm, .../scaled_mm_helper.hpp:17 — reproduces the issue byte-for-byte.
• this branch: all block-FP8 dense linears process and execute; the forward proceeds until DSV4's o_proj fp8_einsum, which calls DeepGEMM directly and is a separate SM120 gap ([Bug]: Block-scaled FP8 (compressed-tensors W8A8) crashes on load on SM120 Blackwell (RTX PRO 6000), v0.24.0 — DeepGEMM "Unknown SF transformation" assertion #47436 family / [New Model][Nvidia] Add SM12x support for DeepSeek V4 Flash with essential fixes #41834 scope). This PR fixes the linear-kernel crash reported in [Bug]: Cutlass C3X dispatch_scaled_mm crashes on SM120 Blackwell (RTX PRO 6000) with FP8 block-scaled model (DeepSeek-V4-Flash) #47818; it does not claim full DSV4-Flash serving on SM120.

Test Result

• 3 new tests pass on the fix; on stock main the two GPU tests fail with exactly the two crash signatures from the issue (CUTLASS scaled_mm_helper.hpp check, Triton KeyError: 'float8_e8m0fnu').
• tests/kernels/quantization/test_block_fp8.py: identical pass/fail set before and after the change on this GPU (the 101 pre-existing failures here are all test_w8a8_block_fp8_deep_gemm_matmul — DeepGEMM on SM120 — plus the N=576 direct-op case above; none are introduced or resolved by this PR).

Accuracy note: outputs for previously-working configs are unchanged by construction (the upcast is exact and only engages for E8M0 checkpoints, which currently cannot run at all on these kernels). A real-weight eval of an ue8m0 checkpoint isn't possible on this box (DeepSeek-V4-Flash is 149 GB vs 96 GB VRAM); the kernel-level tests assert bit-identity against fp32-scale references instead.

---

AI assistance was used for this PR (Claude Code); I reviewed every line and ran the tests and A/B above on my own SM120 hardware.