diff --git a/src/h3_blackwell_runtime/nvfp4_quant.py b/src/h3_blackwell_runtime/nvfp4_quant.py index d3e9bd4..b4e0662 100644 --- a/src/h3_blackwell_runtime/nvfp4_quant.py +++ b/src/h3_blackwell_runtime/nvfp4_quant.py @@ -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