[FEA] Block-Scaled Data Format Support for SM120 in CuTe DSL

[DD-DuDa]

Which component requires the feature?

CuTe DSL

Feature Request

Is your feature request related to a problem? Please describe.
I wish I could use CuTe DSL (Python) to implement block-scaled GEMM operations on SM120 (Blackwell) architecture. Currently, all CuTe DSL Python examples for block-scaled formats (MXFP4, MXFP6, MXFP8) only support SM100.

Describe the solution you'd like
Add CuTe DSL Python support for SM120 block-scaled operations:

1. Python helper functions for SM120 (similar to existing SM100 helpers in blackwell_helpers.py)
2. SM120 MMA operation wrappers for block-scaled formats (MmaMXF4Op, MmaMXF8Op, etc.)
3. Example scripts for dense and grouped block-scaled GEMM on SM120
4. Support for MXFP4 (Float4E2M1FN), MXFP6, and MXFP8 formats with corresponding scale factor types

[Willie-Qu]

I have the same qustion as well. So far, I have implemented a function using inline ptx to use f8 mma instructions on sm_120.

[Willie-Qu]

Here's my implementation for mma.sync.aligned.kind::f8f6f4.m16n8k32.row.col.f32.e4m3.e4m3.f32, just for referrence.

@dsl_user_op
def mma_f8e4m3_16x8x32(
    a0: Uint32, a1: Uint32, a2: Uint32, a3: Uint32,
    b0: Uint32, b1: Uint32,
    c0: Float32, c1: Float32, c2: Float32, c3: Float32,
    *,
    loc=None, ip=None
) -> Tuple[Float32, Float32, Float32, Float32]:
    """
    FP8 MMA: D = A × B + C
    
    PTX: mma.sync.aligned.kind::f8f6f4.m16n8k32.row.col.f32.e4m3.e4m3.f32
    
    A fragment: 4 × u32 (每个 u32 = 4 个 FP8)
    B fragment: 2 × u32 (每个 u32 = 4 个 FP8)
    C/D: 4 × f32
    """
    inputs = [
        a0.ir_value(loc=loc, ip=ip), a1.ir_value(loc=loc, ip=ip),
        a2.ir_value(loc=loc, ip=ip), a3.ir_value(loc=loc, ip=ip),
        b0.ir_value(loc=loc, ip=ip), b1.ir_value(loc=loc, ip=ip),
        c0.ir_value(loc=loc, ip=ip), c1.ir_value(loc=loc, ip=ip),
        c2.ir_value(loc=loc, ip=ip), c3.ir_value(loc=loc, ip=ip),
    ]
    
    result_type = llvm.StructType.get_literal([T.f32(), T.f32(), T.f32(), T.f32()])
    
    result = llvm.inline_asm(
        result_type,
        inputs,
        "mma.sync.aligned.kind::f8f6f4.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
        "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};",
        "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
    )
    
    d0 = Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip))
    d1 = Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip))
    d2 = Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip))
    d3 = Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip))
    
    return (d0, d1, d2, d3)

[DD-DuDa]

Wow, thanks, Willie. I will try to use this currently.