vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync

The tagged mailbox: hand-offs between the CTAs of one GPU (device scope) and between GPUs (system scope), every address a named region (Addr) recorded when the kernel is traced (trace).

Classes:

• Addr –

  A mailbox region's address: the region's name and space, recorded when the

• Mailbox –

  Tagged-pair hand-off between CTAs (device scope) and GPUs (system scope).

Functions:

• poll_sent –

  Poll words idxs of this rank's sentinel region a until none is

• preg –

  Peer-symmetric region region at byte offset of a rank's buffer

• publish –

  This CTA's plain stores made readable by every CTA, then (thread who)

• put_sent –

  Word i of the peer sentinel region a := word (sent_word),

• rearm_sent –

  Words idxs of this rank's sentinel region a := the sentinel (a

• sent_word –

  A pushed bf16-pair word (an f32-typed bf16_pair), never the

• shift –

  a moved by delta bytes (a traced or Python integer), same region.

• sreg –

  Scratch region region at byte offset of base.

• stall –

  The fault hook (tests only): CTA 0 of rank fault_rank stalls ticks

Addr dataclass

A mailbox region's address: the region's name and space, recorded when the kernel is traced (trace), and the traced address. The space also sets the hand-off's scope: device for SCRATCH, system for PEER.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py

@dataclass(frozen=True)
class Addr:
    """A mailbox region's address: the region's name and space, recorded when the
    kernel is traced (``trace``), and the traced address. The space also
    sets the hand-off's scope: device for SCRATCH, system for PEER."""

    region: str
    space: Space
    value: object

Mailbox

Tagged-pair hand-off between CTAs (device scope) and GPUs (system scope).

Every 32-bit value is stored next to this launch's tag; a consumer polls the payload until all tags match, so a hand-off costs one round trip (no store drain, no separate flag). The tag is layer + 1: it tells the layers of one step apart, and the runner zeroes every mailbox at a step's start (StepMailboxes.begin_step), so a pair from an earlier step never carries a live tag. Every address is an Addr: its space picks the scope.

Methods:

• __init__ –

  health: this rank's health words (FAILED, then a wait record a

• poll –

  Batched poll of pairs specs = [(Addr, pair index, npairs in {1, 2})],

• put_bf –

  Elements i .. i+len(vs) (2 or 4, i aligned) as packed bf16 pairs, one

• put_words –

  Raw 32-bit words (1 or 2) at pairs i.. as (word, tag). Integer payloads

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py

