diff --git a/single_shot_inference.py b/single_shot_inference.py index c85465a5..20e8a762 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -84,7 +84,9 @@ def dequant_nvfp4(weight, weight_scale, weight_scale_2=None, input_scale=None): w = torch.stack([lo_f, hi_f], -1).reshape(O, I) s = weight_scale.float().repeat_interleave(16, 1) if weight_scale_2 is not None: s = s * weight_scale_2.float() - if input_scale is not None: s = s * input_scale.float() + # NOTE: input_scale is intentionally NOT used. It's the activation + # quantization scale (for FP8 inputs). Since we use BF16 activations, + # the weight dequant is: lut[weight] * weight_scale * weight_scale_2. return (w * s).bfloat16() def nvfp4_linear(x, weight, weight_scale, weight_scale_2=None, input_scale=None):