Skip to content

vllm.models.deepseek_v41.amd.mono.common.ops

Device-side primitives of the mono kernels: wave reductions and shuffles, the approximate hardware math, bf16 / fp8 packing, buffer resources, the execution-model constants and the kernel symbol naming.

Ported from the FlyDSL GLM-5 layer kernel (ROCm/FlyDSL branch glm5-mla-moe-layer-monokernel @ 1e28fd8a, kernels/mla_moe_layer/shared_reuse_moe_kernel.py). Only published flydsl APIs and aiter's vendored buffer_ops / dpp_utils are used.

Every function but kernel_symbol is traced inside a @flyc.kernel body.

Raw boundaries kept, for want of an equivalent on the pinned FlyDSL: buffer_ops loads / stores (the kernels take raw Int64 addresses, no tensors for make_buffer_tensor); raw MFMA intrinsics (hand-laid operands: permlane transposes, fp8 packing, MXFP4 scales); the LLVM intrinsics here (expect, s_memrealtime, the approximate rsq / rcp / exp2) and the struct unpacking of permlane*_swap.

Also the TP group seen from inside a kernel (every rank's peer buffer, an all-reduce's rank-ordered sum) and a timeline build's phase stamps.

Functions:

  • ballot –

    The wave's lanes where pred holds, a bit a lane.

  • batched_rounds –

    A CTA's strided loop over elements v = tid + THREADS i < n (n a Python

  • bf16_pair –

    Two f32 -> one f32-typed word holding (bf16(a), bf16(b)).

  • bf2_f32 –

    Packed bf16 pair word -> (f32 low, f32 high).

  • bf_hi –

    The f32 of a packed bf16 pair word's high half.

  • bf_lo –

    The f32 of a packed bf16 pair word's low half.

  • butterfly –

    Butterfly reduction over the given xor offsets, in the given order.

  • div_rn –

    x / d rounded to nearest, bit-identical to the IEEE division, given

  • fp8_pack4 –

    Four f32 (already in range) -> one dword of four E4M3 bytes, RNE.

  • fresh –

    v (a per-lane Int32) through an opaque move, here: the lane math built

  • hw_rsq –

    v_rsq_f32 (~1 ulp, no libm range fixups).

  • lane_gather –

    Lane src_lane's value of v, a source lane a lane (ds_bpermute).

  • lanes_below –

    How many of the lanes in ballot bal are below lane.

  • load_ptr64 –

    Entry i of an i64 table in memory, wave-uniform (each half through

  • memrealtime –

    s_memrealtime: the 100 MHz device clock (no FlyDSL wrapper).

  • peer_bases –

    Every rank's peer buffer base from the peers table (an i64 a rank),

  • permlane_swap –

    v_permlane{32,16}_swap on registers (x, y): off 32 trades x's upper-half

  • row_sum –

    Row r, token t of the waves' 16x16 MFMA accumulators (wave w's at red

  • sum_partials –

    The tp ranks' bf16 pair partials at pair_of(src) of the mailbox

  • traced –

    Apply flydsl's AST rewriting to a helper, so its while / if over

  • xred –

    op(v, lane ^ off) for a symmetric op; 32 / 16 take both halves of one swap.

  • xshfl –

    Value of lane lane ^ off: offsets 32 / 16 via v_permlane*_swap, in-row

ballot(pred)

The wave's lanes where pred holds, a bit a lane.

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
220
221
222
def ballot(pred):
    """The wave's lanes where ``pred`` holds, a bit a lane."""
    return fx.Int64(rocdl.ballot(T.i64, as_ir_value(pred)))

batched_rounds(tid, n, load, use, batch=LOAD_BATCH)

A CTA's strided loop over elements v = tid + THREADS i < n (n a Python int): load(v) of batch rounds issued before any use(v, loaded). A round whose use stores to global memory keeps the next round's loads behind it (they may alias), so the rounds' loads would go out in series. A round past n loads element n - 1 (in range) and uses nothing.

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
@traced
def batched_rounds(tid, n, load, use, batch=LOAD_BATCH):
    """A CTA's strided loop over elements v = tid + THREADS i < n (n a Python
    int): ``load(v)`` of ``batch`` rounds issued before any ``use(v, loaded)``.
    A round whose ``use`` stores to global memory keeps the next round's loads
    behind it (they may alias), so the rounds' loads would go out in series. A
    round past n loads element n - 1 (in range) and uses nothing."""
    rounds = -(-n // THREADS)
    for b0 in range_constexpr(0, rounds, batch):
        idx = range(b0, min(b0 + batch, rounds))
        got = [load(fx.min(tid + THREADS * i, n - 1)) for i in idx]
        for i, g in zip(idx, got):
            v = tid + THREADS * i
            if v < n:
                use(v, g)

bf16_pair(a, b)

Two f32 -> one f32-typed word holding (bf16(a), bf16(b)).

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
245
246
247
248
249
250
251
def bf16_pair(a, b):
    """Two f32 -> one f32-typed word holding (bf16(a), bf16(b))."""
    return (
        fx.Vector.from_elements([a, b], fx.Float32)
        .to(fx.BFloat16)
        .bitcast(fx.Float32)[0]
    )

bf2_f32(w)

Packed bf16 pair word -> (f32 low, f32 high).

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
264
265
266
def bf2_f32(w):
    """Packed bf16 pair word -> (f32 low, f32 high)."""
    return bf_lo(w), bf_hi(w)

bf_hi(w)

The f32 of a packed bf16 pair word's high half.

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
259
260
261
def bf_hi(w):
    """The f32 of a packed bf16 pair word's high half."""
    return (w & fx.Int32(-65536)).bitcast(fx.Float32)

bf_lo(w)

The f32 of a packed bf16 pair word's low half.

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
254
255
256
def bf_lo(w):
    """The f32 of a packed bf16 pair word's low half."""
    return (w << 16).bitcast(fx.Float32)

butterfly(v, offsets, op=_add)

Butterfly reduction over the given xor offsets, in the given order.

The order is the numerical contract: the aiter per-head norm reduces 32 lanes with offsets 16, 8, 4, 2, 1, and matching it keeps rstd bit-exact.

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
170
171
172
173
174
175
176
177
def butterfly(v, offsets, op=_add):
    """Butterfly reduction over the given xor offsets, in the given order.

    The order is the numerical contract: the aiter per-head norm reduces 32
    lanes with offsets 16, 8, 4, 2, 1, and matching it keeps rstd bit-exact."""
    for off in offsets:
        v = xred(v, off, op)
    return v

div_rn(x, d, r)

x / d rounded to nearest, bit-identical to the IEEE division, given r = 1.0 / d from one true division per divisor (Markstein: a quotient within 1 ulp plus its exact fma residual times RN(1/d) rounds correctly).

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
124
125
126
127
128
129
def div_rn(x, d, r):
    """``x / d`` rounded to nearest, bit-identical to the IEEE division, given
    ``r = 1.0 / d`` from one true division per divisor (Markstein: a quotient
    within 1 ulp plus its exact fma residual times RN(1/d) rounds correctly)."""
    q = x * r
    return fx.Float32(fmath.fma(fx.Float32(fmath.fma(-q, d, x)), r, q))

fp8_pack4(a, b, c, d)

Four f32 (already in range) -> one dword of four E4M3 bytes, RNE.

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
269
270
271
272
def fp8_pack4(a, b, c, d):
    """Four f32 (already in range) -> one dword of four E4M3 bytes, RNE."""
    lo = fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, a, b, fx.Int32(0), False))
    return fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, c, d, lo, True))

