Skip to content

vllm.models.minimax_m3.amd.mono.drafter.attention

Attention of the DSpark drafter kernels.

Split-KV partitions with the q / k head norms and RoPE fused in, and merge tasks (a row, 4 heads) that quantize the output to MXFP8 for o_proj. The semantics are vLLM's draft attention: non-causal within a request's 7 rows, keys below seq_len, sliding window 4096, FP8 q.k / sqrt(128), P in e4m3, FP8 PV.

LDS tiles are 128 rows x 128 bytes (4096 words); 16-byte chunks are XOR-swizzled by row % 8 against bank conflicts:

  • TQ: q rows m = 16 j + head (j < 7);
  • TK0 / TK1: the K rows (keys) of a block;
  • TV0 / TV1: V transposed, a row per dim with its 128 keys in the PV operand's order (byte 32 g + 4 i + e holds key 16 i + 4 g + e).

Block math, one wave a query row (its 16 heads), after AITER's gfx950 scheme: S^T = K Q^T puts keys 16 i + 4 g .. + 3 of a head in one lane, which is the PV operand at 4 keys a word, so P goes from the QK accumulator to O^T = V^T P^T in registers. exp2 folds log2(e) / sqrt(128) and vLLM's P x 256 (as +8 in the exponent): l and o carry the factor 256, which o / l cancels.

Functions:

  • block_acc –

    Fold one 128-key block into the running (m, l, o) of an online softmax.

  • block_math –

    One 128-key block for query row j (this wave), head l16.

  • kv_load –

    Block blk of request r (keys outside [lo, hi] zero), issued at the layer's start:

  • kv_store –

    The thread's K words -> tile tk (row-major); its V words, transposed in

  • perm –

    v_perm_b32: byte k of the result = byte sel[k] of (a << 32 | b).

  • q_operand –

    Wave wave's Q^T operand (query row wave, head l16).

  • qkv_pre –

    The loads of qkv_rows that do not depend on this layer's QKV.

  • tw –

    Swizzled word index of word w of row r of tile t.

  • window –

    Request r: first / last query position, its window union's first / last key and

Attributes:

  • merge_task –

    Merge row t, heads 4 hq .. 4 hq + 3 (a thread a head and dim).

  • qkv_rows –

    Request r's q (7 rows x 16 heads: head norm, RoPE, FP8) -> TQ, and its new

  • store_partial –

    (m, l, o) of query rows 16 wave + l16 (< 112) -> the partial at ap:

merge_task = traced(merge_task) module-attribute

Merge row t, heads 4 hq .. 4 hq + 3 (a thread a head and dim).

The request's partials 0..na-1 at ap0 are merged in order, rounded to bf16 and quantized to MXFP8 into a8 / as8. The partials' (m, l) of the 4 heads are staged through LDS ml (512 floats). maxp: static bound of na.

qkv_rows = traced(qkv_rows) module-attribute

Request r's q (7 rows x 16 heads: head norm, RoPE, FP8) -> TQ, and its new K / V rows -> the cache and, where they fall in the CTA's first blocks bA (K0 / V0, nfirst >= 1) or bB (K1 / V1, nfirst >= 2), the tiles.

store_partial = traced(store_partial) module-attribute

(m, l, o) of query rows 16 wave + l16 (< 112) -> the partial at ap: m[128] f32 at 0, l at 512 B, o[112][128] at 1 KB.

_lds8c(lds, t, r, w0)

Words w0 .. w0 + 7 (two consecutive 16-byte chunks) of row r of tile t.

Source code in vllm/models/minimax_m3/amd/mono/drafter/attention.py 
230
231
232
233
234
235
236
237
def _lds8c(lds, t, r, w0):
    """Words w0 .. w0 + 7 (two consecutive 16-byte chunks) of row r of tile t."""
    v4 = fx.Vector.make_type(4, fx.Int32)
    x = fx.Vector(fx.ptr_load(lds + tw(t, r, w0), result_type=v4))
    y = fx.Vector(fx.ptr_load(lds + tw(t, r, w0 + 4), result_type=v4))
    return fx.Vector.from_elements(
        [x[e] for e in range(4)] + [y[e] for e in range(4)], fx.Int32
    )

_lds8w(lds, t, r, w0)

Words w0 .. w0 + 3 and w0 + 16 .. w0 + 19 of row r of tile t: an operand.

