vllm.models.muse_spark.amd.mono.kernels.dev

Device-side primitives of the Muse Spark monokernel.

Adapted from ATOM's atom/mono/device (ops.py, sync.py; MIT, AMD). Differences from ATOM that matter:

• Mailbox tags are epoch based (tag = epoch * 64 + slot + 1 with the epoch a device word bumped once per forward inside the CUDA graph) instead of ATOM's zero-at-step-start discipline: nothing is ever zeroed.
• Every spin is bounded by a wall-clock timeout (s_memrealtime, 100 MHz) that writes a diagnostic record and sets a sticky FAILED word, so a bug turns into a report instead of a hang.

Functions:

• bf16_pair –

  Two f32 -> one i32 word holding (bf16(a) low, bf16(b) high), RNE.

• bpermute –

  Lane src_lane's i32 v (ds_bpermute: the LDS crossbar, no LDS memory).

• e8m0_f32 –

  E8M0 biased exponent (an i32 holding 1..254) -> the f32 power of two.

• fp4x8_bf16 –

  One dword of 8 E2M1 (low nibble first) times scale_f32 -> bf16x8

• fresh –

  An opaque v_mov: stops LLVM hoisting / CSE of lane math into long-lived

• ld_vec_if –

  words dwords at dword i, or zeros without a memory access when not ok

• mb_poll –

  [(base, pair index, npairs in {1, 2})] -> value words, all re-loaded

