diff --git a/vllm/patches/deepseek_v4_attention.py b/vllm/patches/deepseek_v4_attention.py index 501898d1..53fdcb07 100644 --- a/vllm/patches/deepseek_v4_attention.py +++ b/vllm/patches/deepseek_v4_attention.py @@ -372,8 +372,12 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): wkv_wgate_weight = compressor.fused_wkv_wgate.weight if wkv_wgate_weight.dtype == torch.uint8: # NVFP4 packed weights — use forward() for dequant+matmul - score, _ = compressor.fused_wkv_wgate(hidden_states) - return score.to(torch.float32) + result = compressor.fused_wkv_wgate(hidden_states) + # MergedColumnParallelLinear may return (output, bias) or + # just output depending on quantization method. + if isinstance(result, tuple): + result = result[0] + return result.to(torch.float32) return torch.mm( hidden_states, wkv_wgate_weight.T, @@ -393,8 +397,10 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): def indexer_compressor_kv_score() -> torch.Tensor: wkv_wgate_weight = indexer.compressor.fused_wkv_wgate.weight if wkv_wgate_weight.dtype == torch.uint8: - score, _ = indexer.compressor.fused_wkv_wgate(hidden_states) - return score.to(torch.float32) + result = indexer.compressor.fused_wkv_wgate(hidden_states) + if isinstance(result, tuple): + result = result[0] + return result.to(torch.float32) return torch.mm( hidden_states, wkv_wgate_weight.T,