Take the linear weight amax as an inf-norm

The same reduction as the grouped weight amax, on the dense path: the
fp32 upcast materialized the whole weight to read a single scalar from
it. One inf-norm over the bf16 weight gives a bitwise identical amax.

Per weight at DeepSeek-V3 671B MLP shapes (7168x18432), the fp32 copy,
abs and fp32 reduce cost 548.7 us together; the bf16 inf-norm costs
86.2 us. Over a 3-linear MLP block at 8192 tokens that is 18.6% of the
CuteDSL block's device time (8111 -> 6603 us) and 16.4% of Triton's
(9313 -> 7783 us), and the saving does not shrink with batch size.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: rdspring1

torchao/prototype/moe_training/nvfp4_training/nvfp4_linear.py

@@ -116,7 +116,10 @@ def _weight_quantize_2d(x: torch.Tensor, use_cutedsl: bool):
     use_cutedsl selects the CuteDSL kernel (plain transpose-quantize via an identity Hadamard)
     over the Triton 2D weight kernel. Neither path applies RHT or SR.
     """
-    global_amax = x.float().abs().max()
+    # An inf-norm is the amax in one reduction. Upcasting first would instead
+    # materialize the whole weight in fp32, for the same result: bf16 -> fp32 is
+    # exact, so it reduces over the same values.
+    global_amax = torch.linalg.vector_norm(x, ord=float("inf")).float()
     quantize = cutedsl_weight_quantize_2d if use_cutedsl else triton_weight_quantize_2d
     codes, sf, t_codes, t_sf = quantize(x, global_amax)
     return (