mx4 quantize: compute per-group scale via exact FP32 bit-construction instead of FP64 exp2

[qchip]

Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/2986

_kernel_quantize_mx4 computes the per-group scale as scale = tl.exp2(group_exp.to(tl.float64)).to(tl.float32) — an FP64 transcendental evaluated once per group. FP64 throughput on datacenter GPUs (including Blackwell / GB300) is a small fraction of FP32, and this cost scales with the number of groups (numel / group_size), so it dominates the kernel and is worst at small group sizes.

group_exp is an integer-valued shared exponent clamped to [-127, 125], so 2^group_exp is always an exact power of two. For group_exp >= -126 the result is a normal FP32 with zero mantissa, which we build directly by setting the exponent field: ((group_exp + 127) << 23) bitcast to FP32. group_exp == -127 corresponds to the subnormal 2^-127 and is handled with a constant. This removes the per-group FP64 transcendental entirely.

This is bit-identical to the previous behavior: both the old exp2(fp64) -> fp32 path and the new bit-construction produce the exact FP32 value of 2^group_exp for every group_exp in [-127, 125], including the 2^-127 subnormal edge.

Measured on a GB300 (NVIDIA GB300, sm_103, CUDA 13.0) with tritonbench fp32_to_mx4 on real MX4-comm production tensor sizes (forward group_size=8, backward group_size=16):

  numel          group  before (FP64)   after (this)   speedup
  1,409,400,832  16     12.06 ms        2.60 ms        4.6x
  2,329,119,232  8      43.92 ms        9.27 ms        4.7x
  2,381,231,104  8      44.92 ms        9.49 ms        4.7x
  2,421,580,288  8      45.70 ms        9.65 ms        4.7x
  1,409,286,144  8      24.09 ms        3.18 ms        7.6x

The forward group_size=8 kernels used in the quantized All2All_Pooled comm path drop from ~44-46 ms to ~9.5 ms. The win applies to all group sizes (the FP64 exp2 was a global cost); smaller groups benefit most because the per-group exp2 count is higher.

Differential Revision: D107485047

[meta-codesync]

@qchip has exported this pull request. If you are a Meta employee, you can view the originating Diff in D107485047.