fresh(v)

v (a per-lane Int32) through an opaque move, here: the lane math built on it is computed at its use, where LLVM would otherwise hoist it out of a task loop or CSE it across stages into registers held for the whole kernel (spills, and a kernel with scratch is refused: runtime.widths).

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def fresh(v):
    """``v`` (a per-lane Int32) through an opaque move, here: the lane math built
    on it is computed at its use, where LLVM would otherwise hoist it out of a task
    loop or CSE it across stages into registers held for the whole kernel (spills,
    and a kernel with scratch is refused: ``runtime.widths``)."""
    return fx.Int32(
        llvm.inline_asm(
            T.i32,
            [fx.Int32(v).ir_value()],
            "v_mov_b32 $0, $1",
            "=v,v",
            has_side_effects=True,
        )
    )

hw_rsq(x)

v_rsq_f32 (~1 ulp, no libm range fixups).

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
111
112
113
def hw_rsq(x):
    """v_rsq_f32 (~1 ulp, no libm range fixups)."""
    return _hw_f32("llvm.amdgcn.rsq.f32", x)

lane_gather(v, src_lane)

Lane src_lane's value of v, a source lane a lane (ds_bpermute).

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
73
74
75
76
77
def lane_gather(v, src_lane):
    """Lane ``src_lane``'s value of v, a source lane a lane (ds_bpermute)."""
    return fx.Int32(
        rocdl.ds_bpermute(T.i32, fx.Int32(src_lane) * 4, fx.Int32(v).ir_value())
    )

lanes_below(bal, lane)

How many of the lanes in ballot bal are below lane.

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
229
230
231
def lanes_below(bal, lane):
    """How many of the lanes in ballot ``bal`` are below ``lane``."""
    return popcount(bal & ((fx.Int64(1) << fx.Int64(lane)) - 1))

load_ptr64(table, i)

Entry i of an i64 table in memory, wave-uniform (each half through uniform): a peer's buffer base, a kernel argument block's pointer.

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
234
235
236
237
238
def load_ptr64(table, i):
    """Entry ``i`` of an i64 table in memory, wave-uniform (each half through
    ``uniform``): a peer's buffer base, a kernel argument block's pointer."""
    pw = fx.Vector(bo.buffer_load(rsrc(table), i * 2, vec_width=2, dtype=T.i32))
    return (fx.Int64(uniform(pw[1])) << 32) | fx.Int64(fx.Uint32(uniform(pw[0])))

memrealtime()

s_memrealtime: the 100 MHz device clock (no FlyDSL wrapper).

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
80
81
82
def memrealtime():
    """s_memrealtime: the 100 MHz device clock (no FlyDSL wrapper)."""
    return fx.Int64(llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], []))

peer_bases(peers, tp)

Every rank's peer buffer base from the peers table (an i64 a rank), loaded once: a put's address must not wait on a load of its own.

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
276
277
278
279
def peer_bases(peers, tp):
    """Every rank's peer buffer base from the ``peers`` table (an i64 a rank),
    loaded once: a put's address must not wait on a load of its own."""
    return [load_ptr64(peers, p) for p in range(tp)]

permlane_swap(off, x, y)

v_permlane{32,16}_swap on registers (x, y): off 32 trades x's upper-half lanes with y's lower-half lanes; off 16 trades x's odd 16-lane rows with y's even rows (each within its half). -> (x', y').

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
def permlane_swap(off, x, y):
    """v_permlane{32,16}_swap on registers (x, y): off 32 trades x's upper-half
    lanes with y's lower-half lanes; off 16 trades x's odd 16-lane rows with y's
    even rows (each within its half). -> (x', y')."""
    swap = rocdl.permlane32_swap if off == 32 else rocdl.permlane16_swap
    pr = swap(
        llvm.StructType.get_literal([T.i32, T.i32]),
        as_ir_value(fx.Int32(x)),
        as_ir_value(fx.Int32(y)),
        False,
        False,
    )
    return tuple(fx.Int32(llvm.extractvalue(T.i32, pr, [j])) for j in range(2))

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

