daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Match Torch NVFP4 scale precision

Browse source
This commit is contained in:
Daniel Maddern 2026-08-15 01:57:25 +07:00
parent 0fda84502a
commit 684b645546
1 changed files with 2 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
src/h3_blackwell_runtime/nvfp4_quant.py
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") raise ValueError("Vortex NVFP4 scale backend requires contiguous input")
try: try:
extension = _vortex_scale_extension() 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: except Exception:
if os.getenv("H3_NVFP4_SCALE_STRICT", "").lower() in {"1", "true", "yes", "on"}: if os.getenv("H3_NVFP4_SCALE_STRICT", "").lower() in {"1", "true", "yes", "on"}:
raise raise