diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 2e88e35a..fa244fd7 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -1711,8 +1711,11 @@ class DeepseekV4Model(nn.Module): because the stacking weight_loader corrupts NVFP4 uint8 data. """ # All attention projections to dequantize to BF16 - bf16_proj_names = {"wq_a", "wq_b", "wkv", "wo_a", "wo_b"} + # wo_a is excluded — it uses fp8_einsum (no input_scale, weight-only FP8) + bf16_proj_names = {"wq_a", "wq_b", "wkv", "wo_b"} + fp8_proj_names = {"wo_a"} bf16_converted = 0 + fp8_converted = 0 compressor_converted = 0 _shard_index = self._build_shard_index("/model") if os.path.isdir("/model") else None @@ -1721,7 +1724,7 @@ class DeepseekV4Model(nn.Module): for layer_idx, layer in tqdm(enumerate(self.layers), total=len(self.layers), desc=" (upcast)NVFP4→BF16 attn projs", unit="layer"): attn = layer.attn - # BF16 dequantization: all attention projections + # BF16 dequantization: attention projections (except wo_a) for proj_name in bf16_proj_names: if not hasattr(attn, proj_name): continue @@ -1733,6 +1736,19 @@ class DeepseekV4Model(nn.Module): self._dequant_nvfp4_to_bf16(mod, E2M1_LUT) bf16_converted += 1 + # FP8 conversion: wo_a (used by fp8_einsum, no input_scale) + FP8_MAX = torch.finfo(torch.float8_e4m3fn).max + for proj_name in fp8_proj_names: + if not hasattr(attn, proj_name): + continue + mod = getattr(attn, proj_name) + if not hasattr(mod, "weight"): + continue + if mod.weight.dtype in (torch.uint8, torch.int8): + E2M1_LUT = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16) + self._convert_nvfp4_to_fp8(mod, E2M1_LUT, FP8_MAX) + fp8_converted += 1 + # Compressor: still needs BF16 reconstruction mla_attn = getattr(attn, "mla_attn", None) if mla_attn is not None: