[Feature Request] Support `mxfp/nvfp` data types and blockscaled GEMM APIs

[Rachmanino]

Required prerequisites

• I have searched the Issue Tracker that this hasn't already been reported. (comment there if it has.)

Motivation

mxfp/nvfp types with group scaling have been widely used in efficient AI training/inference. Currently TileLang support dequantgemm gemm for mxfp/nvfp via storing them in uint8, and requires to perform manual SIMT scaling and T.gemm, which maybe kinda compilcated to upper users. What's more, hardware native blockscaled gemm is supported on modern GPU architecture, e.g. SM100.

Solution

I propose that we may expose a unified T.blockscaled_gemm API in the frontend, which will be lowered to native blockscaled ptx instructions on supported architecture. For older architecture, we can implement this by appending SIMT scaling operations to mma macros.

Alternatives

No response

Additional context

No response

[qqq-tao]

I'd like to take this issue. Development in progress first on sm 120a.

[Rachmanino]

I'm already working on this and passed a simple mxfp8 blockscaled example on sm100 :)

[Rachmanino]

As #1945 merged into TileLang main branch, the mininal framework for supporting low-bit Tensor Core kernels on Blackwell architecture has been established. This will be a beginning, rather than the end of our development. We will focus on introducing related new features, including:

• support more dtypes and their blockscaled gemm lowering
• support generalized cases (e.g. transposed, grouped gemm, tcgen05-rs)
• reconsider and refactor relevent APIs for better clarity and usability