Row r, token t of the waves' 16x16 MFMA accumulators (wave w's at red + (w * 64 + lane) * 4), summed in wave order from 0 (waves: which, in order; all of them by default).

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
210
211
212
213
214
215
216
217
def row_sum(red, r, t, waves=None):
    """Row r, token t of the waves' 16x16 MFMA accumulators (wave w's at ``red``
    + (w * 64 + lane) * 4), summed in wave order from 0 (``waves``: which, in
    order; all of them by default)."""
    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

sum_partials(poll, own, pair_of, tp)

The tp ranks' bf16 pair partials at pair_of(src) of the mailbox region own, summed in fp32 in rank order 0 .. tp - 1: one poll batch -> (lo sum, hi sum). Every rank sums in the same order, so they agree.

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
282
283
284
285
286
287
288
289
290
291
292
def sum_partials(poll, own, pair_of, tp):
    """The ``tp`` ranks' bf16 pair partials at ``pair_of(src)`` of the mailbox
    region ``own``, summed in fp32 in rank order 0 .. tp - 1: one poll batch ->
    (lo sum, hi sum). Every rank sums in the same order, so they agree."""
    ws = poll([(own, pair_of(src), 1) for src in range(tp)])
    acc0, acc1 = bf2_f32(ws[0][0])
    for src in range_constexpr(1, tp):
        lo, hi = bf2_f32(ws[src][0])
        acc0 = acc0 + lo
        acc1 = acc1 + hi
    return acc0, acc1

traced(fn)

Apply flydsl's AST rewriting to a helper, so its while / if over traced values lower to scf ops. Kernels rewrite only their own body.

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
180
181
182
183
def traced(fn):
    """Apply flydsl's AST rewriting to a helper, so its ``while`` / ``if`` over
    traced values lower to scf ops. Kernels rewrite only their own body."""
    return ASTRewriter.transform(fn)

xred(v, off, op)

op(v, lane ^ off) for a symmetric op; 32 / 16 take both halves of one swap.

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
154
155
156
157
158
159
160
161
162
163
def xred(v, off, op):
    """op(v, lane ^ off) for a symmetric op; 32 / 16 take both halves of one swap."""
    if off < 16:
        return op(v, xshfl(v, off))
    is_f = isinstance(v, fx.Float32)
    x = v.bitcast(fx.Int32) if is_f else fx.Int32(v)
    a, b = permlane_swap(off, x, x)
    if is_f:
        return op(a.bitcast(fx.Float32), b.bitcast(fx.Float32))
    return op(type(v)(a), type(v)(b))

xshfl(v, off)

Value of lane lane ^ off: offsets 32 / 16 via v_permlane*_swap, in-row offsets via DPP (VALU latency, not LDS). Powers of two only: any other offset would take the lane ^ 1 path.

Source code in vllm/models/deepseek_v41/amd/mono/common/ops.py 
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
def xshfl(v, off):
    """Value of lane ``lane ^ off``: offsets 32 / 16 via v_permlane*_swap, in-row
    offsets via DPP (VALU latency, not LDS). Powers of two only: any other
    offset would take the lane ^ 1 path."""
    assert off in (1, 2, 4, 8, 16, 32), off
    if off >= 16:
        return v.shuffle_xor(off, 64)
    is_f = isinstance(v, fx.Float32)
    x = v.bitcast(fx.Int32) if is_f else fx.Int32(v)
    if off == 8:
        y = fx.Int32(update_dpp_i32(x, x, 0x118, 0xF, 0xC, False))
        y = fx.Int32(update_dpp_i32(y, x, 0x108, 0xF, 0x3, False))
    elif off == 4:
        y = fx.Int32(update_dpp_i32(x, x, 0x114, 0xF, 0xA, False))
        y = fx.Int32(update_dpp_i32(y, x, 0x104, 0xF, 0x5, False))
    elif off == 2:
        y = fx.Int32(update_dpp_i32(x, x, 0x4E, 0xF, 0xF, False))
    else:
        y = fx.Int32(update_dpp_i32(x, x, 0xB1, 0xF, 0xF, False))
    return y.bitcast(fx.Float32) if is_f else y