biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix unpack error: handle both tuple and tensor returns from NVFP4 forward()

Browse Source
This commit is contained in:
biondizzle
2026-05-19 00:33:43 +00:00
parent d9dc042ff7
commit 358830925a
1 changed files with 10 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
vllm/patches/deepseek_v4_attention.py
View File

@@ -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,