• mb_put –

  Pair i := (v, tag), device scope (CM_SYS: system scope, a peer GPU's buffer).

• mb_put2 –

  Pairs i, i+1 := (w0, tag), (w1, tag): one 16 B store (i even).

• mb_put_if –

  Pair i := (v, tag) when ok (branch-free, see st_if).

• mfma_bf16 –

  16x16x32 bf16 MFMA: a, b bf16x8 vectors, c f32x4.

• st_if –

  buffer_store of v (i32 or i32 vector) at dword i, dropped (out-of-range offset, no

• traced –

  FlyDSL AST rewriting for a helper (if / while on traced values).

_spin_bounded(load_all, pending, health, region_id, index, tag)

Re-load until pending clears or TIMEOUT_TICKS pass; on a timeout the CTA's first record goes to health (word 0 = sticky FAILED, words 8 + 8 bid .. = record) and later waits of this launch exit at once.

Source code in vllm/models/muse_spark/amd/mono/kernels/dev.py

@traced
def _spin_bounded(load_all, pending, health, region_id, index, tag):
    """Re-load until ``pending`` clears or TIMEOUT_TICKS pass; on a timeout the
    CTA's first record goes to ``health`` (word 0 = sticky FAILED, words
    8 + 8 bid .. = record) and later waits of this launch exit at once."""
    t0 = memrealtime()
    v = load_all()
    late = fx.Int32(0)
    failed = ld_i32(health, 0, CM_DEV)
    while pending(v) & (late == 0) & (failed == 0):
        rocdl.s_sleep(1)
        v = load_all()
        late = (memrealtime() - t0 > TIMEOUT_TICKS).select(fx.Int32(1), fx.Int32(0))
    if late != 0:
        bid = fx.block_idx.x
        at = fx.Int32(memrealtime() & 0x7FFFFFFF)
        record = [1, region_id, index, tag, v[1], bid, at, 0]
        for h in range_constexpr(2):
            bo.buffer_store(
                fx.Vector.from_elements(
                    [fx.Int32(w) for w in record[4 * h : 4 * h + 4]], fx.Int32
                ),
                rsrc(health),
                DIAG_WORDS + bid * DIAG_WORDS + 4 * h,
                cache_modifier=CM_DEV,
            )
        bo.buffer_store(fx.Int32(1), rsrc(health), 0, cache_modifier=CM_DEV)
    return v

bf16_pair(a, b)

Two f32 -> one i32 word holding (bf16(a) low, bf16(b) high), RNE.

Source code in vllm/models/muse_spark/amd/mono/kernels/dev.py

def bf16_pair(a, b):
    """Two f32 -> one i32 word holding (bf16(a) low, bf16(b) high), RNE."""
    return (
        fx.Vector.from_elements([a, b], fx.Float32).to(fx.BFloat16).bitcast(fx.Int32)[0]
    )

bpermute(v, src_lane)

Lane src_lane's i32 v (ds_bpermute: the LDS crossbar, no LDS memory).

Source code in vllm/models/muse_spark/amd/mono/kernels/dev.py

def bpermute(v, src_lane):
    """Lane ``src_lane``'s i32 v (ds_bpermute: the LDS crossbar, no LDS memory)."""
    return fx.Int32(
        rocdl.ds_bpermute(T.i32, fx.Int32(src_lane) * 4, fx.Int32(v).ir_value())
    )

e8m0_f32(code)

E8M0 biased exponent (an i32 holding 1..254) -> the f32 power of two.

Source code in vllm/models/muse_spark/amd/mono/kernels/dev.py

def e8m0_f32(code):
    """E8M0 biased exponent (an i32 holding 1..254) -> the f32 power of two."""
    return (fx.Int32(code) << 23).bitcast(fx.Float32)

fp4x8_bf16(word, scale_f32)

One dword of 8 E2M1 (low nibble first) times scale_f32 -> bf16x8 (exact: v_cvt_scalef32_pk_bf16_fp4, the aiter Muse a16w4 decode).

Source code in vllm/models/muse_spark/amd/mono/kernels/dev.py

def fp4x8_bf16(word, scale_f32):
    """One dword of 8 E2M1 (low nibble first) times ``scale_f32`` -> bf16x8
    (exact: v_cvt_scalef32_pk_bf16_fp4, the aiter Muse a16w4 decode)."""
    parts = []
    for sel in range_constexpr(4):
        pair = rocdl.cvt_scalef32_pk_bf16_fp4(
            T.vec(2, T.bf16), as_ir_value(fx.Int32(word)), as_ir_value(scale_f32), sel
        )
        parts.append(fx.Vector(pair).bitcast(fx.Int32)[0])
    return fx.Vector.from_elements(parts, fx.Int32).bitcast(fx.BFloat16)

fresh(v)

An opaque v_mov: stops LLVM hoisting / CSE of lane math into long-lived registers (ATOM ops.fresh).

Source code in vllm/models/muse_spark/amd/mono/kernels/dev.py

def fresh(v):
    """An opaque v_mov: stops LLVM hoisting / CSE of lane math into long-lived
    registers (ATOM ops.fresh)."""
    return fx.Int32(
        llvm.inline_asm(
            T.i32, [fx.Int32(v).ir_value()], "v_mov_b32 $0, $1", "=v,v",
            has_side_effects=True
        )
    )  # fmt: skip

ld_vec_if(ok, ptr, i, words, nbytes, cm=0)

words dwords at dword i, or zeros without a memory access when not ok (out-of-range buffer offset; nbytes must bound the buffer).

Source code in vllm/models/muse_spark/amd/mono/kernels/dev.py

def ld_vec_if(ok, ptr, i, words, nbytes, cm=0):
    """``words`` dwords at dword ``i``, or zeros without a memory access when not ``ok``
    (out-of-range buffer offset; ``nbytes`` must bound the buffer)."""
    return fx.Vector(
        bo.buffer_load(
            rsrc(ptr, nbytes),
            ok.select(i, fx.Int32(OOB_DW)),
            vec_width=words,
            dtype=T.i32,
            cache_modifier=cm,
        )
    )

mb_poll(tag, health, specs, region_id=0, batch=16, oks=None, cm=CM_DEV)

[(base, pair index, npairs in {1, 2})] -> value words, all re-loaded together until every tag matches (one round trip after the last lands). oks (optional, one traced bool a spec): a false spec is not loaded (an out-of-range buffer offset, no memory access), returns zeros, and is not waited for.

Source code in vllm/models/muse_spark/amd/mono/kernels/dev.py

def mb_poll(tag, health, specs, region_id=0, batch=16, oks=None, cm=CM_DEV):
    """[(base, pair index, npairs in {1, 2})] -> value words, all re-loaded
    together until every tag matches (one round trip after the last lands).
    ``oks`` (optional, one traced bool a spec): a false spec is not loaded (an
    out-of-range buffer offset, no memory access), returns zeros, and is not
    waited for."""
    if len(specs) == 0:
        return []
    if len(specs) > batch:
        o1 = None if oks is None else oks[:batch]
        o2 = None if oks is None else oks[batch:]
        return mb_poll(tag, health, specs[:batch], region_id, batch, o1, cm) + mb_poll(
            tag, health, specs[batch:], region_id, batch, o2, cm
        )

    def load_all():
        words = []
        for si, (b, i, n) in enumerate(specs):
            idx = fx.Int32(i) * 2
            if oks is not None:
                idx = oks[si].select(idx, fx.Int32(OOB_DW))
            w = fx.Vector(
                bo.buffer_load(
                    rsrc(b, OOB_BOUND),
                    idx,
                    vec_width=2 * n,
                    dtype=T.i32,
                    cache_modifier=cm,
                )
            )
            words += [w[e] for e in range(2 * n)]
        return fx.Vector.from_elements(words, fx.Int32)

    def pending(v):
        bad = None
        e = 0
        for si, (_, _, n) in enumerate(specs):
            for q in range(n):
                b1 = v[e + 2 * q + 1] != tag
                if oks is not None:
                    b1 = b1 & oks[si]
                bad = b1 if bad is None else (bad | b1)
            e += 2 * n
        return bad

    v = _spin_bounded(load_all, pending, health, region_id, specs[0][1], tag)
    outs, e = [], 0
    for _, _, n in specs:
        outs.append([v[e + 2 * q] for q in range(n)])
        e += 2 * n
    return outs

mb_put(tag, base, i, v, cm=CM_DEV)

Pair i := (v, tag), device scope (CM_SYS: system scope, a peer GPU's buffer).

Source code in vllm/models/muse_spark/amd/mono/kernels/dev.py

def mb_put(tag, base, i, v, cm=CM_DEV):
    """Pair i := (v, tag), device scope (CM_SYS: system scope, a peer GPU's buffer)."""
    bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v)
    bo.buffer_store(
        fx.Vector.from_elements([bits, tag], fx.Int32),
        rsrc(base),
        i * 2,
        cache_modifier=cm,
    )

mb_put2(tag, base, i, w0, w1)

Pairs i, i+1 := (w0, tag), (w1, tag): one 16 B store (i even).

Source code in vllm/models/muse_spark/amd/mono/kernels/dev.py

def mb_put2(tag, base, i, w0, w1):
    """Pairs i, i+1 := (w0, tag), (w1, tag): one 16 B store (i even)."""
    bo.buffer_store(
        fx.Vector.from_elements([fx.Int32(w0), tag, fx.Int32(w1), tag], fx.Int32),
        rsrc(base),
        i * 2,
        cache_modifier=CM_DEV,
    )

mb_put_if(ok, tag, base, i, v, cm=CM_DEV)

Pair i := (v, tag) when ok (branch-free, see st_if).

Source code in vllm/models/muse_spark/amd/mono/kernels/dev.py

def mb_put_if(ok, tag, base, i, v, cm=CM_DEV):
    """Pair i := (v, tag) when ``ok`` (branch-free, see st_if)."""
    bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v)
    st_if(ok, fx.Vector.from_elements([bits, tag], fx.Int32), base, fx.Int32(i) * 2, cm)

mfma_bf16(a, b, c)

16x16x32 bf16 MFMA: a, b bf16x8 vectors, c f32x4.

Source code in vllm/models/muse_spark/amd/mono/kernels/dev.py

def mfma_bf16(a, b, c):
    """16x16x32 bf16 MFMA: a, b bf16x8 vectors, c f32x4."""
    return fx.Vector(rocdl.mfma_f32_16x16x32_bf16(T.vec(4, T.f32), [a, b, c]))

st_if(ok, v, base, i, cm=0)

buffer_store of v (i32 or i32 vector) at dword i, dropped (out-of-range offset, no branch) when not ok: keeps the waitcnt pass from draining at a join.

Source code in vllm/models/muse_spark/amd/mono/kernels/dev.py

def st_if(ok, v, base, i, cm=0):
    """buffer_store of v (i32 or i32 vector) at dword i, dropped (out-of-range offset, no
    branch) when not ``ok``: keeps the waitcnt pass from draining at a join."""
    bo.buffer_store(
        v,
        rsrc(base, OOB_BOUND),
        ok.select(fx.Int32(i), fx.Int32(OOB_DW)),
        cache_modifier=cm,
    )

traced(fn)

FlyDSL AST rewriting for a helper (if / while on traced values).

Source code in vllm/models/muse_spark/amd/mono/kernels/dev.py

def traced(fn):
    """FlyDSL AST rewriting for a helper (``if`` / ``while`` on traced values)."""
    return ASTRewriter.transform(fn)