Skip to content

vllm.models.deepseek_v41.amd.mono.common.mx

FP8 / MX arithmetic the mono kernels share: the E4M3 range, E8M0 scale codes, the byte of a scale in aiter's shuffle_scale layout, an FP4 weight tile's operand loads and the scaled MFMA that takes them.

An E8M0 code is a biased exponent; pow2(code) is the scale it stands for.

Functions:

  • clamp_fp8 –

    v clamped to the E4M3 range.

  • code_ceil –

    Biased exponent of the smallest power of two >= v (v >= 0, finite),

  • fp4_tile_load –

    A lane's 16 B of an FP4 weight's (16, 16)-shuffled 16-row group rg

  • mfma_scaled –

    One 16x16x128 scaled MFMA (gfx950 f8f6f4): A a (a_fmt) with E8M0

  • rcp_pow2 –

    1 / pow2(code), exact (code in [1, 253]): x times it is x / scale

  • scale_index –

    Byte of e8m0 scale (row, col) in aiter's shuffle_scale layout:

  • udiv –

    x // d by a constant d > 0 as a shift (d a power of two, any

  • umod –

    x % d as a mask or an unsigned remainder (udiv).

clamp_fp8(v)

v clamped to the E4M3 range.

Source code in vllm/models/deepseek_v41/amd/mono/common/mx.py 
28
29
30
def clamp_fp8(v):
    """``v`` clamped to the E4M3 range."""
    return fx.max(fx.min(v, fx.Float32(FP8_MAX)), fx.Float32(-FP8_MAX))

code_ceil(v)

Biased exponent of the smallest power of two >= v (v >= 0, finite), clamped to [1, 254]: the codes a scale may take.

Source code in vllm/models/deepseek_v41/amd/mono/common/mx.py 
33
34
35
36
37
38
39
40
def code_ceil(v):
    """Biased exponent of the smallest power of two >= ``v`` (v >= 0, finite),
    clamped to [1, 254]: the codes a scale may take."""
    bits = v.bitcast(fx.Int32)
    e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).select(
        fx.Int32(1), fx.Int32(0)
    )
    return fx.max(fx.min(e, fx.Int32(254)), fx.Int32(1))

fp4_tile_load(r_w, base, rg, k, st, at, live=None)

A lane's 16 B of an FP4 weight's (16, 16)-shuffled 16-row group rg (K = k) at 128-k step st: dword base + the group's 2 k + 256 a step + at (lane * 4: 32 fp4, K block lane / 16 of row lane % 16). Nontemporal: a weight a step reads once. live: oob.

Source code in vllm/models/deepseek_v41/amd/mono/common/mx.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def fp4_tile_load(r_w, base, rg, k, st, at, live=None):
    """A lane's 16 B of an FP4 weight's (16, 16)-shuffled 16-row group ``rg``
    (K = ``k``) at 128-k step ``st``: dword ``base`` + the group's 2 k + 256 a
    step + ``at`` (``lane * 4``: 32 fp4, K block lane / 16 of row lane % 16).
    Nontemporal: a weight a step reads once. ``live``: ``oob``."""
    return fx.Vector(
        bo.buffer_load(
            r_w,
            oob(live, base + rg * (2 * k) + st * 256 + at),
            vec_width=4,
            dtype=T.i32,
            cache_modifier=CM_NT,
        )
    )

mfma_scaled(a, b, c, sa, sb, a_fmt=FP8, b_fmt=FP8, sa_byte=0, sb_byte=0)

One 16x16x128 scaled MFMA (gfx950 f8f6f4): A a (a_fmt) with E8M0 sa (its byte sa_byte), B b (b_fmt) with E8M0 sb (its byte sb_byte).

Source code in vllm/models/deepseek_v41/amd/mono/common/mx.py 
108
109
110
111
112
113
114
115
116
def mfma_scaled(a, b, c, sa, sb, a_fmt=FP8, b_fmt=FP8, sa_byte=0, sb_byte=0):
    """One 16x16x128 scaled MFMA (gfx950 f8f6f4): A ``a`` (``a_fmt``) with E8M0
    ``sa`` (its byte ``sa_byte``), B ``b`` (``b_fmt``) with E8M0 ``sb`` (its
    byte ``sb_byte``)."""
    return fx.Vector(
        rocdl.mfma_scale_f32_16x16x128_f8f6f4(
            T.vec(4, T.f32), [a, b, c, a_fmt, b_fmt, sa_byte, sa, sb_byte, sb]
        )
    )

rcp_pow2(code)

1 / pow2(code), exact (code in [1, 253]): x times it is x / scale bit for bit, a multiply where the division is a dozen instructions.

Source code in vllm/models/deepseek_v41/amd/mono/common/mx.py 
43
44
45
46
def rcp_pow2(code):
    """1 / ``pow2(code)``, exact (code in [1, 253]): x times it is x / scale
    bit for bit, a multiply where the division is a dozen instructions."""
    return ((254 - code) << 23).bitcast(fx.Float32)

scale_index(row, col, cols, sm=False)

Byte of e8m0 scale (row, col) in aiter's shuffle_scale layout: view(rows/32, 2, 16, cols/8, 2, 4).permute(0, 3, 5, 2, 4, 1) (sm: in shifts and masks, udiv).

Source code in vllm/models/deepseek_v41/amd/mono/common/mx.py 
49
50
51
52
53
54
55
def scale_index(row, col, cols, sm=False):
    """Byte of e8m0 scale (row, col) in aiter's ``shuffle_scale`` layout:
    ``view(rows/32, 2, 16, cols/8, 2, 4).permute(0, 3, 5, 2, 4, 1)`` (``sm``:
    in shifts and masks, ``udiv``)."""
    r32, a, b = udiv(row, 32, sm), umod(udiv(row, 16, sm), 2, sm), umod(row, 16, sm)
    c8, d, e = udiv(col, 8, sm), umod(udiv(col, 4, sm), 2, sm), umod(col, 4, sm)
    return ((((r32 * (cols // 8) + c8) * 4 + e) * 16 + b) * 2 + d) * 2 + a

udiv(x, d, unsigned=True)

x // d by a constant d > 0 as a shift (d a power of two, any sign) or an unsigned division (x >= 0), where a traced // is a signed floor division with its fix-ups. unsigned false: //.

Source code in vllm/models/deepseek_v41/amd/mono/common/mx.py 
58
59
60
61
62
63
64
65
66
def udiv(x, d, unsigned=True):
    """``x // d`` by a constant ``d > 0`` as a shift (``d`` a power of two, any
    sign) or an unsigned division (``x >= 0``), where a traced ``//`` is a
    signed floor division with its fix-ups. ``unsigned`` false: ``//``."""
    if isinstance(x, int) or not unsigned:
        return x // d
    if d & (d - 1) == 0:
        return x >> (d.bit_length() - 1)
    return fx.Int32(fx.Uint32(x) // d)

umod(x, d, unsigned=True)

x % d as a mask or an unsigned remainder (udiv).

Source code in vllm/models/deepseek_v41/amd/mono/common/mx.py 
69
70
71
72
73
74
75
def umod(x, d, unsigned=True):
    """``x % d`` as a mask or an unsigned remainder (``udiv``)."""
    if isinstance(x, int) or not unsigned:
        return x % d
    if d & (d - 1) == 0:
        return x & (d - 1)
    return fx.Int32(fx.Uint32(x) % d)