P4: Fix dequantize_nvfp4 bridge — handle float8_e4m3fn dtype
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user