[bugfix] nvfp4: cast fp32 inputs to bf16 instead of asserting

NVFP4QuantizeMethod.quantize_input() and apply() asserted bf16/fp16 inputs
and crashed on the fp32 pre-attention norm output in eager mode (without the
torch.compile fusion that keeps it bf16):

  AssertionError: only allow bf16/fp16 inputs to fp4 linear, got torch.float32

The FP4 linear emits bf16 regardless (mm_fp4 out dtype is bf16), and the
sibling fastvideo/layers/fp4linear.py already casts non-bf16/fp16 inputs, so
cast fp32 -> bf16 here too instead of failing. Lets NVFP4Config inference run
in eager mode (repro: LTX2-distilled on DGX Spark / sm_121, no compile).
This commit is contained in:
Raghav
2026-07-06 23:10:17 -07:00
committed by SolitaryThinker
parent d921832cd2
commit 2e8eb75c11
@@ -285,8 +285,12 @@ class NVFP4QuantizeMethod(QuantizeMethodBase):
def quantize_input(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
SfLayout, _, _ = _require_flashinfer()
assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, (
f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}")
# The pre-attention norm can emit fp32 (e.g. in eager mode, without the
# torch.compile fusion that keeps it bf16). FP4 linear emits bf16 regardless
# (see apply(): _mm_fp4 out dtype is bf16), so cast fp32 inputs to bf16
# instead of failing — matching the sibling fastvideo/layers/fp4linear.py.
if x.dtype not in (torch.bfloat16, torch.float16):
x = x.to(torch.bfloat16)
x_2d = x.view(-1, x.shape[-1])
x_fp4, x_scale = _nvfp4_quantize(
x_2d,
@@ -332,8 +336,8 @@ class NVFP4QuantizeMethod(QuantizeMethodBase):
if x_scale.dim() > 2:
x_scale = x_scale.view(-1, x_scale.shape[-1])
else:
assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, (
f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}")
if x.dtype not in (torch.bfloat16, torch.float16):
x = x.to(torch.bfloat16)
x = x.view(-1, x.shape[-1])
x_global_sf = self.x_global_sf
x_fp4, x_scale = _nvfp4_quantize(