Swizzle scales

[nastya236]

Improvements to scales swizzling kernel.

Overview:

• Each thread loads 4 rows × 4 bytes = 16 bytes with stride 32, so a warp covers one full 128×4 tile.
• If K % 4 == 0: vectorized int32 loads
• If not: byte-by-byte loads with zero-padding ( reduction dim is typically divisible by group_size x 4, so we won't be there often).
• load: global-> registers -> shared
• coalesced store: shared -> global
• XOR swizzle to avoid bank conflicts

Some numbers for qqmm with new swizzling (variance is <0.5%):

M	N	K	Improvement (%)
16384	4096	11008	10.19
32768	4096	11008	9.40
65536	4096	11008	9.41
16384	11008	4096	13.54
32768	11008	4096	9.47
65536	11008	4096	13.79
16384	4096	12288	5.59
32768	4096	12288	5.06
65536	4096	12288	6.44
16384	12288	4096	13.48
32768	12288	4096	11.86
65536	12288	4096	15.02
16384	5120	27648	7.59
32768	5120	27648	5.12
65536	5120	27648	6.09
16384	27648	5120	13.18
32768	27648	5120	12.42
65536	27648	5120	11.34

Probably can be improved further!

[awni]

Awesome! Left a couple more nits. Please take a look then we can merge.

As a side-note I'd like to find a way to test the QQ ops in CI 🤔 but I'm drawing blanks. It might at some point make sense to implement slow fallbacks for the sm < 10.0, since we just need the matmuls. Otherwise it could be a while before CI has blackwell available.

[nastya236]

Thanks for the review! Regarding tests, this is a very good point. I tested it with examples/python/qqmm.py on B200. I’ll implement a fallback so we can test correctness