Skip to content

vllm.models.deepseek_v41.amd.mono.common.gfx942

gfx942 (MI300X, MI325X) replacements for the gfx950 instructions the mono kernels use.

FP8 on gfx942 is FNUZ: the same bit patterns as OCP e4m3 stand for half the OCP value, and 0x80 is NaN instead of -0. The kernels quantize with an FP8 maximum of 224 there (arch.GFX942), which gives the same bit patterns as OCP quantization with 448 and a scale code one higher. A product of two bytes then needs 2^(code_a + code_b - 254) with the FNUZ codes, which is what mx_factor returns.

gfx942 has no scaled MFMA. A GEMV runs one v_mfma_f32_16x16x32_fp8_fp8 per 32-wide K block (one scale block) with a zero accumulator and adds the result times the block's scale factor (mfma_fp8_scaled). A lane holds 8 bytes of each block: K 8 (lane / 16) .. + 8 of the block, for its row (A) or column (B).

Functions:

  • acc_scaled –

    acc + d * f for the 4 accumulator values of a lane.

  • bf16_hi_pair –

    Two f32 that a bf16 holds exactly -> one dword of their bf16 (a low).

  • fp4x8_fnuz –

    Eight e2m1 values in a dword (value i in bits 4 i .. 4 i + 3) -> two

  • fp8x2_f32 –

    Two FNUZ bytes of a dword as f32: bytes 0 and 1, or 2 and 3 when hi.

  • mfma_bf16_k32 –

    The 16x16x32 bf16 product of gfx950's mfma_f32_16x16x32_bf16 as two

  • mfma_fp8 –

    One v_mfma_f32_16x16x32_fp8_fp8: A and B are a lane's 8 FNUZ bytes as

  • mfma_fp8_scaled –

    acc plus one 32-wide K block's product: the FP8 MFMA with a zero

  • mx_factor –

    2^(code_sum - 254) as f32, for code_sum the sum of two E8M0 codes. The

  • ocp_to_fnuz –

    Four OCP e4m3 bytes -> the FNUZ bytes for half their values. Every bit

  • pair_i64 –

    Two dwords as the 64-bit operand of an FP8 MFMA, lo in the low half.

  • pin_acc –

    The 4 accumulator values of a lane through opaque moves with side

acc_scaled(acc, d, f)

acc + d * f for the 4 accumulator values of a lane.

Source code in vllm/models/deepseek_v41/amd/mono/common/gfx942.py 
53
54
55
56
57
def acc_scaled(acc, d, f):
    """``acc`` + ``d`` * ``f`` for the 4 accumulator values of a lane."""
    return fx.Vector.from_elements(
        [fx.Float32(fmath.fma(d[i], f, acc[i])) for i in range(4)], fx.Float32
    )

bf16_hi_pair(a, b)

Two f32 that a bf16 holds exactly -> one dword of their bf16 (a low). The bf16 of such a value is the top half of its f32 bits.

Source code in vllm/models/deepseek_v41/amd/mono/common/gfx942.py 
126
127
128
129
130
131
def bf16_hi_pair(a, b):
    """Two f32 that a bf16 holds exactly -> one dword of their bf16 (a low).
    The bf16 of such a value is the top half of its f32 bits."""
    return fx.Int32(
        rocdl.perm_b32(b.bitcast(fx.Int32), a.bitcast(fx.Int32), fx.Int32(0x07060302))
    )

fp4x8_fnuz(w)

Eight e2m1 values in a dword (value i in bits 4 i .. 4 i + 3) -> two dwords of FNUZ e4m3 bytes: (values 0, 2, 4, 6) and (values 1, 3, 5, 7).

Each byte is the e2m1 value times 2^-7, exactly, the subnormal 0.5 included: the e2m1 sign goes to bit 7 and its exponent and mantissa bits to bits 4 .. 2, which are the low two exponent bits and the top mantissa bit of e4m3 with bias 8. A scale code for these bytes is the e8m0 code plus 7. A negative zero nibble (0x8) would give 0x80, the FNUZ NaN, so the weights must not hold one: weights942.fp4_drop_negative_zero_ rewrites them to +0 at load, which changes no product.

