[Model] Support DeepSeek-V4.1-Flash (#56214)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Signed-off-by: andyluo7 <andy.luo@amd.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: andyluo7 <andy.luo@amd.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

vllm/model_executor/layers/quantization/__init__.py

@@ -115,7 +115,20 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
     # lazy import to avoid triggering `torch.compile` too early
     from vllm.config.quantization import _ONLINE_SHORTHANDS
     from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
-    from vllm.models.deepseek_v4 import DeepseekV4FP8Config
+
+    if current_platform.is_cuda() or current_platform.is_rocm():
+        # The v4.1 class is the v4 one extended to accept model_type
+        # "deepseek_v41" and its 32x32 MXFP8 linear layout. V4.1 has
+        # platform-specific implementations for both CUDA and ROCm.
+        from vllm.models.deepseek_v4_1 import (
+            DeepseekV4FP8Config as DeepseekV41FP8Config,
+        )
+
+        deepseek_config: type[QuantizationConfig] = DeepseekV41FP8Config
+    else:
+        from vllm.models.deepseek_v4 import DeepseekV4FP8Config
+
+        deepseek_config = DeepseekV4FP8Config
 
     from .auto_awq import AutoAWQConfig
     from .auto_gptq import AutoGPTQConfig
@@ -161,7 +174,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
         "inc": INCConfig,
         "mxfp4": Mxfp4Config,
         "gpt_oss_mxfp4": GptOssMxfp4Config,
-        "deepseek_v4_fp8": DeepseekV4FP8Config,
+        "deepseek_v4_fp8": deepseek_config,
         "humming": HummingConfig,
         "online": OnlineQuantizationConfig,
         # MiniMax-style checkpoints tag `quant_method: "mxfp8"`; load with the

vllm/model_executor/layers/quantization/utils/fp8_utils.py

@@ -937,14 +937,13 @@ def w8a8_triton_block_scaled_mm(
     assert len(block_size) == 2
     block_n, block_k = block_size[0], block_size[1]
 
-    # Triton cannot currently bind E8M0 scale tensors directly. On ROCm,
-    # DeepSeek-V4 checkpoints store block scales in exponent-only E8M0 format,
-    # so decode them to fp32 before launching the kernel.
-    if current_platform.is_rocm() or current_platform.is_xpu():
-        if As.dtype == torch.float8_e8m0fnu:
-            As = _upcast_e8m0_to_fp32(As).contiguous()
-        if Bs.dtype == torch.float8_e8m0fnu:
-            Bs = _upcast_e8m0_to_fp32(Bs).contiguous()
+    # Triton cannot currently bind E8M0 scale tensors directly. Checkpoints
+    # with exponent-only UE8M0 scales (e.g. DeepSeek-V4) are decoded to fp32
+    # before launching the kernel.
+    if As.dtype == torch.float8_e8m0fnu:
+        As = _upcast_e8m0_to_fp32(As).contiguous()
+    if Bs.dtype == torch.float8_e8m0fnu:
+        Bs = _upcast_e8m0_to_fp32(Bs).contiguous()
 
     assert A.shape[-1] == B.shape[-1]
     assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
@@ -1089,8 +1088,6 @@ def deepgemm_post_process_weight_scale_block(
     is_sfa: bool = False,
 ) -> torch.Tensor:
     if ws.dtype in (torch.float8_e8m0fnu, torch.uint8):
-        # Scales already in E8M0 from checkpoint; upcast to fp32 and let
-        # DeepGEMM pack the layout expected by the target architecture.
         ws = _upcast_e8m0_to_fp32(ws)
     else:
         assert ws.dtype == torch.float32, (
@@ -1121,12 +1118,7 @@ def deepgemm_post_process_fp8_weight_block(
         f"to be torch.float8_e4m3fn, got {wq.dtype} instead."
     )
 
-    if ws.dtype in (torch.float8_e8m0fnu, torch.uint8):
-        # Scales already in E8M0 from checkpoint (float8_e8m0fnu, or raw E8M0
-        # bits as uint8 for MXFP8) - upcast to fp32 and skip requantization
-        # (weights already have power-of-two scales).
-        ws = _upcast_e8m0_to_fp32(ws)
-    else:
+    if ws.dtype not in (torch.float8_e8m0fnu, torch.uint8):
         assert ws.dtype == torch.float32, (
             f"Expected tensor scales dtype to be torch.float32 or "
             f"torch.float8_e8m0fnu or torch.uint8, got {ws.dtype} instead"
@@ -1137,7 +1129,7 @@ def deepgemm_post_process_fp8_weight_block(
     if is_bmm:
         # Reshape 2D weight/scale to 3D for grouped BMM (einsum):
         # wq: (g*r, d) -> (g, r, d)
-        # ws: (g*r/128, d/128) -> (g, r/128, d/128)
+        # Scale dimensions follow quant_block_shape.
         g = bmm_batch_size
         assert wq.ndim == 2 and ws.ndim == 2
         d = wq.size(1)

vllm/model_executor/layers/quantization/utils/mxfp8_utils.py

@@ -126,8 +126,14 @@ def _mxfp8_quant_triton_kernel():
         amax = tl.maximum(tl.max(tl.abs(x), axis=1), TINY)  # [BLOCK_M]
         sb = tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0
         sb = tl.minimum(tl.maximum(sb, 0.0), 254.0)
-        descale = tl.exp2(sb - 127.0)
-        xq = (x / descale[:, None]).to(xq_ptr.dtype.element_ty)
+        # Scale by the reciprocal rather than dividing by ``exp2(sb - 127)``:
+        # sb == 0 (an all-zero / denormal block) makes that divisor 2**-127,
+        # which is subnormal in fp32 and flushes to zero on CDNA, turning the
+        # block's zeros into 0/0 == NaN. The reciprocal is normal for every
+        # reachable sb, and both forms are exact powers of two, so nothing
+        # else changes.
+        rescale = tl.exp2(127.0 - sb)
+        xq = (x * rescale[:, None]).to(xq_ptr.dtype.element_ty)
         tl.store(
             xq_ptr + offs_m[:, None] * sqm + offs_k[None, :] * sqk,
             xq,