Multiply with inverse scale in loop inside of dividing to improve perf

[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

---

Multiply with inverse scale in loop inside of dividing to improve perf

• NCU showed the kernel was now compute bound, which confused me since this should be a mem bw bound kernel.
• Debugging, I've learned doing __fdiv_rn 32 times in an unrolled loop is much slower than doing it once to calculate the inverse scale, then multiplying data * inverse_scale 32 times inside the loop inside.

Test plan

• pytest test_kernels.py -k cuda_mx -s

Performance

• Memory bandwidth utilization and runtime improved by ~47% with this change.

Before change:

[[email protected] ~/private-torchao/benchmarks/mx_formats (3631dea7)]$ CUDA_VISIBLE_DEVICES=2 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 234.592005610466
mem_bw_gbps 3468.5537296233347

After change:

[[email protected] ~/private-torchao/benchmarks/mx_formats (work)]$ CUDA_VISIBLE_DEVICES=2 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 159.10400450229645
mem_bw_gbps 5114.23316179484

Benchmarking against Triton

CUDA dim1 cast now achieves ~23% faster runtime and ~23% higher mem bw utilization than the Triton kernel with this change.

CUDA: 5114 gbps, 159us

[[email protected] ~/private-torchao/benchmarks/mx_formats (work)]$ CUDA_VISIBLE_DEVICES=2 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 159.10400450229645
mem_bw_gbps 5114.23316179484

Triton: 4137gbps, 197us

[[email protected] ~/private-torchao/benchmarks/mx_formats (work)]$ CUDA_VISIBLE_DEVICES=2 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 196.6720074415207
mem_bw_gbps 4137.31972630598

[danielvegamyhre]

fyi @drisspg @vkuzo dim1 cast perf is looking good now!