class Mailbox:
    """Tagged-pair hand-off between CTAs (device scope) and GPUs (system scope).

    Every 32-bit value is stored next to this launch's tag; a consumer
    polls the payload until all tags match, so a hand-off costs one round trip
    (no store drain, no separate flag). The tag is ``layer + 1``: it tells the
    layers of one step apart, and the runner zeroes every mailbox at a step's
    start (``StepMailboxes.begin_step``), so a pair from an earlier step never
    carries a live tag. Every address is an ``Addr``: its space picks the scope.
    """

    def __init__(self, layer, health, alarm, region_ids):
        """``health``: this rank's health words (FAILED, then a wait record a
        CTA); ``alarm``: this rank's flag word on the output rank;
        ``region_ids``: region name -> the id a record names it by."""
        self.tag = layer + 1
        self.health = health
        self.alarm = alarm
        self.region_ids = region_ids

    def put(self, a, i, v):
        bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v)
        bo.buffer_store(
            fx.Vector.from_elements([bits, self.tag], fx.Int32),
            rsrc(_noted(a, "put")),
            i * 2,
            cache_modifier=_SCOPE_CM[a.space],
        )

    def put_bf(self, a, i, vs):
        """Elements i .. i+len(vs) (2 or 4, i aligned) as packed bf16 pairs, one
        store."""
        words = []
        for j in range_constexpr(len(vs) // 2):
            words += [bf16_pair(vs[2 * j], vs[2 * j + 1]).bitcast(fx.Int32), self.tag]
        bo.buffer_store(
            fx.Vector.from_elements(words, fx.Int32),
            rsrc(_noted(a, "put")),
            i,
            cache_modifier=_SCOPE_CM[a.space],
        )

    def put_words(self, a, i, words):
        """Raw 32-bit words (1 or 2) at pairs i.. as (word, tag). Integer payloads
        (fp8 words) must come this way: through an f32 value a NaN bit pattern is
        not kept."""
        vals = []
        for w in words:
            vals += [fx.Int32(w), self.tag]
        bo.buffer_store(
            fx.Vector.from_elements(vals, fx.Int32),
            rsrc(_noted(a, "put")),
            i * 2,
            cache_modifier=_SCOPE_CM[a.space],
        )

    def poll(self, specs, batch=POLL_MAX):
        """Batched poll of pairs ``specs`` = [(Addr, pair index, npairs in {1, 2})],
        all of one space.

        All pairs are re-loaded together while any tag is stale, so a batch costs
        one round trip after its last producer lands. Returns the value words."""
        if const_expr(len(specs) == 0):
            return []
        if const_expr(len(specs) > batch):
            return self.poll(specs[:batch], batch) + self.poll(specs[batch:], batch)
        spaces = {a.space for a, _, _ in specs}
        assert len(spaces) == 1, f"one poll batch spans {spaces}"
        cm = _SCOPE_CM[spaces.pop()]
        bases = [_noted(a, "poll") for a, _, _ in specs]
        tag = self.tag

        def load_all():
            words = []
            for b, (_, i, n) in zip(bases, specs):
                w = fx.Vector(
                    bo.buffer_load(
                        rsrc(b),
                        fx.Int32(i) * 2,
                        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)

        nw = sum(2 * n for _, _, n in specs)

        def pending(v):
            bad = v[1] != tag
            for e in range_constexpr(3, nw, 2):
                bad = bad | (v[e] != tag)
            return bad

        rid = self.region_ids.get(specs[0][0].region, -1)
        v = _spin(load_all, pending, self.health, self.alarm, rid, 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

__init__(layer, health, alarm, region_ids)

health: this rank's health words (FAILED, then a wait record a CTA); alarm: this rank's flag word on the output rank; region_ids: region name -> the id a record names it by.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py

def __init__(self, layer, health, alarm, region_ids):
    """``health``: this rank's health words (FAILED, then a wait record a
    CTA); ``alarm``: this rank's flag word on the output rank;
    ``region_ids``: region name -> the id a record names it by."""
    self.tag = layer + 1
    self.health = health
    self.alarm = alarm
    self.region_ids = region_ids

poll(specs, batch=POLL_MAX)

Batched poll of pairs specs = [(Addr, pair index, npairs in {1, 2})], all of one space.

All pairs are re-loaded together while any tag is stale, so a batch costs one round trip after its last producer lands. Returns the value words.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py

def poll(self, specs, batch=POLL_MAX):
    """Batched poll of pairs ``specs`` = [(Addr, pair index, npairs in {1, 2})],
    all of one space.

    All pairs are re-loaded together while any tag is stale, so a batch costs
    one round trip after its last producer lands. Returns the value words."""
    if const_expr(len(specs) == 0):
        return []
    if const_expr(len(specs) > batch):
        return self.poll(specs[:batch], batch) + self.poll(specs[batch:], batch)
    spaces = {a.space for a, _, _ in specs}
    assert len(spaces) == 1, f"one poll batch spans {spaces}"
    cm = _SCOPE_CM[spaces.pop()]
    bases = [_noted(a, "poll") for a, _, _ in specs]
    tag = self.tag

    def load_all():
        words = []
        for b, (_, i, n) in zip(bases, specs):
            w = fx.Vector(
                bo.buffer_load(
                    rsrc(b),
                    fx.Int32(i) * 2,
                    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)

    nw = sum(2 * n for _, _, n in specs)

    def pending(v):
        bad = v[1] != tag
        for e in range_constexpr(3, nw, 2):
            bad = bad | (v[e] != tag)
        return bad

    rid = self.region_ids.get(specs[0][0].region, -1)
    v = _spin(load_all, pending, self.health, self.alarm, rid, 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

put_bf(a, i, vs)

Elements i .. i+len(vs) (2 or 4, i aligned) as packed bf16 pairs, one store.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py

def put_bf(self, a, i, vs):
    """Elements i .. i+len(vs) (2 or 4, i aligned) as packed bf16 pairs, one
    store."""
    words = []
    for j in range_constexpr(len(vs) // 2):
        words += [bf16_pair(vs[2 * j], vs[2 * j + 1]).bitcast(fx.Int32), self.tag]
    bo.buffer_store(
        fx.Vector.from_elements(words, fx.Int32),
        rsrc(_noted(a, "put")),
        i,
        cache_modifier=_SCOPE_CM[a.space],
    )

put_words(a, i, words)

Raw 32-bit words (1 or 2) at pairs i.. as (word, tag). Integer payloads (fp8 words) must come this way: through an f32 value a NaN bit pattern is not kept.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py

def put_words(self, a, i, words):
    """Raw 32-bit words (1 or 2) at pairs i.. as (word, tag). Integer payloads
    (fp8 words) must come this way: through an f32 value a NaN bit pattern is
    not kept."""
    vals = []
    for w in words:
        vals += [fx.Int32(w), self.tag]
    bo.buffer_store(
        fx.Vector.from_elements(vals, fx.Int32),
        rsrc(_noted(a, "put")),
        i * 2,
        cache_modifier=_SCOPE_CM[a.space],
    )

_spin(load_all, pending, health, alarm, region_id, index, tag)

Re-issue load_all until pending clears, for at most WAIT_TICKS, and not at all once this rank has failed (health word FAILED), so a failed step runs to its end instead of hanging. A side-effecting op in the loop keeps the loads from being hoisted.

A CTA's first give-up is recorded in its health record; it then raises the output rank's flag for this rank (alarm) and drains that store before it sets FAILED, so no later store of this rank's can overtake the flag.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py

@traced
def _spin(load_all, pending, health, alarm, region_id, index, tag):
    """Re-issue ``load_all`` until ``pending`` clears, for at most WAIT_TICKS, and
    not at all once this rank has failed (``health`` word FAILED), so a failed
    step runs to its end instead of hanging. A side-effecting op in the loop keeps
    the loads from being hoisted.

    A CTA's first give-up is recorded in its ``health`` record; it then raises the
    output rank's flag for this rank (``alarm``) and drains that store before it
    sets FAILED, so no later store of this rank's can overtake the flag."""
    t0 = memrealtime()
    v = load_all()
    why = fx.Int32(0)  # 1: timed out, 2: the rank had failed
    while pending(v) & (why == 0):
        rocdl.s_sleep(1)
        v = load_all()
        failed = ld_health(health, FAILED) != 0
        late = memrealtime() - t0 > WAIT_TICKS
        why = failed.select(fx.Int32(2), late.select(fx.Int32(1), fx.Int32(0)))
    if why == 1:
        bid = fx.block_idx.x
        rec = (bid + 1) * RECORD_WORDS
        if ld_health(health, rec) == 0:
            at = fx.Int32(memrealtime() & 0x7FFFFFFF)
            record = [1, region_id, index, tag, v[1], bid, at, 0]
            for h in range_constexpr(2):  # 16 B a store
                bo.buffer_store(
                    fx.Vector.from_elements(
                        [fx.Int32(w) for w in record[4 * h : 4 * h + 4]], fx.Int32
                    ),
                    rsrc(health),
                    rec + 4 * h,
                    cache_modifier=CM_SYS,
                )
        bo.buffer_store(fx.Int32(1), rsrc(alarm), 0, cache_modifier=CM_SYS)
        rocdl.s_waitcnt(vmcnt=0)
        bo.buffer_store(fx.Int32(1), rsrc(health), FAILED, cache_modifier=CM_SYS)
    return v

poll_sent(mb, a, idxs, lives=None, rearm=True)

Poll words idxs of this rank's sentinel region a until none is the sentinel (with fail-stop), then re-arm them. lives: dead indices are not waited for or re-armed. A read is destructive: one thread a word. rearm=False: the caller re-arms later (rearm_sent).

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py

def poll_sent(mb, a, idxs, lives=None, rearm=True):
    """Poll words ``idxs`` of this rank's sentinel region ``a`` until none is
    the sentinel (with fail-stop), then re-arm them. ``lives``: dead indices
    are not waited for or re-armed. A read is destructive: one thread a word.
    ``rearm=False``: the caller re-arms later (``rearm_sent``)."""
    base = _noted(a, "poll")
    cm = _SCOPE_CM[a.space]
    n = len(idxs)
    if const_expr(n == 1):
        idxs = [idxs[0], idxs[0]]  # _spin's record reads v[1]
        if const_expr(lives is not None):
            lives = [lives[0], lives[0]]

    def load_all():
        return fx.Vector.from_elements(
            [
                fx.Int32(
                    bo.buffer_load(
                        rsrc(base),
                        fx.Int32(i),
                        vec_width=1,
                        dtype=T.i32,
                        cache_modifier=cm,
                    )
                )
                for i in idxs
            ],
            fx.Int32,
        )

    def pending(v):
        bad = None
        for e in range_constexpr(len(idxs)):
            b = v[e] == fx.Int32(SENT)
            if const_expr(lives is not None):
                b = b & lives[e]
            bad = b if bad is None else bad | b
        return bad

    rid = mb.region_ids.get(a.region, -1)
    v = _spin(load_all, pending, mb.health, mb.alarm, rid, fx.Int32(idxs[0]), SENT)
    if const_expr(rearm):
        rearm_sent(a, idxs[:n], None if lives is None else lives[:n], noted=False)
    return [v[e] for e in range(n)]

preg(base, offset, region)

Peer-symmetric region region at byte offset of a rank's buffer base.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py

def preg(base, offset, region):
    """Peer-symmetric region ``region`` at byte ``offset`` of a rank's buffer
    ``base``."""
    return Addr(region, Space.PEER, base + fx.Int64(offset))

publish(put, a, i, v, who)

This CTA's plain stores made readable by every CTA, then (thread who) the mailbox pair (a, i) := v. The stores are drained first: a flag may not overtake the data it announces, and a workgroup-scope release fence does not wait for them (LLVM's AMDGPU memory model omits vmcnt outside tgsplit mode: the workgroup's waves share one CU, the readers here do not).

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py

@traced
def publish(put, a, i, v, who):
    """This CTA's plain stores made readable by every CTA, then (thread ``who``)
    the mailbox pair (``a``, i) := v. The stores are drained first: a flag may
    not overtake the data it announces, and a workgroup-scope release fence does
    not wait for them (LLVM's AMDGPU memory model omits vmcnt outside tgsplit
    mode: the workgroup's waves share one CU, the readers here do not)."""
    rocdl.s_waitcnt(vmcnt=0)
    gpu.barrier()
    if who:
        put(a, i, v)

put_sent(a, i, word)

Word i of the peer sentinel region a := word (sent_word), system scope.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py

def put_sent(a, i, word):
    """Word ``i`` of the peer sentinel region ``a`` := ``word`` (``sent_word``),
    system scope."""
    bo.buffer_store(
        fx.Int32(word), rsrc(_noted(a, "put")), i, cache_modifier=_SCOPE_CM[a.space]
    )

rearm_sent(a, idxs, lives=None, noted=True)

Words idxs of this rank's sentinel region a := the sentinel (a dead index skipped), after this thread read them (poll_sent).

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py

def rearm_sent(a, idxs, lives=None, noted=True):
    """Words ``idxs`` of this rank's sentinel region ``a`` := the sentinel (a
    dead index skipped), after this thread read them (``poll_sent``)."""
    base = _noted(a, "put") if noted else a.value
    cm = _SCOPE_CM[a.space]
    for e in range_constexpr(len(idxs)):
        if const_expr(lives is None):
            bo.buffer_store(
                fx.Int32(SENT), rsrc(base), fx.Int32(idxs[e]), cache_modifier=cm
            )
        else:
            _rearm(base, idxs[e], lives[e], cm)

sent_word(w)

A pushed bf16-pair word (an f32-typed bf16_pair), never the sentinel.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py

def sent_word(w):
    """A pushed bf16-pair word (an f32-typed ``bf16_pair``), never the
    sentinel."""
    w = w.bitcast(fx.Int32) if isinstance(w, fx.Float32) else fx.Int32(w)
    return (w == fx.Int32(SENT)).select(fx.Int32(SENT_NAN), w)

shift(a, delta)

a moved by delta bytes (a traced or Python integer), same region.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py

def shift(a, delta):
    """``a`` moved by ``delta`` bytes (a traced or Python integer), same region."""
    return Addr(a.region, a.space, a.value + delta)

sreg(base, offset, region)

Scratch region region at byte offset of base.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py

def sreg(base, offset, region):
    """Scratch region ``region`` at byte ``offset`` of ``base``."""
    return Addr(region, Space.SCRATCH, base + fx.Int64(offset))

stall(fault_rank, fault_epoch, fault_layer, ticks, rank, epoch, layer)

The fault hook (tests only): CTA 0 of rank fault_rank stalls ticks once, in the launch for slot fault_layer of the step at epoch fault_epoch.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py

@traced
def stall(fault_rank, fault_epoch, fault_layer, ticks, rank, epoch, layer):
    """The fault hook (tests only): CTA 0 of rank ``fault_rank`` stalls ``ticks``
    once, in the launch for slot ``fault_layer`` of the step at epoch
    ``fault_epoch``."""
    if (
        (rank == fault_rank)
        & (uniform(ld_i32(epoch, 0)) == fault_epoch)
        & (layer == fault_layer)
        & (fx.block_idx.x == 0)
    ):
        t0 = memrealtime()
        now = memrealtime()
        while now - t0 < ticks:
            rocdl.s_sleep(127)
            now = memrealtime()