vllm.models.minimax_m3.amd.ops.flydsl_utils

Helpers shared by the MiniMax-M3 FlyDSL kernels (gfx950).

Functions:

• cvt_fp4x2_f16 –

  The two e2m1 codes of byte sel of word as two fp16 (exact).

• cvt_fp8x2_f16 –

  E4M3 bytes 0, 1 (hi: 2, 3) of word as two fp16 (exact).

• dequant8 –

  8 e2m1 codes (low nibble the even dim) and the group's e4m3 table ->

• exp2 –

  v_exp_f32 without denormal range reduction.

• key_order4 –

  4 dwords of e4m3 (16 dims) -> even dims, then odd dims per 8 (key order).

• maxf_nn –

  v_max_f32 with nnan; without it LLVM canonicalizes both inputs first.

• mfma_16x16x32 –

  v_mfma_f32_16x16x32_{bf16,f16} on 4-dword operands.

• mfma_fp8_16x16x128 –

  v_mfma_f32_16x16x128_f8f6f4 on e4m3 operands; zero scales select the

• pack8 –

  Two 4 x i32 halves -> the 8 x i32 operand of the fp8 16x16x128 MFMA.

• perm –

  v_perm_b32: byte i of the result is byte sel_i of (hi:lo).

• permlane_swap –

  v_permlane{16,32}_swap; fx has no wrapper.

• pk_bf16 –

  Two fp32 -> one dword of two bf16 (RNE), a low.

• pk_f16 –

  Two fp32 -> one dword of two fp16 (RNE), a low.

• run_compiled –

  Compiles and runs exe on the first call, then dispatches the cache.

• xlane4 –

  op over lanes L, L^16, L^32, L^48, broadcast to all four.

cvt_fp4x2_f16(word, sel)

The two e2m1 codes of byte sel of word as two fp16 (exact).

Source code in vllm/models/minimax_m3/amd/ops/flydsl_utils.py

def cvt_fp4x2_f16(word, sel):
    """The two e2m1 codes of byte ``sel`` of ``word`` as two fp16 (exact)."""
    return Vec(
        fx.rocdl.cvt_scalef32_pk_f16_fp4(
            Vec.make_type(2, fx.Float16),
            fx.as_ir_value(fx.Int32(word)),
            fx.as_ir_value(fx.Float32(1.0)),
            sel,
        )
    )

cvt_fp8x2_f16(word, hi)

E4M3 bytes 0, 1 (hi: 2, 3) of word as two fp16 (exact).

Source code in vllm/models/minimax_m3/amd/ops/flydsl_utils.py

def cvt_fp8x2_f16(word, hi):
    """E4M3 bytes 0, 1 (``hi``: 2, 3) of ``word`` as two fp16 (exact)."""
    return Vec(
        fx.rocdl.cvt_scalef32_pk_f16_fp8(
            Vec.make_type(2, fx.Float16),
            fx.as_ir_value(fx.Int32(word)),
            fx.as_ir_value(fx.Float32(1.0)),
            hi,
        )
    )

dequant8(c, t_lo, t_hi)

8 e2m1 codes (low nibble the even dim) and the group's e4m3 table -> the e4m3 keys of the even dims, then of the odd dims.

Source code in vllm/models/minimax_m3/amd/ops/flydsl_utils.py

def dequant8(c, t_lo, t_hi):
    """8 e2m1 codes (low nibble the even dim) and the group's e4m3 table ->
    the e4m3 keys of the even dims, then of the odd dims."""
    ml = perm(t_hi, t_lo, c & fx.Int32(0x07070707))
    mh = perm(t_hi, t_lo, (c >> 4) & fx.Int32(0x07070707))
    sign = fx.Int32(0x80808080 - (1 << 32))
    return ml ^ ((c << 4) & sign), mh ^ (c & sign)

exp2(x)

v_exp_f32 without denormal range reduction.

Source code in vllm/models/minimax_m3/amd/ops/flydsl_utils.py

def exp2(x):
    """v_exp_f32 without denormal range reduction."""
    return fx.Float32(fx.rocdl.exp2(fx.Float32.ir_type, fx.Float32(x).ir_value()))

key_order4(d)

4 dwords of e4m3 (16 dims) -> even dims, then odd dims per 8 (key order).

Source code in vllm/models/minimax_m3/amd/ops/flydsl_utils.py

def key_order4(d):
    """4 dwords of e4m3 (16 dims) -> even dims, then odd dims per 8 (key order)."""
    return Vec.from_elements(
        [
            perm(d[1], d[0], PERM_EVEN),
            perm(d[1], d[0], PERM_ODD),
            perm(d[3], d[2], PERM_EVEN),
            perm(d[3], d[2], PERM_ODD),
        ],
        fx.Int32,
    )

maxf_nn(a, b)

v_max_f32 with nnan; without it LLVM canonicalizes both inputs first.

Source code in vllm/models/minimax_m3/amd/ops/flydsl_utils.py

def maxf_nn(a, b):
    """v_max_f32 with nnan; without it LLVM canonicalizes both inputs first."""
    fm = ir.Attribute.parse("#arith.fastmath<nnan>")
    return fx.Float32(
        arith.MaxNumFOp(fx.as_ir_value(a), fx.as_ir_value(b), fastmath=fm).result
    )

mfma_16x16x32(a4, b4, acc, fp16)