Source code in vllm/models/deepseek_v41/amd/mono/common/gfx942.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def fp4x8_fnuz(w):
    """Eight e2m1 values in a dword (value i in bits 4 i .. 4 i + 3) -> two
    dwords of FNUZ e4m3 bytes: (values 0, 2, 4, 6) and (values 1, 3, 5, 7).

    Each byte is the e2m1 value times 2^-7, exactly, the subnormal 0.5
    included: the e2m1 sign goes to bit 7 and its exponent and mantissa bits
    to bits 4 .. 2, which are the low two exponent bits and the top mantissa
    bit of e4m3 with bias 8. A scale code for these bytes is the e8m0 code
    plus 7. A negative zero nibble (0x8) would give 0x80, the FNUZ NaN, so
    the weights must not hold one: ``weights942.fp4_drop_negative_zero_``
    rewrites them to +0 at load, which changes no product."""
    w = fx.Int32(w)
    even = ((w << 2) & 0x1C1C1C1C) | ((w << 4) & SIGN_BYTES)
    odd = ((w >> 2) & 0x1C1C1C1C) | (w & SIGN_BYTES)
    return even, odd

fp8x2_f32(word, hi)

Two FNUZ bytes of a dword as f32: bytes 0 and 1, or 2 and 3 when hi.

Source code in vllm/models/deepseek_v41/amd/mono/common/gfx942.py 
120
121
122
123
def fp8x2_f32(word, hi):
    """Two FNUZ bytes of a dword as f32: bytes 0 and 1, or 2 and 3 when ``hi``."""
    v = fx.Vector(rocdl.cvt_pk_f32_fp8(T.vec(2, T.f32), fx.Int32(word), hi))
    return v[0], v[1]

mfma_bf16_k32(a, b, c)

The 16x16x32 bf16 product of gfx950's mfma_f32_16x16x32_bf16 as two 16x16x16 bf16 MFMAs: lane l holds K 8 (l // 16) .. + 8 in a and b (8 bf16 each). The first MFMA takes elements 0 .. 3 of every lane, the second elements 4 .. 7, so together they cover all 32 K.

Source code in vllm/models/deepseek_v41/amd/mono/common/gfx942.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
def mfma_bf16_k32(a, b, c):
    """The 16x16x32 bf16 product of gfx950's ``mfma_f32_16x16x32_bf16`` as two
    16x16x16 bf16 MFMAs: lane l holds K 8 (l // 16) .. + 8 in ``a`` and
    ``b`` (8 bf16 each). The first MFMA takes elements 0 .. 3 of every lane,
    the second elements 4 .. 7, so together they cover all 32 K."""
    a16 = fx.Vector(a).bitcast(fx.Int16)
    b16 = fx.Vector(b).bitcast(fx.Int16)
    for h in range(2):
        ah = fx.Vector.from_elements([a16[4 * h + i] for i in range(4)], fx.Int16)
        bh = fx.Vector.from_elements([b16[4 * h + i] for i in range(4)], fx.Int16)
        c = fx.Vector(
            rocdl.mfma_f32_16x16x16bf16_1k(T.vec(4, T.f32), [ah, bh, c, 0, 0, 0])
        )
    return c

mfma_fp8(a, b, c)

One v_mfma_f32_16x16x32_fp8_fp8: A and B are a lane's 8 FNUZ bytes as an i64 (pair_i64), C a Vector of 4 f32.

Source code in vllm/models/deepseek_v41/amd/mono/common/gfx942.py 
38
39
40
41
42
43
def mfma_fp8(a, b, c):
    """One v_mfma_f32_16x16x32_fp8_fp8: A and B are a lane's 8 FNUZ bytes as
    an i64 (``pair_i64``), C a Vector of 4 f32."""
    return fx.Vector(
        rocdl.mfma_f32_16x16x32_fp8_fp8(T.vec(4, T.f32), [a, b, c, 0, 0, 0])
    )

mfma_fp8_scaled(acc, a, b, code_sum)

acc plus one 32-wide K block's product: the FP8 MFMA with a zero accumulator, then times the block's scale factor (mx_factor).

