[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:
co-authored by
Jee Jee Li
andyluo7
OpenAI Codex
zjy0516
Claude Opus 5
parent
e6821ceac9
commit
e77daef89e
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user