Skip to content

vllm.models.deepseek_v41.amd.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).

Classes:

  • Addr –

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

  • Mailbox –

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

Functions:

  • 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)

  • sreg –

    Scratch region region at byte offset of base.

Addr dataclass

A mailbox region's address: the region's name and space, and the traced address. The space sets the hand-off's scope: device for SCRATCH, system for PEER.

Source code in vllm/models/deepseek_v41/amd/mono/common/sync.py 
44
45
46
47
48
49
50
51
52
@dataclass(frozen=True)
class Addr:
    """A mailbox region's address: the region's name and space, and the traced
    address. The space 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 mono layer's launches tag their hand-offs with their epoch (layer.py), so a pair from an earlier step never carries a live tag. Every address is an Addr: its space picks the scope.

Methods:

  • 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/deepseek_v41/amd/mono/common/sync.py 
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 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
114
115
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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
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 mono layer's launches tag their
    hand-offs with their epoch (``layer.py``), 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, tag):
        self.tag = tag

    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(a.value),
            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(a.value),
            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(a.value),
            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 = [a.value 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

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

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/deepseek_v41/amd/mono/common/sync.py 
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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
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 = [a.value 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

    v = _spin(load_all, pending)
    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/deepseek_v41/amd/mono/common/sync.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
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(a.value),
        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/deepseek_v41/amd/mono/common/sync.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
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(a.value),
        i * 2,
        cache_modifier=_SCOPE_CM[a.space],
    )

_spin(load_all, pending)

Re-issue load_all until pending clears; a side-effecting op in the loop keeps the loads from being hoisted.

Source code in vllm/models/deepseek_v41/amd/mono/common/sync.py 
28
29
30
31
32
33
34
35
36
@traced
def _spin(load_all, pending):
    """Re-issue ``load_all`` until ``pending`` clears; a side-effecting op in the
    loop keeps the loads from being hoisted."""
    v = load_all()
    while pending(v):
        rocdl.s_sleep(1)
        v = load_all()
    return v

preg(base, offset, region)

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

Source code in vllm/models/deepseek_v41/amd/mono/common/sync.py 
60
61
62
63
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/deepseek_v41/amd/mono/common/sync.py 
167
168
169
170
171
172
173
174
175
176
177
@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)

sreg(base, offset, region)

Scratch region region at byte offset of base.

Source code in vllm/models/deepseek_v41/amd/mono/common/sync.py 
55
56
57
def sreg(base, offset, region):
    """Scratch region ``region`` at byte ``offset`` of ``base``."""
    return Addr(region, Space.SCRATCH, base + fx.Int64(offset))