[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
co-authored by
Jee Jee Li
andyluo7
OpenAI Codex
zjy0516
Claude Opus 5
Jee Jee Li
andyluo7
OpenAI Codex
zjy0516
Claude Opus 5
parent
e6821ceac9
commit
e77daef89e
@@ -37,6 +37,7 @@ from vllm.model_executor.layers.quantization.fp8 import (
|
||||
Fp8MoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
|
||||
from vllm.model_executor.layers.quantization.modelopt import ModelOptLinearMethod
|
||||
from vllm.model_executor.layers.quantization.online.fp8 import (
|
||||
Fp8PerTensorOnlineLinearMethod,
|
||||
)
|
||||
@@ -63,6 +64,383 @@ MODELS = [
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="DSV4.1 requires CUDA")
|
||||
@pytest.mark.parametrize("tp_rank", [0, 1])
|
||||
@pytest.mark.parametrize("scale_dtype", [torch.uint8, torch.float8_e8m0fnu])
|
||||
def test_deepseek_v41_mxfp8_scale_loading(
|
||||
dist_init, default_vllm_config, monkeypatch, tp_rank, scale_dtype
|
||||
):
|
||||
"""Expand linear scales before TP slicing, including padded shared experts."""
|
||||
from vllm.model_executor.layers import linear as linear_module
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
MergedColumnParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from vllm.models.deepseek_v4_1 import quant_config as quant_module
|
||||
from vllm.models.deepseek_v4_1.nvidia import model as model_module
|
||||
|
||||
default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16)
|
||||
quant_config = quant_module.DeepseekV4FP8Config(
|
||||
is_checkpoint_fp8_serialized=True, weight_block_size=[32, 32]
|
||||
)
|
||||
quant_config._resolved_expert_dtype = "fp4"
|
||||
for module in (model_module, linear_module):
|
||||
monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)
|
||||
monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: tp_rank)
|
||||
tp_args = dict(quant_config=quant_config, bias=False)
|
||||
block = torch.nn.Module()
|
||||
block.attn = torch.nn.Module()
|
||||
block.attn.wq_b = ColumnParallelLinear(128, 128, **tp_args)
|
||||
block.attn.fused_wqa_wkv = MergedColumnParallelLinear(128, [128, 64], **tp_args)
|
||||
block.ffn = torch.nn.Module()
|
||||
block.ffn.shared_experts = torch.nn.Module()
|
||||
block.ffn.shared_experts.gate_up_proj = MergedColumnParallelLinear(
|
||||
128, [128, 128], **tp_args
|
||||
)
|
||||
block.ffn.shared_experts.down_proj = RowParallelLinear(128, 128, **tp_args)
|
||||
block.ffn.experts = torch.nn.Module()
|
||||
block.ffn.experts.register_parameter(
|
||||
"weight_scale", torch.nn.Parameter(torch.empty(3, 4, dtype=torch.uint8), False)
|
||||
)
|
||||
block.engram = torch.nn.Module()
|
||||
block.engram.register_parameter(
|
||||
"weight_scale_inv",
|
||||
torch.nn.Parameter(torch.empty(3, 4, dtype=scale_dtype), False),
|
||||
)
|
||||
|
||||
def load_expert_scale(param, weight, *args, **kwargs):
|
||||
param.data.copy_(weight)
|
||||
return True
|
||||
|
||||
block.ffn.experts.weight_scale.weight_loader = load_expert_scale
|
||||
model = model_module.DeepseekV4Model.__new__(model_module.DeepseekV4Model)
|
||||
torch.nn.Module.__init__(model)
|
||||
model.layers = torch.nn.ModuleList([block])
|
||||
model.config = SimpleNamespace(num_attention_heads=2)
|
||||
model.quant_config = quant_config
|
||||
model.use_sequence_parallel = False
|
||||
model.get_expert_mapping = lambda: [
|
||||
("experts.weight_scale", "experts.0.w1.weight_scale", 0, "w1")
|
||||
]
|
||||
|
||||
checkpoints = [
|
||||
("attn.wq_b", "attn.wq_b", 128, 128, 0, None),
|
||||
("attn.wq_a", "attn.fused_wqa_wkv", 128, 128, 0, slice(0, 64)),
|
||||
("attn.wkv", "attn.fused_wqa_wkv", 64, 128, 0, slice(64, 96)),
|
||||
(
|
||||
"ffn.shared_experts.w1",
|
||||
"ffn.shared_experts.gate_up_proj",
|
||||
96,
|
||||
128,
|
||||
0,
|
||||
slice(0, 64),
|
||||
),
|
||||
(
|
||||
"ffn.shared_experts.w3",
|
||||
"ffn.shared_experts.gate_up_proj",
|
||||
96,
|
||||
128,
|
||||
0,
|
||||
slice(64, 128),
|
||||
),
|
||||
(
|
||||
"ffn.shared_experts.down_proj",
|
||||
"ffn.shared_experts.down_proj",
|
||||
128,
|
||||
96,
|
||||
1,
|
||||
None,
|
||||
),
|
||||
]
|
||||
weights = []
|
||||
expected = []
|
||||
for source, target, n, k, axis, shard in checkpoints:
|
||||
weight = torch.randint(-4, 5, (n, k)).to(torch.float8_e4m3fn)
|
||||
scale_bytes = torch.randint(124, 131, (n // 32, k // 32), dtype=torch.uint8)
|
||||
weights.extend(
|
||||
[
|
||||
(f"layers.0.{source}.weight", weight),
|
||||
(f"layers.0.{source}.scale", scale_bytes.view(scale_dtype)),
|
||||
]
|
||||
)
|
||||
dequant = (
|
||||
weight.float().reshape(n // 32, 32, k // 32, 32)
|
||||
* torch.exp2(scale_bytes.float() - 127)[:, None, :, None]
|
||||
).reshape(n, k)
|
||||
if "shared_experts" in source:
|
||||
padded = torch.zeros(128, 128)
|
||||
padded[:n, :k] = dequant
|
||||
dequant = padded
|
||||
expected.append((target, shard, dequant.chunk(2, dim=axis)[tp_rank]))
|
||||
|
||||
expert_scale = torch.full((3, 4), 125, dtype=torch.uint8)
|
||||
weights.append(
|
||||
("layers.0.ffn.experts.0.w1.weight_scale", expert_scale.view(scale_dtype))
|
||||
)
|
||||
weights.append(("layers.0.engram.weight_scale_inv", expert_scale.view(scale_dtype)))
|
||||
mapper = model_module._make_deepseek_v4_weights_mapper("fp4", "weight_scale")
|
||||
loaded = model.load_weights(
|
||||
(name.removeprefix("model."), weight) for name, weight in mapper.apply(weights)
|
||||
)
|
||||
assert "layers.0.ffn.experts.weight_scale" in loaded
|
||||
assert torch.equal(block.ffn.experts.weight_scale, expert_scale)
|
||||
assert torch.equal(block.engram.weight_scale_inv.view(torch.uint8), expert_scale)
|
||||
for target, shard, reference in expected:
|
||||
linear = block.get_submodule(target)
|
||||
assert isinstance(linear.quant_method, ModelOptLinearMethod)
|
||||
assert f"layers.0.{target}.weight_scale" in loaded
|
||||
weight = linear.weight if shard is None else linear.weight[shard]
|
||||
scale = linear.weight_scale if shard is None else linear.weight_scale[shard]
|
||||
actual = (
|
||||
weight.float().unflatten(-1, (-1, 32))
|
||||
* torch.exp2(scale.float() - 127).unsqueeze(-1)
|
||||
).flatten(-2)
|
||||
torch.testing.assert_close(actual, reference, rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="DSV4.1 requires CUDA")
|
||||
@pytest.mark.parametrize(
|
||||
"weight_block_size,expert_dtype,scale_name",
|
||||
[
|
||||
([32, 32], "fp4", "weight_scale"),
|
||||
([128, 128], "fp4", "weight_scale_inv"),
|
||||
([32, 32], "fp8", "weight_scale_inv"),
|
||||
],
|
||||
)
|
||||
def test_deepseek_v41_vl_mapper_routes_linear_scales(
|
||||
weight_block_size, expert_dtype, scale_name
|
||||
):
|
||||
"""The VL wrapper must map ``.scale`` keys to the parameter the linear
|
||||
quant method registers, as the text model does. A hardcoded
|
||||
``weight_scale_inv`` raised KeyError for native MXFP8 checkpoints."""
|
||||
from vllm.models.deepseek_v4_1.nvidia import model as model_module
|
||||
from vllm.models.deepseek_v4_1.nvidia import vl_model as vl_module
|
||||
|
||||
vllm_config = SimpleNamespace(
|
||||
quant_config=SimpleNamespace(weight_block_size=weight_block_size)
|
||||
)
|
||||
resolved = model_module._linear_scale_param_name(vllm_config, expert_dtype)
|
||||
assert resolved == scale_name
|
||||
|
||||
mapper = vl_module._make_deepseek_v4_vl_weights_mapper(expert_dtype, resolved)
|
||||
weight = torch.empty(0)
|
||||
mapped = [
|
||||
name
|
||||
for name, _ in mapper.apply(
|
||||
[("layers.0.attn.wq_a.scale", weight), ("layers.0.attn.wkv.weight", weight)]
|
||||
)
|
||||
]
|
||||
assert mapped == [
|
||||
f"language_model.model.layers.0.attn.wq_a.{scale_name}",
|
||||
"language_model.model.layers.0.attn.wkv.weight",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="DeepGEMM requires CUDA")
|
||||
@pytest.mark.parametrize("scale_dtype", [torch.uint8, torch.float8_e8m0fnu])
|
||||
@pytest.mark.parametrize(
|
||||
"weight_shape,is_bmm",
|
||||
[
|
||||
((129, 160), False),
|
||||
((387, 160), True),
|
||||
((3, 129, 160), False),
|
||||
((3, 128, 512), False),
|
||||
],
|
||||
)
|
||||
def test_deepgemm_mxfp8_preserves_weight_and_scale_values(
|
||||
scale_dtype, weight_shape, is_bmm
|
||||
):
|
||||
"""DeepGEMM layout conversion preserves native weights and E8M0 scales."""
|
||||
from vllm.model_executor.layers.quantization.utils import fp8_utils
|
||||
from vllm.utils.deep_gemm import is_deep_gemm_supported
|
||||
|
||||
if not is_deep_gemm_supported() or not current_platform.is_device_capability_family(
|
||||
100
|
||||
):
|
||||
pytest.skip("DeepGEMM MXFP8 requires Blackwell")
|
||||
|
||||
weight = torch.randn(weight_shape, device="cuda").to(torch.float8_e4m3fn)
|
||||
scales = (
|
||||
torch.randint(
|
||||
1,
|
||||
255,
|
||||
(weight.numel() // 32 + 1,),
|
||||
device="cuda",
|
||||
dtype=torch.uint8,
|
||||
)[1:]
|
||||
.view(*weight_shape[:-1], weight_shape[-1] // 32)
|
||||
.view(scale_dtype)
|
||||
)
|
||||
original_weight = weight.view(torch.uint8).clone()
|
||||
original_scales = scales.view(torch.uint8).clone()
|
||||
processed_weight, packed = fp8_utils.deepgemm_post_process_fp8_weight_block(
|
||||
weight,
|
||||
scales,
|
||||
quant_block_shape=(1, 32),
|
||||
use_e8m0=True,
|
||||
is_bmm=is_bmm,
|
||||
bmm_batch_size=3 if is_bmm else 0,
|
||||
)
|
||||
assert packed.dtype == torch.int32
|
||||
unpacked = torch.stack(
|
||||
[(packed >> (8 * i)) & 0xFF for i in range(4)], dim=-1
|
||||
).flatten(-2)
|
||||
unpacked = unpacked[..., : scales.shape[-1]].to(torch.uint8)
|
||||
torch.testing.assert_close(
|
||||
unpacked.reshape_as(original_scales), original_scales, rtol=0, atol=0
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
processed_weight.view(torch.uint8).reshape_as(original_weight),
|
||||
original_weight,
|
||||
rtol=0,
|
||||
atol=0,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="DeepGEMM requires CUDA")
|
||||
@pytest.mark.parametrize("prequantized", [False, True])
|
||||
@pytest.mark.parametrize("config_source", ["deepseek", "mxfp8"])
|
||||
@pytest.mark.parametrize("num_tokens", [1, 7, 128])
|
||||
def test_mxfp8_bmm_loads_and_projects_grouped_weights(
|
||||
dist_init, default_vllm_config, prequantized, config_source, num_tokens
|
||||
):
|
||||
"""BMM metadata set after construction selects grouped weight processing."""
|
||||
from vllm.model_executor.kernels.linear.mxfp8.deep_gemm import (
|
||||
DeepGemmMxfp8BmmLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.linear import ColumnParallelLinear
|
||||
from vllm.model_executor.layers.quantization import get_quantization_config
|
||||
from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
|
||||
fused_inv_rope_fp8_quant,
|
||||
)
|
||||
from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
|
||||
compute_fp8_einsum_recipe,
|
||||
deep_gemm_fp8_o_proj,
|
||||
)
|
||||
from vllm.models.deepseek_v4_1.nvidia.model import DeepseekV4Model
|
||||
from vllm.models.deepseek_v4_1.quant_config import DeepseekV4FP8Config
|
||||
from vllm.utils.deep_gemm import is_deep_gemm_supported
|
||||
|
||||
if not is_deep_gemm_supported() or not current_platform.is_device_capability_family(
|
||||
100
|
||||
):
|
||||
pytest.skip("DeepGEMM MXFP8 BMM requires Blackwell")
|
||||
|
||||
default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16)
|
||||
quant_config = DeepseekV4FP8Config(
|
||||
is_checkpoint_fp8_serialized=True, weight_block_size=[32, 32]
|
||||
)
|
||||
quant_config._resolved_expert_dtype = "fp4"
|
||||
if config_source == "mxfp8":
|
||||
quant_config = get_quantization_config("mxfp8").from_config(
|
||||
{"quant_method": "mxfp8"}
|
||||
)
|
||||
with torch.device("cuda"):
|
||||
linear = ColumnParallelLinear(
|
||||
512,
|
||||
256,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
return_bias=False,
|
||||
)
|
||||
linear.is_bmm = True
|
||||
linear.bmm_batch_size = 2
|
||||
weight = torch.randn(256, 512, device="cuda").to(torch.float8_e4m3fn)
|
||||
scales = torch.randint(124, 131, (8, 16), device="cuda", dtype=torch.uint8)
|
||||
checkpoint_scales = (
|
||||
scales.view(torch.float8_e8m0fnu)
|
||||
if config_source == "deepseek"
|
||||
else scales.repeat_interleave(32, dim=0)
|
||||
)
|
||||
model = DeepseekV4Model.__new__(DeepseekV4Model)
|
||||
torch.nn.Module.__init__(model)
|
||||
block = torch.nn.Module()
|
||||
block.attn = torch.nn.Module()
|
||||
block.attn.wo_a = linear
|
||||
model.layers = torch.nn.ModuleList([block])
|
||||
model.config = SimpleNamespace(num_attention_heads=2)
|
||||
model.quant_config = quant_config
|
||||
model.use_sequence_parallel = False
|
||||
model.get_expert_mapping = lambda: []
|
||||
model.load_weights(
|
||||
[
|
||||
("layers.0.attn.wo_a.weight", weight),
|
||||
("layers.0.attn.wo_a.weight_scale", checkpoint_scales),
|
||||
]
|
||||
)
|
||||
reference_weight = (
|
||||
weight.float().reshape(8, 32, 16, 32)
|
||||
* torch.exp2(scales.float() - 127)[:, None, :, None]
|
||||
).reshape(2, 128, 512)
|
||||
linear.quant_method.process_weights_after_loading(linear)
|
||||
linear.quant_method.process_weights_after_loading(linear)
|
||||
assert isinstance(linear.quant_method.kernel, DeepGemmMxfp8BmmLinearKernel)
|
||||
assert linear.weight.shape == (2, 128, 512)
|
||||
torch.testing.assert_close(
|
||||
linear.weight.view(torch.uint8).flatten(),
|
||||
weight.view(torch.uint8).flatten(),
|
||||
rtol=0,
|
||||
atol=0,
|
||||
)
|
||||
unpacked_scales = torch.stack(
|
||||
[(linear.weight_scale >> (8 * i)) & 0xFF for i in range(4)], dim=-1
|
||||
).flatten(-2)
|
||||
torch.testing.assert_close(
|
||||
unpacked_scales.to(torch.uint8).reshape(256, 16),
|
||||
scales.repeat_interleave(32, dim=0),
|
||||
rtol=0,
|
||||
atol=0,
|
||||
)
|
||||
|
||||
x = torch.randn(2, num_tokens, 512, device="cuda", dtype=torch.bfloat16)
|
||||
reference = torch.einsum("gmk,gnk->mgn", x.float(), reference_weight)
|
||||
inputs = x.transpose(0, 1)
|
||||
if prequantized:
|
||||
cache = torch.cat(
|
||||
(
|
||||
torch.ones(num_tokens, 32, device="cuda"),
|
||||
torch.zeros(num_tokens, 32, device="cuda"),
|
||||
),
|
||||
dim=1,
|
||||
)
|
||||
inputs = fused_inv_rope_fp8_quant(
|
||||
x.permute(1, 0, 2),
|
||||
torch.arange(num_tokens, device="cuda"),
|
||||
cache,
|
||||
n_groups=2,
|
||||
heads_per_group=1,
|
||||
nope_dim=448,
|
||||
rope_dim=64,
|
||||
quant_group_size=32,
|
||||
tma_aligned_scales=current_platform.has_device_capability(100),
|
||||
)
|
||||
with torch.no_grad():
|
||||
output = linear(inputs)
|
||||
assert output.shape == reference.shape
|
||||
assert (output.float() - reference).norm() / reference.norm() < 0.06
|
||||
if prequantized:
|
||||
recipe, tma_aligned_scales = compute_fp8_einsum_recipe(block_size=32)
|
||||
projected = deep_gemm_fp8_o_proj(
|
||||
x.permute(1, 0, 2),
|
||||
torch.arange(num_tokens, device="cuda"),
|
||||
cache,
|
||||
linear,
|
||||
torch.nn.Identity(),
|
||||
n_groups=2,
|
||||
heads_per_group=1,
|
||||
nope_dim=448,
|
||||
rope_dim=64,
|
||||
o_lora_rank=128,
|
||||
einsum_recipe=recipe,
|
||||
tma_aligned_scales=tma_aligned_scales,
|
||||
)
|
||||
torch.testing.assert_close(projected, output.flatten(1), rtol=0, atol=0)
|
||||
compiled = torch.compile(linear, backend="eager", fullgraph=True)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(compiled(inputs), output, rtol=0, atol=0)
|
||||
|
||||
|
||||
def test_prepare_gated_trtllm_fp8_moe_weights_pads_each_projection(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
flashinfer_utils,
|
||||
|
||||
@@ -127,6 +127,23 @@ def test_modelopt_nvfp4_quantizes_parallel_lm_head():
|
||||
assert method.spec.activation is kNvfp4Dynamic
|
||||
|
||||
|
||||
def test_modelopt_mxfp8_preserves_per_row_checkpoint_scales(dist_init, monkeypatch):
|
||||
"""Standard MXFP8 checkpoints already have one scale row per weight row."""
|
||||
from vllm.model_executor.layers.linear import ReplicatedLinear
|
||||
|
||||
kernel = Mock()
|
||||
kernel.input_quant_key.return_value = None
|
||||
monkeypatch.setattr(
|
||||
"vllm.model_executor.layers.quantization.modelopt.init_mxfp8_linear_kernel",
|
||||
lambda **kwargs: kernel,
|
||||
)
|
||||
config = ModelOptMxFp8Config.from_config({"quant_method": "mxfp8"})
|
||||
linear = ReplicatedLinear(64, 64, bias=False, quant_config=config)
|
||||
scales = torch.arange(128, dtype=torch.uint8).reshape(64, 2)
|
||||
linear.weight_scale.weight_loader(linear.weight_scale, scales)
|
||||
assert torch.equal(linear.weight_scale, scales)
|
||||
|
||||
|
||||
def test_modelopt_fp8_updates_weight_dims_after_transpose():
|
||||
"""Humming reads weight.input_dim/output_dim. Swapping the
|
||||
ModelWeightParameter for a plain Parameter drops them, so the per-tensor
|
||||
|
||||
Reference in New Issue
Block a user