[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>
This commit is contained in:
Yongye Zhu
2026-09-11 02:11:19 -07:00
committed by GitHub
co-authored by Jee Jee Li andyluo7 OpenAI Codex zjy0516 Claude Opus 5
parent e6821ceac9
commit e77daef89e
49 changed files with 3996 additions and 138 deletions
@@ -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)
@@ -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,