Skip to content

vllm.models.deepseek_v41.amd.mono.attention.gemv

MXFP8 GEMV of the mega kernel: 16 rows of a row-major e4m3 weight [N, K] (vLLM's loaded layout, E8M0 scales in 32x32 blocks [N / 32, K / 32]) against every token of a tile, its MXFP8 rows in LDS, on the gfx950 scaled MFMA.

A task may cover a window of K (kbeg .. kbeg + nsteps 128-steps): a split-K part, its activations in LDS holding just that window. The weight operands are loaded first (all in flight, before the task waits on its activations); wave w takes the window's steps w, w + 8, ... (or, split, a contiguous range of waves 0 .. split - 1), its 16 x 16 partial in red, summed in wave order by row_sum. Adapted from ATOM's attn_pre.gemv_fp8_*, which reads aiter's (16, 16)-preshuffled layout.

Functions:

  • gemv_loads –

    The weight operands of rows row0 .. +16 (tiled: row0 + 16 (wave //

  • gemv_mfmas –

    gemv_loads' operands against rows tokens' MXFP8 rows of the

  • lds_at –

    Element u of a span of rows row wide, at the lds_row stride.

  • row_sum –

    Row r, token t of the waves' accumulators, summed in wave order.

  • rows_loads –

    Tokens t0 .. t0 + n's published MXFP8 row slices -- words i32 from

  • rows_store –

    rows_loads' registers into LDS rows 0 .. n (lds_row strides).

  • span_items –

    A [n, width] span of 16 B items: rows a round of THREADS threads, rounds.

_step_load(c, r_w, r_ws, k, row0, st, pred=None)

Global step st's weight operand of rows row0 .. +16: two 16 B chunks a lane (K bytes 16 j .. and 64 + 16 j .. of the step, row row0 + lane % 16) and the E8M0 byte of its K block. pred false: masked (out of range: no memory traffic, zeros).

Source code in vllm/models/deepseek_v41/amd/mono/attention/gemv.py 
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
def _step_load(c, r_w, r_ws, k, row0, st, pred=None):
    """Global step ``st``'s weight operand of rows row0 .. +16: two 16 B chunks a
    lane (K bytes 16 j .. and 64 + 16 j .. of the step, row row0 + lane % 16)
    and the E8M0 byte of its K block. ``pred`` false: masked (out of range: no
    memory traffic, zeros)."""
    lane = c["lane"]
    j = lane // 16
    row = row0 + lane % 16
    if const_expr(GFX942):
        return _step_load_942(r_w, r_ws, k, row0, lane, st)
    base = row * k + 128 * st + 16 * j
    lo = fx.Vector(bo.buffer_load(r_w, base // 4, vec_width=4, dtype=T.i32))
    hi = fx.Vector(bo.buffer_load(r_w, (base + 64) // 4, vec_width=4, dtype=T.i32))
    kg = k // 32
    wsi = (row0 // 32) * kg + st * 4 + j
    word = fx.Int32(bo.buffer_load(r_ws, wsi // 4, vec_width=1, dtype=T.i32))
    return lo, hi, (word >> (wsi % 4 * 8)) & 0xFF

_step_load_942(r_w, r_ws, k, row0, lane, st)

gfx942: step st of rows row0 .. + 16 of the runner's tile-major FP8 copy of the weight (weights942.dense_copy). Lane 16 j + r takes row r's four 8-byte chunks of lane group j (K 32 b + 8 j .. + 8 for blocks b = 0 .. 3) with two 16-byte loads, blocks 0 and 1, then blocks 2 and 3. Each load of the wave reads 1 KB that is contiguous. The scale codes are the FNUZ codes (e8m0 + 1) in 32 x 32 blocks. The step's four codes are one dword, as K / 32 is a multiple of 4.

Source code in vllm/models/deepseek_v41/amd/mono/attention/gemv.py 
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
def _step_load_942(r_w, r_ws, k, row0, lane, st):
    """gfx942: step ``st`` of rows row0 .. + 16 of the runner's tile-major FP8
    copy of the weight (``weights942.dense_copy``). Lane 16 j + r takes row
    r's four 8-byte chunks of lane group j (K 32 b + 8 j .. + 8 for blocks
    b = 0 .. 3) with two 16-byte loads, blocks 0 and 1, then blocks 2 and 3.
    Each load of the wave reads 1 KB that is contiguous. The scale codes are
    the FNUZ codes (e8m0 + 1) in 32 x 32 blocks. The step's four codes are
    one dword, as K / 32 is a multiple of 4."""
    # A step of the 16 rows is 2 KB, the first halves of the lanes and then
    # the second halves.
    base = (row0 * k + 2048 * st) // 4 + 4 * lane
    lo = fx.Vector(bo.buffer_load(r_w, base, vec_width=4, dtype=T.i32))
    hi = fx.Vector(bo.buffer_load(r_w, base + 256, vec_width=4, dtype=T.i32))
    wsi = (row0 // 32) * (k // 32) + st * 4
    word = fx.Int32(bo.buffer_load(r_ws, wsi // 4, vec_width=1, dtype=T.i32))
    return lo, hi, word

_step_mfma_942(xl, xsl, x_row, s_row, st, col, j, ops, acc, live)

gfx942: one K step as four FP8 MFMAs, one a 32-wide scale block, each added to acc times its block's scale factor. Block b takes the lane's weight chunk b (_step_load_942) and the activation bytes K 32 b + 8 j .. + 8 of its token, read from the natural LDS row. A step past the end (live false) gets the code sum 127, whose factor is 0.

Source code in vllm/models/deepseek_v41/amd/mono/attention/gemv.py 
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
def _step_mfma_942(xl, xsl, x_row, s_row, st, col, j, ops, acc, live):
    """gfx942: one K step as four FP8 MFMAs, one a 32-wide scale block, each
    added to ``acc`` times its block's scale factor. Block b takes the lane's
    weight chunk b (``_step_load_942``) and the activation bytes K 32 b + 8 j
    .. + 8 of its token, read from the natural LDS row. A step past the end
    (``live`` false) gets the code sum 127, whose factor is 0."""
    lo, hi, sw = ops
    xw = [
        fx.Vector(
            fx.ptr_load(
                xl + (col * x_row + st * 32 + 8 * b + 2 * j),
                result_type=fx.Vector.make_type(2, fx.Int32),
            )
        )
        for b in range(4)
    ]
    sx = fx.Vector(
        fx.ptr_load(
            xsl + (col * s_row + st * 4),
            result_type=fx.Vector.make_type(4, fx.Int32),
        )
    )
    for b in range_constexpr(4):
        w = lo if b < 2 else hi
        a = pair_i64(w[2 * (b % 2)], w[2 * (b % 2) + 1])
        code = ((sw >> (8 * b)) & 0xFF) + sx[b]
        if const_expr(live is not None):
            code = live.select(code, fx.Int32(127))
        acc = mfma_fp8_scaled(acc, a, pair_i64(xw[b][0], xw[b][1]), code)
    return pin_acc(acc)

_steps(c, nsteps, split, tiled)

This wave's steps of the window: [(local step, live)].

Source code in vllm/models/deepseek_v41/amd/mono/attention/gemv.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
def _steps(c, nsteps, split, tiled):
    """This wave's steps of the window: [(local step, live)]."""
    wave = c["wave"]
    if const_expr(split is None):
        per = (nsteps + WAVES - 1) // WAVES
        ragged = nsteps % WAVES != 0
        return [
            (wave + WAVES * i, (wave + WAVES * i < nsteps) if ragged else None)
            for i in range(per)
        ]
    assert split <= WAVES and nsteps % split == 0
    per = nsteps // split
    part = wave % split if tiled else fx.min(wave, split - 1)
    return [(part * per + i, None) for i in range(per)]

gemv_loads(c, w, ws, k, row0, split=None, tiled=False, kbeg=0, nsteps=None, pred=None)

The weight operands of rows row0 .. +16 (tiled: row0 + 16 (wave // split)) over steps kbeg .. kbeg + nsteps for this wave, all in flight (pred false: masked off, no memory traffic).

Source code in vllm/models/deepseek_v41/amd/mono/attention/gemv.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
def gemv_loads(
    c, w, ws, k, row0, split=None, tiled=False, kbeg=0, nsteps=None, pred=None
):
    """The weight operands of rows row0 .. +16 (``tiled``: row0 + 16 (wave //
    split)) over steps kbeg .. kbeg + nsteps for this wave, all in flight
    (``pred`` false: masked off, no memory traffic)."""
    nsteps = k // 128 if nsteps is None else nsteps
    if const_expr(pred is None):
        r_w, r_ws = rsrc(w), rsrc(ws)
    else:
        # a masked-off CTA reads through a zero-size buffer: zeros, no traffic
        nb = pred.select(fx.Int64(0xFFFFFFFF), fx.Int64(0))
        r_w, r_ws = rsrc(w, nb), rsrc(ws, nb)
    if const_expr(tiled):
        row0 = row0 + 16 * (c["wave"] // split)
    return [
        _step_load(
            c,
            r_w,
            r_ws,
            k,
            row0,
            kbeg + (st if live is None else fx.min(st, nsteps - 1)),
            pred,
        )
        for st, live in _steps(c, nsteps, split, tiled)
    ]

gemv_mfmas(c, k, xl, xsl, red, ops, split=None, tiled=False, rows=16, nsteps=None)

gemv_loads' operands against rows tokens' MXFP8 rows of the window in LDS (xl words, xsl codes, rows lds_row of the window) -> red[wave][lane] (a 16 x 16 C: row 4 (lane // 16) + i, token lane % 16).

Source code in vllm/models/deepseek_v41/amd/mono/attention/gemv.py 
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
def gemv_mfmas(c, k, xl, xsl, red, ops, split=None, tiled=False, rows=16, nsteps=None):
    """``gemv_loads``' operands against ``rows`` tokens' MXFP8 rows of the
    window in LDS (``xl`` words, ``xsl`` codes, rows ``lds_row`` of the window)
    -> ``red[wave][lane]`` (a 16 x 16 C: row 4 (lane // 16) + i, token lane % 16)."""
    lane, wave = c["lane"], c["wave"]
    nsteps = k // 128 if nsteps is None else nsteps
    acc = fx.Vector.filled(4, 0.0, fx.Float32)
    for i, (st, live) in enumerate(_steps(c, nsteps, split, tiled)):
        acc = _step_mfma(c, nsteps, xl, xsl, st, ops[i], acc, live, rows)
    if const_expr(split is not None and not tiled):
        acc = fx.Vector.from_elements(
            [(wave < split).select(acc[e], fx.Float32(0.0)) for e in range(4)],
            fx.Float32,
        )
    fx.ptr_store(acc, red + (wave * 64 + lane) * 4)

lds_at(u, row)

Element u of a span of rows row wide, at the lds_row stride.

Source code in vllm/models/deepseek_v41/amd/mono/attention/gemv.py 
35
36
37
def lds_at(u, row):
    """Element u of a span of rows ``row`` wide, at the ``lds_row`` stride."""
    return u + u // row * LDS_PAD

row_sum(red, r, t, waves=None)

Row r, token t of the waves' accumulators, summed in wave order.

Source code in vllm/models/deepseek_v41/amd/mono/attention/gemv.py 
210
211
212
213
214
215
def row_sum(red, r, t, waves=None):
    """Row r, token t of the waves' accumulators, summed in wave order."""
    tot = fx.Float32(0.0)
    for w in waves if waves is not None else range(WAVES):
        tot = tot + fx.ptr_load(red + ((w * 64 + 16 * (r // 4) + t) * 4 + r % 4))
    return tot

rows_loads(c, words_base, codes_base, stride_w, off_w, words, stride_g, off_g, groups, t0, n, pred=None)

Tokens t0 .. t0 + n's published MXFP8 row slices -- words i32 from off_w and groups codes (an i32 each) from off_g, rows stride_w / stride_g i32 apart -- loaded (issued, not waited on), 16 B a lane-item: a thread one column, rows a round apart. Plain loads: their flags are seen, and nothing in the launch read these lines before. pred false: masked off (zeros, no memory traffic).

Source code in vllm/models/deepseek_v41/amd/mono/attention/gemv.py 
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
def rows_loads(
    c,
    words_base,
    codes_base,
    stride_w,
    off_w,
    words,
    stride_g,
    off_g,
    groups,
    t0,
    n,
    pred=None,
):
    """Tokens t0 .. t0 + n's published MXFP8 row slices -- ``words`` i32 from
    ``off_w`` and ``groups`` codes (an i32 each) from ``off_g``, rows
    ``stride_w`` / ``stride_g`` i32 apart -- loaded (issued, not waited on),
    16 B a lane-item: a thread one column, rows a round apart. Plain loads:
    their flags are seen, and nothing in the launch read these lines before.
    ``pred`` false: masked off (zeros, no memory traffic)."""
    tid = c["tid"]
    nb = None if pred is None else pred.select(fx.Int64(0xFFFFFFFF), fx.Int64(0))
    out = []
    for base, stride, off, width in (
        (words_base, stride_w, off_w, words),
        (codes_base, stride_g, off_g, groups),
    ):
        wq, rpr, rounds = span_items(width, n)
        row0, col = tile_map(tid, wq)
        a0 = col * 4 + off + t0 * stride
        r = rsrc(base, nb)
        regs = []
        for i in range(rounds):
            row = row0 + rpr * i
            if (i + 1) * rpr > n or rpr * wq < THREADS:
                row = fx.min(row, n - 1)
            regs.append(
                fx.Vector(
                    bo.buffer_load(r, row * stride + a0, vec_width=4, dtype=T.i32)
                )
            )
        out.append(regs)
    return out

rows_store(c, regs, words, groups, xl, xsl, n)

rows_loads' registers into LDS rows 0 .. n (lds_row strides).

Source code in vllm/models/deepseek_v41/amd/mono/attention/gemv.py 
286
287
288
289
290
def rows_store(c, regs, words, groups, xl, xsl, n):
    """``rows_loads``' registers into LDS rows 0 .. n (``lds_row`` strides)."""
    _store_span(c, regs[0], words, xl, n)
    _store_span(c, regs[1], groups, xsl, n)
    gpu.barrier()

span_items(width, n)

A [n, width] span of 16 B items: rows a round of THREADS threads, rounds.

Source code in vllm/models/deepseek_v41/amd/mono/attention/gemv.py 
218
219
220
221
222
223
def span_items(width, n):
    """A [n, width] span of 16 B items: rows a round of THREADS threads, rounds."""
    wq = width // 4
    rpr = THREADS // wq
    assert rpr >= 1, width
    return wq, rpr, (n + rpr - 1) // rpr