P4: Fix dequantize_nvfp4 bridge — handle float8_e4m3fn dtype

This commit is contained in:
2026-06-02 16:31:56 +00:00
parent 29f836d711
commit 57ab4b9d4c

View File

@@ -367,6 +367,11 @@ def dequantize_nvfp4(x_fp4, x_sf, gsa, shape=None):
mod = get_cuda_module("dequant_nvfp4", ["dequant_nvfp4.cu"])
if gsa.dim() == 2:
gsa = gsa.squeeze(1) # (M, 1) → (M,)
# dequant kernel expects uint8 for both fp4 and sf
if x_fp4.dtype != torch.uint8:
x_fp4 = x_fp4.view(torch.uint8)
if x_sf.dtype != torch.uint8:
x_sf = x_sf.view(torch.uint8)
return mod.dequant_nvfp4(x_fp4, x_sf, gsa)