Fix nvfp4 online scale with pcg (#32246)

2026-07-23 14:42:37 -07:00
parent 410ab4fde5
commit 378aea1385
1 changed files with 7 additions and 1 deletions
@@ -965,9 +965,15 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
        ):
            e4m3_max = 256.0

+        global_scale_inv = torch.full(
+            (1,),
+            1.0 / (e4m3_max * 6.0),
+            dtype=torch.float32,
+            device=hidden_states.device,
+        )
        hs_fp4_bytes, hs_sf_bytes, per_token_scale = nvfp4_quantize(
            hidden_states,
-            1.0 / (e4m3_max * 6.0),
+            global_scale_inv,
            sfLayout=SfLayout.layout_linear,
            per_token_activation=True,
            backend="cute-dsl",