[float8] Add support for blockwise fp8 quantization scheme used in DeepSeek v3

[danielvegamyhre]

DeepSeek v3 uses a blockwise fp8 quantization strategy, where the scaling factor is computed independently for each block, rather than for each tensor/row/etc. The code is available here.

It would be useful for torchao to support this as well, for users wishing to do research or development with this same quantization strategy.

cc @drisspg @vkuzo

[gau-nernst]

Just want to add some of my observations here. I played around abit with block-wise FP8 on my consumer GPU (4070Ti SUPER, sm89). A simple triton kernel does not perform really well, only 1.4x speedup over BF16 (for reference, row-wise FP8 is ~1.9x speedup). With dynamic quant overhead, e2e speedup won't be too attractive. (ofc optimizing for Hopper will be completely different).

Also tried block-wise INT8 (which is the main idea of JetFire). A simple triton kernel performs somewhat ok on consumer GPU (1.9x speedup over BF16, compared to row-wise INT8 is 2.9x speedup - note that INT8 matmul is 4x faster than BF16 on consumer GPUs), but on A100, couldn't get any speedup (speedup < 1). Probably because in the case of block-wise INT8, there is a dtype conversion from INT32 to FP32 when scaling MMA accumulate results, while FP8 does not.

For quantization BLOCK_SIZE_K (number of elements along K dim that share 1 scale value), I think only K<=128 would have a simple and performant implementation, since if BLOCK_SIZE_K is too big, we will use too much shared memory. Tried a few ways around it, such as loading tiles smaller than quantization BLOCK_SIZE_K, but couldn't make it fast.

[drisspg]

We should also take a look at the new blockwise fp8 gemm added in cutlass 3.7

cc @alexsamardzic

[the-tuning-machine]

A PR has been created for this issue: PR #1668

[danielvegamyhre]

#1763 has been merged w/ fp8 blockwise quant/dequant, fp8 blockwise gemm, and simple linear wrapper for anyone interested in experimenting with this.

[yiakwy-xpu-ml-framework-team]

Hi @danielvegamyhre I have added native support in PR#925 to quantize bf16 weights. Later I will added optimized quantized kernel in SGLang and flashFloat, let me know if it helps :

[danielvegamyhre]

Hi @danielvegamyhre I have added native support in PR#925 to quantize bf16 weights. Later I will added optimized quantized kernel in SGLang and flashFloat, let me know if it helps :

@yiakwy-xpu-ml-framework-team we already have kernels for this in torchao, see https://github.com/pytorch/ao/blob/main/torchao/prototype/blockwise_fp8/blockwise_quantization.py

We need an e2e float8 blockwise training experience using these kernels now :)

[yiakwy-xpu-ml-framework-team]

https://github.com/pytorch/ao/blob/main/torchao/prototype/blockwise_fp8/blockwise_quantization.py

That is nice!.

Note this line is wrong, cause, K is not alwasy divied by 128 :

ao/torchao/prototype/blockwise_fp8/blockwise_quantization.py

Line 204 in 19237e3

assert x.size(0) % block_size == 0 and x.size(1) % block_size == 0, (

Later I will udpate my algorithm directly to torchao, so that my people can have refrence to this work.