add bfloat16 support for colwise scaling

[danielvegamyhre]

Stacked PRs:

• Multiply with inverse scale in loop inside of dividing to improve perf #12
• Fix uncoalesced global accesses #11
• ->add bfloat16 support for colwise scaling #10
• Fix numerics for colwise scaling #9
• fix scale shapes #8
• update impls of some helpers and add comments #7

---

Add bfloat16 support for colwise scaling

• Add templating to enable TMA loads input data correctly for the given input data type (fp32 or bf16)
• Always convert amax to fp32. This way we can use the same "amax to scale" function for both fp32 and bf16 input tensors.
• Add tests for bfloat16 inputs
• Update benchmark script to operate on bfloat16 inputs for dim1 CUDA cast, so we can compare triton vs CUDA implementations head to head.

Test plan

• pytest test_kernels.py -k cuda_mx

Performance

Perf of this initial implementation is slightly worse than Triton (slightly lower mem bw utilization). I'm guessing it is at least in part due to the fact that I'm currently casting the bfloat 16 inputs to fp32 when moving from SMEM to registers. This was for simplicity to re-use the same "amax to scale" func as fp32. As a next step, I'll see if I can add a custom scale calc func for bf16 to avoid this.

Triton dim1 cast with bf16 input: ~2127gbps mem bw

(ao) [danvm@devgpu031.atn1 ~/private-torchao/benchmarks/mx_formats (work)]$ python cast_bench.py --mode dim1_mx_triton
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.9.0.dev20250702+cu128
triton version: 3.3.1
mode: dim1_mx_triton
time_us 382.4959993362427
mem_bw_gbps 2127.3293770706896

Cuda dim1 cast with bf16 input: ~1825gbps mem bw

(ao) [danvm@devgpu031.atn1 ~/private-torchao/benchmarks/mx_formats (work)]$ python cast_bench.py --mode dim1_cuda
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.9.0.dev20250702+cu128
triton version: 3.3.1
mode: dim1_cuda
time_us 445.82399725914
mem_bw_gbps 1825.148446477705


## Next steps
- Add custom scale calc func for bf16 to avoid unnecessary cast and improve perf 
- Investigate segfaults when inputs dims aren't divisible by 32