From 3445bd24c11d60932fdd3c0f03ab3518af6f2bd7 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 05:36:34 +0000 Subject: [PATCH] =?UTF-8?q?feat:=20keep=20attention=20weights=20native=20N?= =?UTF-8?q?VFP4=20=E2=80=94=20stop=20dequantizing=20to=20BF16?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _convert_nvfp4_post_load() was converting wq_b, wo_b, fused_wqa_wkv from NVFP4→BF16. These layers already have FlashInferCutlassNvFp4LinearKernel registered as their quant_method — they CAN run native NVFP4. Now only wo_a gets FP8 conversion (fp8_einsum requires FP8) and compressor gets BF16 reconstruction (weight_loader issue). Everything else stays NVFP4 native — Blackwell FP4 acceleration for the full model, not just the MoE experts. This also eliminates the 5-minute NVFP4→BF16 conversion loop. --- vllm/patches/deepseek_v4.py | 45 +++++++++---------------------------- 1 file changed, 11 insertions(+), 34 deletions(-) diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 9054a18c..e3a01fdf 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -1609,37 +1609,24 @@ class DeepseekV4Model(nn.Module): def _convert_nvfp4_post_load(self): """Post-load conversion of NVFP4 weights for vLLM compatibility. - Strategy: - - wo_a: Convert to FP8 (attention forward reads weight/weight_scale_inv - directly and passes to deepseek_v4_fp8_einsum, bypassing quant_method) - - fused_wqa_wkv, wq_b, wo_b: Dequant NVFP4->bf16 (called via - .forward() which goes through quant_method; FP8 would dtype-mismatch) - - compressor.fused_wkv_wgate: Dequant NVFP4->bf16 (used via direct - torch.mm in attention parallel stream) - - shared_experts (gate_up_proj, down_proj): Stay native NVFP4 via DeepGEMM fp8_fp4_gemm - - MoE experts: Handled by DeepseekV4MegaMoEExperts (NVFP4 native) + Only wo_a needs FP8 conversion (attention forward uses fp8_einsum + which requires FP8 inputs). All other NVFP4 weights stay native — + vLLM's FlashInferCutlassNvFp4LinearKernel handles them directly. + + Compressor weights are reconstructed from checkpoint sub-weights + because the stacking weight_loader corrupts NVFP4 uint8 data. """ - E2M1_LUT = torch.tensor( - [0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16 - ) FP8_MAX = torch.finfo(torch.float8_e4m3fn).max - # wo_a: attention forward reads .weight and .weight_scale_inv directly - # for fp8_einsum. Only layer that needs FP8 conversion. + # Only wo_a needs conversion — fp8_einsum requires FP8 weight + scale fp8_proj_names = {"wo_a"} - # No BF16 dequant paths active — cuBLAS BF16 is broken on Blackwell. - # wo_a goes NVFP4→FP8; compressor gets reconstructed from checkpoint; - # MoE experts stay native NVFP4 via CUTLASS kernel. - fp8_converted = 0 - fp8_from_bf16 = 0 compressor_converted = 0 - # Build shard index once for compressor reconstruction (avoids N×M full-shard loads) _shard_index = self._build_shard_index("/model") if os.path.isdir("/model") else None from tqdm import tqdm - for layer_idx, layer in tqdm(enumerate(self.layers), total=len(self.layers), desc=" (upcast)NVFP4→FP8/BF16 convert", unit="layer"): + for layer_idx, layer in tqdm(enumerate(self.layers), total=len(self.layers), desc=" (upcast)NVFP4→FP8 wo_a only", unit="layer"): attn = layer.attn # FP8 conversion: only wo_a @@ -1650,28 +1637,18 @@ class DeepseekV4Model(nn.Module): if not hasattr(mod, "weight"): continue if mod.weight.dtype in (torch.uint8, torch.int8): - # NVFP4 -> dequant to bf16 -> requant to FP8 + 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 - elif mod.weight.dtype == torch.bfloat16: - self._convert_bf16_to_fp8(mod, FP8_MAX) - fp8_from_bf16 += 1 - # Compressor: fused_wkv_wgate used via direct torch.mm - # Compressor weights were SKIPPED during loading (skip patterns) - # because the stacking weight_loader corrupts NVFP4 uint8 data. - # We reconstruct the bf16 weight from the individual sub-weights - # that were loaded separately before stacking. - # Note: compressor.kv_proj.weight and compressor.gate_proj.weight - # are skipped, so fused_wkv_wgate.weight is zeros (empty tensor). - # We need to manually create it. + # Compressor: still needs BF16 reconstruction mla_attn = getattr(attn, "mla_attn", None) if mla_attn is not None: + E2M1_LUT = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16) compressor = getattr(mla_attn, "compressor", None) if compressor is not None and hasattr(compressor, "fused_wkv_wgate"): compressor_converted += self._reconstruct_compressor_weight( compressor.fused_wkv_wgate, attn, layer_idx, E2M1_LUT, _shard_index=_shard_index) - # Indexer compressor (C4A layers only) indexer = getattr(mla_attn, "indexer", None) if indexer is not None: idx_compressor = getattr(indexer, "compressor", None)