v_mfma_f32_16x16x32_{bf16,f16} on 4-dword operands.

Source code in vllm/models/minimax_m3/amd/ops/flydsl_utils.py

def mfma_16x16x32(a4, b4, acc, fp16: bool):
    """v_mfma_f32_16x16x32_{bf16,f16} on 4-dword operands."""
    dtype = fx.Float16 if fp16 else fx.BFloat16
    op = fx.rocdl.mfma_f32_16x16x32_f16 if fp16 else fx.rocdl.mfma_f32_16x16x32_bf16
    a = Vec(a4).bitcast(dtype)
    b = Vec(b4).bitcast(dtype)
    return Vec(op(Vec.make_type(4, fx.Float32), [a, b, acc, 0, 0, 0]))

mfma_fp8_16x16x128(a8, b8, acc4)

v_mfma_f32_16x16x128_f8f6f4 on e4m3 operands; zero scales select the unscaled instruction.

Source code in vllm/models/minimax_m3/amd/ops/flydsl_utils.py

def mfma_fp8_16x16x128(a8, b8, acc4):
    """v_mfma_f32_16x16x128_f8f6f4 on e4m3 operands; zero scales select the
    unscaled instruction."""
    return Vec(
        fx.rocdl.mfma_scale_f32_16x16x128_f8f6f4(
            Vec.make_type(4, fx.Float32), [a8, b8, acc4, 0, 0, 0, 0, 0, 0]
        )
    )

pack8(h0, h1)

Two 4 x i32 halves -> the 8 x i32 operand of the fp8 16x16x128 MFMA.

Source code in vllm/models/minimax_m3/amd/ops/flydsl_utils.py

def pack8(h0, h1):
    """Two 4 x i32 halves -> the 8 x i32 operand of the fp8 16x16x128 MFMA."""
    a, b = Vec(h0), Vec(h1)
    return Vec.from_elements(
        [a[k] for k in range(4)] + [b[k] for k in range(4)], fx.Int32
    )

perm(hi, lo, sel)

v_perm_b32: byte i of the result is byte sel_i of (hi:lo).

Source code in vllm/models/minimax_m3/amd/ops/flydsl_utils.py

def perm(hi, lo, sel):
    """v_perm_b32: byte i of the result is byte sel_i of (hi:lo)."""
    return fx.Int32(rocdl.perm_b32(fx.Int32(hi), fx.Int32(lo), fx.Int32(sel)))

permlane_swap(op, d_a, d_b)

v_permlane{16,32}_swap; fx has no wrapper.

Source code in vllm/models/minimax_m3/amd/ops/flydsl_utils.py

def permlane_swap(op, d_a, d_b):
    """v_permlane{16,32}_swap; fx has no wrapper."""
    pair_ty = ir.Type.parse("!llvm.struct<(i32, i32)>")
    res = op(pair_ty, fx.as_ir_value(d_a), fx.as_ir_value(d_b), False, False)
    return (
        fx.Int32(llvm.extractvalue(T.i32, res, [0])),
        fx.Int32(llvm.extractvalue(T.i32, res, [1])),
    )

pk_bf16(a, b)

Two fp32 -> one dword of two bf16 (RNE), a low.

Source code in vllm/models/minimax_m3/amd/ops/flydsl_utils.py

def pk_bf16(a, b):
    """Two fp32 -> one dword of two bf16 (RNE), ``a`` low."""
    return fx.Int32(
        fx.rocdl.cvt_pk_bf16_f32(
            fx.as_ir_value(fx.Float32(a)), fx.as_ir_value(fx.Float32(b))
        )
    )

pk_f16(a, b)

Two fp32 -> one dword of two fp16 (RNE), a low.

Source code in vllm/models/minimax_m3/amd/ops/flydsl_utils.py

def pk_f16(a, b):
    """Two fp32 -> one dword of two fp16 (RNE), ``a`` low."""
    return (
        Vec.from_elements([fx.Float32(a), fx.Float32(b)], fx.Float32)
        .to(fx.Float16)
        .bitcast(fx.Int32)[0]
    )

run_compiled(exe, *args, waves_per_eu=None)

Compiles and runs exe on the first call, then dispatches the cache.

Source code in vllm/models/minimax_m3/amd/ops/flydsl_utils.py

def run_compiled(exe, *args, waves_per_eu: int | None = None):
    """Compiles and runs ``exe`` on the first call, then dispatches the cache."""
    import flydsl.compiler as flyc
    from flydsl.compiler.kernel_function import CompilationContext

    cf = getattr(exe, "_cf", None)
    if cf is None:
        hints = {"waves_per_eu": waves_per_eu} if waves_per_eu else {}
        with CompilationContext.compile_hints(hints):
            exe._cf = flyc.compile(exe, *args)
    else:
        cf(*args)

xlane4(x, op)

op over lanes L, L^16, L^32, L^48, broadcast to all four.

Source code in vllm/models/minimax_m3/amd/ops/flydsl_utils.py

def xlane4(x, op):
    """``op`` over lanes L, L^16, L^32, L^48, broadcast to all four."""
    a, b = permlane_swap(rocdl.permlane32_swap, bits(x), bits(x))
    m = op(as_f32(a), as_f32(b))
    a, b = permlane_swap(rocdl.permlane16_swap, bits(m), bits(m))
    return op(as_f32(a), as_f32(b))