[WIP] Blockscaled SM90

[KareemMusleh]

I'm mainly doing this because I'm interested in mxfp8 sonicmoe. Here are some of my thoughts:

1. FP8 WGMMA requires the input to be k-major, which works fine for the forward but a naive implementation to compute the activation gradient requires us to transpose the weights. I think this can be avoided by using swap_AB. We can quant + transpose the activations in a single kernel
2. I should at least attempt doing fp8 gather_A it's possible that cp.async will be fast enough for loading the sfa with 32bit alignment. Edit: after thinking about it some more I think it might be possible to fuse the quant of the inputs with a permute of the scales. Which will allow us to both use gather_A for the input and use TMA for the SFA
3. Should probably swizzle the sfa. Maybe should also try storing the scales as e8m0

[tridao]

are you doing block size 128? mxfp8 block size 32 will be very slow on hopper imo.

[KareemMusleh]

are you doing block size 128? mxfp8 block size 32 will be very slow on hopper imo.

Yes block size (1, 128) for activations and (128, 128) for weights. Just like DeepGEMM

[tridao]

there's only 1 scale for B in the whole k_tile?

[KareemMusleh]

For now the kernel assumes tile_n == tile_k == 128. I'll first be adding support for tile_n == 192 so that we can have better auto tuning.

I think that tile_k != 128 is not gonna improve the perf. But if needed I'll implement it

[tridao]

no, i mean for SFA you have 1 scale per row of A. But for SFB you have 1 scale per "tile_n" columns of B?

[KareemMusleh]

yes because right now I hardcoded tile_n to be 128. I'll remove this constraint later.

Sorry this whole thing is messy I'll try to clean it up

[KareemMusleh]

Leaving this here as it might help with understanding the code

Link to PTX Docs

[KareemMusleh]

@tridao this seems to be getting closer to being done. I'm getting similar perf to DG (around 98-100% of DG).

Will be posting proper benchmarks tomorrow. If the perf is good enough I'll be moving on to implementing the actual fwd + bwd in sonicmoe. Though I still have some ideas that I want to try out. Like adding swizzle to SFA and overlapping the scale FMA with WGMMA

[KareemMusleh]

testing using latest DG with cuda 13.0. The comparison with m_varlen is unfair rn because DG asserts tile_m == 128 rather than tile_m == 256

Should add support for tile_n == 192 (also padding/masking) to have a fair comparison