Source code in vllm/models/minimax_m3/amd/mono/drafter/attention.py 
220
221
222
223
224
225
226
227
def _lds8w(lds, t, r, w0):
    """Words w0 .. w0 + 3 and w0 + 16 .. w0 + 19 of row r of tile t: an operand."""
    v4 = fx.Vector.make_type(4, fx.Int32)
    x = fx.Vector(fx.ptr_load(lds + tw(t, r, w0), result_type=v4))
    y = fx.Vector(fx.ptr_load(lds + tw(t, r, w0 + 16), result_type=v4))
    return fx.Vector.from_elements(
        [x[e] for e in range(4)] + [y[e] for e in range(4)], fx.Int32
    )

block_acc(lds, qb, tk, tv, blk, j, p0, lo, hi, l16, g4, m_run, l_run, o_run, mask)

Fold one 128-key block into the running (m, l, o) of an online softmax.

m is in raw score units, l and o are scaled by 256 at m: m' = max(m, block max), l and o are rescaled by exp(m - m'), and P (x 256 at m', e4m3) V accumulates into o.

Source code in vllm/models/minimax_m3/amd/mono/drafter/attention.py 
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
def block_acc(lds, qb, tk, tv, blk, j, p0, lo, hi, l16, g4, m_run, l_run, o_run, mask):
    """Fold one 128-key block into the running (m, l, o) of an online softmax.

    m is in raw score units, l and o are scaled by 256 at m: m' = max(m, block max),
    l and o are rescaled by exp(m - m'), and P (x 256 at m', e4m3) V accumulates
    into o.
    """
    s = []
    for i in range_constexpr(8):
        c = _mma_unit(_lds8w(lds, tk, 16 * i + l16, g4 * 4), qb, dev.zero4())
        s += [c[e] for e in range(4)]
    if const_expr(mask):
        qp = p0 + j
        for i in range_constexpr(8):
            for e in range_constexpr(4):
                kp = blk * layout.BLOCK + 16 * i + 4 * g4 + e
                ok = (kp > qp - layout.WINDOW) & (kp >= lo) & (kp <= hi)
                s[4 * i + e] = ok.select(s[4 * i + e], F32(NEG_INF))
    mb = s[0]
    for x in range_constexpr(1, 32):
        mb = fx.max(mb, s[x])
    mb = butterfly(mb, (16, 32), fx.max)
    mn = fx.max(m_run, mb)
    msafe = (mn == F32(NEG_INF)).select(F32(0.0), mn)
    alpha = hw_exp2((m_run - msafe) * QK_C)
    off = msafe * QK_C - 8.0
    p = [hw_exp2(s[x] * QK_C - off) for x in range(32)]
    ps = p[0]
    for x in range_constexpr(1, 32):
        ps = ps + p[x]
    l_new = l_run * alpha + butterfly(ps, (16, 32))
    pw = fx.Vector.from_elements(
        [
            fp8_pack4(p[4 * i], p[4 * i + 1], p[4 * i + 2], p[4 * i + 3])
            for i in range(8)
        ],
        fx.Int32,
    )
    o = []
    for n in range_constexpr(8):
        acc = fx.Vector.from_elements(
            [o_run[4 * n + e] * alpha for e in range(4)], fx.Float32
        )
        c = _mma_unit(_lds8c(lds, tv, 16 * n + l16, g4 * 8), pw, acc)
        o += [c[e] for e in range(4)]
    return mn, l_new, o

block_math(lds, qb, tk, tv, blk, j, p0, lo, hi, l16, g4, mask)

One 128-key block for query row j (this wave), head l16.

Returns (m in score units, l, o: 32 f32 of dims 16 n + 4 g + e); l and o are scaled by 256. mask (static): the block is at the window's edge.

Source code in vllm/models/minimax_m3/amd/mono/drafter/attention.py 
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
def block_math(lds, qb, tk, tv, blk, j, p0, lo, hi, l16, g4, mask):
    """One 128-key block for query row ``j`` (this wave), head ``l16``.

    Returns (m in score units, l, o: 32 f32 of dims 16 n + 4 g + e); l and o are
    scaled by 256. ``mask`` (static): the block is at the window's edge.
    """
    s = []
    for i in range_constexpr(8):
        ka = _lds8w(lds, tk, 16 * i + l16, g4 * 4)
        c = _mma_unit(ka, qb, dev.zero4())
        s += [c[e] for e in range(4)]
    if const_expr(mask):
        qp = p0 + j
        for i in range_constexpr(8):
            for e in range_constexpr(4):
                kp = blk * layout.BLOCK + 16 * i + 4 * g4 + e
                ok = (kp > qp - layout.WINDOW) & (kp >= lo) & (kp <= hi)
                s[4 * i + e] = ok.select(s[4 * i + e], F32(NEG_INF))
    m = s[0]
    for x in range_constexpr(1, 32):
        m = fx.max(m, s[x])
    m = butterfly(m, (16, 32), fx.max)
    msafe = m
    if const_expr(mask):
        msafe = (m == F32(NEG_INF)).select(F32(0.0), m)
    off = msafe * QK_C - 8.0
    p = [hw_exp2(s[x] * QK_C - off) for x in range(32)]
    lsum = p[0]
    for x in range_constexpr(1, 32):
        lsum = lsum + p[x]
    lsum = butterfly(lsum, (16, 32))
    pw = fx.Vector.from_elements(
        [
            fp8_pack4(p[4 * i], p[4 * i + 1], p[4 * i + 2], p[4 * i + 3])
            for i in range(8)
        ],
        fx.Int32,
    )
    o = []
    for n in range_constexpr(8):
        va = _lds8c(lds, tv, 16 * n + l16, g4 * 8)
        c = _mma_unit(va, pw, dev.zero4())
        o += [c[e] for e in range(4)]
    return m * SM_SCALE, lsum, o

kv_load(cache, block_table, bt_width, r, blk, lo, hi, tid)

Block blk of request r (keys outside [lo, hi] zero), issued at the layer's start: this thread's K words (key tid / 4, bytes 32 (tid % 4) ..: 8 words) and V words (keys 4 (tid / 16) .. + 3, dims 8 (tid % 16) ..: 4 x 2 words).

Source code in vllm/models/minimax_m3/amd/mono/drafter/attention.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def kv_load(cache, block_table, bt_width, r, blk, lo, hi, tid):
    """Block blk of request r (keys outside [lo, hi] zero), issued at the layer's start:
    this thread's K words (key tid / 4, bytes 32 (tid % 4) ..: 8 words) and V words
    (keys 4 (tid / 16) .. + 3, dims 8 (tid % 16) ..: 4 x 2 words)."""
    cache, nblk = cache
    # padded requests of a FULL graph replay carry stale metadata: keep every
    # address inside the block table row and the cache
    col = fx.min(fx.max(fx.min(blk, hi // layout.BLOCK), I32(0)), bt_width - 1)
    phys = fx.min(fx.max(dev.ld(block_table, r * bt_width + col), I32(0)), nblk - 1)
    base = cache + I64(phys) * (layout.BLOCK * 256)
    km = tid // 4
    kq = tid % 4
    kp = blk * layout.BLOCK + km
    klive = (kp >= lo) & (kp <= hi)
    k0 = dev.ld(base, km * 64 + kq * 8, 4)
    k1 = dev.ld(base, km * 64 + kq * 8 + 4, 4)
    kw = [klive.select(k0[e], I32(0)) for e in range(4)] + [
        klive.select(k1[e], I32(0)) for e in range(4)
    ]
    vg = tid // 16
    vd = tid % 16
    vw = []
    for j in range_constexpr(4):
        m = vg * 4 + j
        vp = blk * layout.BLOCK + m
        vlive = (vp >= lo) & (vp <= hi)
        v = dev.ld(base, m * 64 + 32 + vd * 2, 2)
        vw.append([vlive.select(v[0], I32(0)), vlive.select(v[1], I32(0))])
    return kw, vw

kv_store(lds, tk, tv, kw, vw, tid)

The thread's K words -> tile tk (row-major); its V words, transposed in registers, -> tile tv.

Source code in vllm/models/minimax_m3/amd/mono/drafter/attention.py 
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
def kv_store(lds, tk, tv, kw, vw, tid):
    """The thread's K words -> tile ``tk`` (row-major); its V words, transposed in
    registers, -> tile ``tv``."""
    km = tid // 4
    kq = tid % 4
    for h in range_constexpr(2):
        w0 = kq * 8 + h * 4
        fx.ptr_store(
            fx.Vector.from_elements(kw[4 * h : 4 * h + 4], fx.Int32),
            lds + tw(tk, km, w0),
        )
    # keys 4 vg .. 4 vg + 3 = 16 i + 4 g ..: word 8 g + i of each dim row
    vg = tid // 16
    vd = tid % 16
    vwi = (vg % 4) * 8 + vg // 4
    for h in range_constexpr(2):  # dims 8 vd + 4 h .. + 3
        a = [vw[j][h] for j in range(4)]
        t01l = perm(a[1], a[0], 0x05010400)
        t01h = perm(a[1], a[0], 0x07030602)
        t23l = perm(a[3], a[2], 0x05010400)
        t23h = perm(a[3], a[2], 0x07030602)
        outs = [
            perm(t23l, t01l, 0x05040100),
            perm(t23l, t01l, 0x07060302),
            perm(t23h, t01h, 0x05040100),
            perm(t23h, t01h, 0x07060302),
        ]
        for x in range_constexpr(4):
            d = vd * 8 + h * 4 + x
            fx.ptr_store(outs[x], lds + tw(tv, d, vwi))

perm(a, b, sel)

v_perm_b32: byte k of the result = byte sel[k] of (a << 32 | b).

Source code in vllm/models/minimax_m3/amd/mono/drafter/attention.py 
56
57
58
59
60
61
62
63
64
65
66
def perm(a, b, sel):
    """v_perm_b32: byte k of the result = byte sel[k] of (a << 32 | b)."""
    return I32(
        llvm.call_intrinsic(
            T.i32,
            "llvm.amdgcn.perm",
            [I32(a).ir_value(), I32(b).ir_value(), I32(sel).ir_value()],
            [],
            [],
        )
    )

q_operand(lds, wave, l16, g4)

Wave wave's Q^T operand (query row wave, head l16).

Source code in vllm/models/minimax_m3/amd/mono/drafter/attention.py 
240
241
242
def q_operand(lds, wave, l16, g4):
    """Wave ``wave``'s Q^T operand (query row ``wave``, head ``l16``)."""
    return _lds8w(lds, TQ, wave * 16 + l16, g4 * 4)

qkv_pre(r, qn, kn, cos_sin, positions, slots, tid)

The loads of qkv_rows that do not depend on this layer's QKV.

Source code in vllm/models/minimax_m3/amd/mono/drafter/attention.py 
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
def qkv_pre(r, qn, kn, cos_sin, positions, slots, tid):
    """The loads of ``qkv_rows`` that do not depend on this layer's QKV."""
    li, rope, dA, dB = _dims(tid)
    grp = tid // 16
    ib = rope.select(4 * li, I32(0))
    gq = [
        dev.ld(qn, dA // 2),
        dev.ld(qn, dA // 2 + 1),
        dev.ld(qn, dB // 2),
        dev.ld(qn, dB // 2 + 1),
    ]
    gk = [
        dev.ld(kn, dA // 2),
        dev.ld(kn, dA // 2 + 1),
        dev.ld(kn, dB // 2),
        dev.ld(kn, dB // 2 + 1),
    ]
    css, sls, poss = [], [], []
    for s in range_constexpr(4):
        t = grp + 32 * s
        isq, j, col, row = _slot_rows(r, t)
        pos = ipos(positions, row)
        c0 = (pos * layout.ROT + ib) // 2
        s0 = (pos * layout.ROT + layout.ROT // 2 + ib) // 2
        css.append(
            [
                dev.ld(cos_sin, c0),
                dev.ld(cos_sin, c0 + 1),
                dev.ld(cos_sin, s0),
                dev.ld(cos_sin, s0 + 1),
            ]
        )
        sls.append(I64(dev.ld(slots, row * 2)))
        poss.append(pos)
    return gq, gk, css, sls, poss

tw(t, r, w)

Swizzled word index of word w of row r of tile t.

Source code in vllm/models/minimax_m3/amd/mono/drafter/attention.py 
51
52
53
def tw(t, r, w):
    """Swizzled word index of word w of row r of tile t."""
    return t * TILE_WORDS + r * 32 + ((w // 4) ^ (r % 8)) * 4 + w % 4

window(r, positions, seq_lens)

Request r: first / last query position, its window union's first / last key and blocks (keys < seq_len, from the first row's window start, <= MAXB blocks).

Source code in vllm/models/minimax_m3/amd/mono/drafter/attention.py 
73
74
75
76
77
78
79
80
81
def window(r, positions, seq_lens):
    """Request r: first / last query position, its window union's first / last key and
    blocks (keys < seq_len, from the first row's window start, <= MAXB blocks)."""
    p0 = ipos(positions, r * layout.QPR)
    p6 = ipos(positions, r * layout.QPR + layout.QPR - 1)
    seq = dev.ld(seq_lens, r)
    lo = fx.max(p0 - (layout.WINDOW - 1), I32(0))
    hi = fx.min(seq - 1, (lo // layout.BLOCK + layout.MAXB) * layout.BLOCK - 1)
    return p0, p6, lo, hi, lo // layout.BLOCK, hi // layout.BLOCK