Routes norm(1/2, dim=-1) through the efficient sum_reduction_inner kernel via abs/square load + sqrt finalize, instead of the slow per-output-threadgroup norm kernel with per-element pow.
Example:
import torch
import torch.nn.functional as F
x = torch.randn(4096, 4096, device="mps", dtype=torch.bfloat16)
x.norm(2, dim=-1) # L2 -> norm_l2_reduction_inner kernel
x.norm(1, dim=-1) # L1 -> norm_l1_reduction_inner kernel
F.normalize(x, dim=-1)
Perf:
fp32
| shape | old (µs) | new (µs) | speedup |
|---|---|---|---|
| (32,768) | 6.9 | 3.8 | 1.8x |
| (128,1024) | 17.1 | 4.2 | 4.1x |
| (512,4096) | 171.9 | 20.9 | 8.2x |
| (2048,4096) | 676.8 | 114.5 | 5.9x |
| (4096,4096) | 1348.2 | 238.4 | 5.7x |
| (16384,1024) | 1601.8 | 236.2 | 6.8x |
| (1024,8192) | 666.3 | 112.0 | 5.9x |
| (65536,256) | 1707.3 | 229.9 | 7.4x |
| (8,32,768) | 27.5 | 4.2 | 6.6x |
| (32,128,128) | 73.1 | 7.4 | 9.9x |
bf16
| shape | old (µs) | new (µs) | speedup |
|---|---|---|---|
| (32,768) | 6.9 | 3.6 | 1.9x |
| (128,1024) | 17.0 | 3.7 | 4.6x |
| (512,4096) | 177.5 | 10.3 | 17.2x |
| (2048,4096) | 715.7 | 40.0 | 17.9x |
| (4096,4096) | 1402.8 | 117.6 | 11.9x |
| (16384,1024) | 1650.8 | 112.5 | 14.7x |
| (1024,8192) | 678.0 | 42.4 | 16.0x |
| (65536,256) | 1790.2 | 113.2 | 15.8x |
| (8,32,768) | 27.4 | 3.9 | 7.0x |
| (32,128,128) | 74.5 | 7.5 | 10.0x |
fp16
| shape | old (µs) | new (µs) | speedup |
|---|---|---|---|
| (32,768) | 6.8 | 3.8 | 1.8x |
| (128,1024) | 17.0 | 3.9 | 4.3x |
| (512,4096) | 172.6 | 10.2 | 16.9x |
| (2048,4096) | 678.6 | 40.4 | 16.8x |
| (4096,4096) | 1352.9 | 112.8 | 12.0x |
| (16384,1024) | 1604.2 | 107.6 | 14.9x |
| (1024,8192) | 670.2 | 42.5 | 15.8x |
| (65536,256) | 1711.6 | 108.1 | 15.8x |
| (8,32,768) | 27.5 | 4.0 | 6.8x |
| (32,128,128) | 73.2 | 7.7 | 9.5x |
Pull Request resolved: #186076
Approved by: https://github.com/malfet
SOCIAL SHARE CARD GENERATOR