Source code in vllm/models/deepseek_v41/amd/mono/common/gfx942.py 
84
85
86
87
88
def mfma_fp8_scaled(acc, a, b, code_sum):
    """``acc`` plus one 32-wide K block's product: the FP8 MFMA with a zero
    accumulator, then times the block's scale factor (``mx_factor``)."""
    zero = fx.Vector.filled(4, 0.0, fx.Float32)
    return acc_scaled(acc, mfma_fp8(a, b, zero), mx_factor(code_sum))

mx_factor(code_sum)

2^(code_sum - 254) as f32, for code_sum the sum of two E8M0 codes. The f32 bits of 2^e are (e + 127) << 23, so this is (code_sum - 127) << 23. The sum must stay in [128, 381]: the factor is then a normal f32.

Source code in vllm/models/deepseek_v41/amd/mono/common/gfx942.py 
46
47
48
49
50
def mx_factor(code_sum):
    """2^(code_sum - 254) as f32, for code_sum the sum of two E8M0 codes. The
    f32 bits of 2^e are (e + 127) << 23, so this is (code_sum - 127) << 23.
    The sum must stay in [128, 381]: the factor is then a normal f32."""
    return ((code_sum - 127) << 23).bitcast(fx.Float32)

ocp_to_fnuz(w)

Four OCP e4m3 bytes -> the FNUZ bytes for half their values. Every bit pattern keeps its meaning (halved) except 0x80 and the OCP NaNs. 0x80, OCP's -0, is the FNUZ NaN and becomes 0x00. The NaNs 0x7F and 0xFF stay as they are and read as 240 and -240. A byte's magnitude plus 0x7F carries into its sign bit exactly when the magnitude is not zero, without reaching the next byte, so the sign survives only on nonzero bytes.

Source code in vllm/models/deepseek_v41/amd/mono/common/gfx942.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def ocp_to_fnuz(w):
    """Four OCP e4m3 bytes -> the FNUZ bytes for half their values. Every bit
    pattern keeps its meaning (halved) except 0x80 and the OCP NaNs. 0x80,
    OCP's -0, is the FNUZ NaN and becomes 0x00. The NaNs 0x7F and 0xFF stay
    as they are and read as 240 and -240. A byte's magnitude plus 0x7F
    carries into its sign bit exactly when the magnitude is not zero, without
    reaching the next byte, so the sign survives only on nonzero bytes."""
    w = fx.Int32(w)
    nonzero = ((w & MAG_BYTES) + MAG_BYTES) & SIGN_BYTES
    return w & (nonzero | MAG_BYTES)

pair_i64(lo, hi)

Two dwords as the 64-bit operand of an FP8 MFMA, lo in the low half.

Source code in vllm/models/deepseek_v41/amd/mono/common/gfx942.py 
31
32
33
34
35
def pair_i64(lo, hi):
    """Two dwords as the 64-bit operand of an FP8 MFMA, lo in the low half."""
    return fx.Vector.from_elements([fx.Int32(lo), fx.Int32(hi)], fx.Int32).bitcast(
        fx.Int64
    )[0]

pin_acc(acc)

The 4 accumulator values of a lane through opaque moves with side effects. An unrolled GEMV loop pins its accumulators after each round. The round's FMAs must then finish before the moves, and the next round's LDS reads (and the MFMAs that need them) cannot move above them. Without it LLVM runs every MFMA of the loop first and the scaling FMAs at the end, which needs a register for each MFMA result and spills them all.

Source code in vllm/models/deepseek_v41/amd/mono/common/gfx942.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def pin_acc(acc):
    """The 4 accumulator values of a lane through opaque moves with side
    effects. An unrolled GEMV loop pins its accumulators after each round.
    The round's FMAs must then finish before the moves, and the next round's
    LDS reads (and the MFMAs that need them) cannot move above them. Without
    it LLVM runs every MFMA of the loop first and the scaling FMAs at the end,
    which needs a register for each MFMA result and spills them all."""
    return fx.Vector.from_elements(
        [
            fx.Float32(
                llvm.inline_asm(
                    T.f32,
                    [fx.Float32(acc[i]).ir_value()],
                    "v_mov_b32 $0, $1",
                    "=v,v",
                    has_side_effects=True,
                )
            )
            for i in range(4)
        ],
        fx.Float32,
    )