Match Torch NVFP4 scale precision

This commit is contained in:
Daniel Maddern 2026-08-15 01:57:25 +07:00
parent 0fda84502a
commit 684b645546

View file

@ -55,7 +55,8 @@ def nvfp4_activation_scale(tensor: torch.Tensor, *, timings: dict[str, list[floa
raise ValueError("Vortex NVFP4 scale backend requires contiguous input")
try:
extension = _vortex_scale_extension()
return _record_timing(timings, "vortex_absmax_scale", lambda: extension.nvfp4_activation_scale(tensor, divisor))
scale = _record_timing(timings, "vortex_absmax_scale", lambda: extension.nvfp4_activation_scale(tensor, divisor))
return _record_timing(timings, "scale_compat_cast", lambda: scale.to(tensor.dtype) if tensor.dtype != torch.float32 else scale)
except Exception:
if os.getenv("H3_NVFP4_SCALE_STRICT", "").lower() in {"1", "true", "yes", "on"}:
raise