Skip to content

vllm.models.minimax_m3.amd.specialized_decode.mono.sync

Bounded tagged hand-offs and the sticky error word.

The contract, in the LLVM "AMDGPU Memory Model" (llvm/docs/AMDGPUMemoryModel.md, AMM) and AMDGPUUsage "Execution Barriers":

  • A pair is one naturally aligned 64-bit atomic, (value | tag << 32), monotonic at its space's scope (agent for this GPU's scratch, system for the peer heap). An atomic access is never split, so a load that sees a pair's tag sees that store's value. A triple pair (put_triple, the attention seam) is the same atomic with three bf16 in bits 0 .. 47 under a 16-bit tag, unique between epoch resets (the kernel's EPOCH_LIMIT).
  • Data a pair announces is stored store-available and loaded load-visible at agent scope (AMM: agent atomics of 16 to 64 bits, llvm.amdgcn.av.load. b128; store_dev / load_dev).
  • Happens-before from producer to consumer: the producer's release fence before the pair (publish), the consumer's acquire fence after the poll that saw it (acquire), and within a CTA release, s_barrier, acquire (gpu.barrier; "Execution Barriers": they synchronize). These fences are tagged !{"amdgcn-av", "none"}: they synchronize and keep the compiler from moving accesses across them, but make nothing available or visible.
  • AMM visibility: a store-available write that happens-before a load-visible read of the same agent is visible to it, so the read returns that write (the previous launch's writes to the location are ordered before it by the kernel boundary).

The av operations are lowered with the agent-scope cache bits (sc1) and the av-none fences to waits only, so no L2 write-back or invalidate is emitted. The GFX942 workgroup release waits lgkmcnt(0) only outside tgsplit mode, so a publishing CTA also has every wave wait for its own stores before the barrier (cta_sync(drain=True)). The step-end fault handshake uses plain system release / acquire once a step.

A poll gives up after ticks of the 100 MHz clock or once any rank has set the error word; a timeout sets [1, rank, epoch, site] on every rank.

Classes:

  • Mailbox –

    Tagged pairs of one launch (tag), every address an Addr.

Functions:

  • cta_sync –

    gpu.barrier: workgroup release fence, s_barrier, acquire fence

  • load128 –

    Dwords i .. i + 3 (i a multiple of 4): one load-visible at scope.

  • load_dev –

    Dwords i .. i + n of this GPU's scratch (n in {1, 2, 4}, i aligned),

  • store_dev –

    Dwords i .. of this GPU's scratch := words (1 or 2, i aligned),

  • wave_sync –

    Every lane of the wave: what each lane did before (its acquires)

Mailbox

Tagged pairs of one launch (tag), every address an Addr.

Methods:

  • acquire –

    After a poll of published pairs, before what they announce.

  • poll –

    Pairs specs = [(Addr, pair index, npairs in {1, 2})], all of one

  • publish –

    Pair i of a := (v, tag) after a synchronizing release: what

  • put –

    Pair i of a := (v, tag), monotonic: a value no other data hangs on.

  • put_release_all –

    Pair i of every addrs := (v, tag) after one release fence.

  • put_triple –

    Pair i of a := (48 value bits | tag16 << 48), monotonic.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/sync.py 
231
232
233
234
235
236
237
238
239
240
241
242
243
244
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
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
class Mailbox:
    """Tagged pairs of one launch (``tag``), every address an ``Addr``."""

    def __init__(
        self, tag, err_base, peer_bases, rank, epoch, ticks=WAIT_TICKS, tag16=None
    ):
        self.tag = tag
        # a triple pair's 16-bit tag (``put_triple``), if the launch has one
        self.tag16 = tag16
        self.err_base = err_base
        self.peer_bases = peer_bases
        self.rank = rank
        self.epoch = epoch
        self.ticks = ticks

    def put(self, a, i, v):
        """Pair i of ``a`` := (v, tag), monotonic: a value no other data hangs on."""
        bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v)
        store64(a.value, i, pack(bits, self.tag), _SCOPE[a.space])

    def put_words(self, a, i, words):
        for k in range_constexpr(len(words)):
            store64(a.value, i + k, pack(words[k], self.tag), _SCOPE[a.space])

    def publish(self, a, i, v):
        """Pair i of ``a`` := (v, tag) after a synchronizing release: what
        happened before it (the CTA's stores past a ``cta_sync(drain=True)``,
        pairs it acquired) happens before an acquire of the pair."""
        sync_release(_SCOPE[a.space])
        self.put(a, i, v)

    def acquire(self, space):
        """After a poll of published pairs, before what they announce."""
        sync_acquire(_SCOPE[space])

    def put_release_all(self, addrs, i, v):
        """Pair i of every ``addrs`` := (v, tag) after one release fence."""
        release(_SCOPE[addrs[0].space])
        for a in addrs:
            self.put(a, i, v)

    def acquire_system(self):
        acquire_fence(None)

    def fail(self, site):
        lo = (fx.Int64(fx.Int32(self.rank)) << 32) | fx.Int64(1)
        hi = (fx.Int64(fx.Int32(site)) << 32) | fx.Int64(
            fx.Uint32(fx.Int32(self.epoch))
        )
        for base in self.peer_bases:
            store64(base, 1, hi, None)
            store64(base, 0, lo, None)

    def put_triple(self, a, i, bits48):
        """Pair i of ``a`` := (48 value bits | tag16 << 48), monotonic."""
        word = (bits48 & ((1 << 48) - 1)) | (fx.Int64(self.tag16) << 48)
        store64(a.value, i, word, _SCOPE[a.space])

    def poll(self, specs, batch=POLL_MAX, site=0, triples=False):
        """Pairs ``specs`` = [(Addr, pair index, npairs in {1, 2})], all of one
        space, re-loaded together until every tag is this launch's. Returns the
        value words; ``triples``: of triple pairs (``put_triple``), both
        words."""
        if const_expr(len(specs) == 0):
            return []
        if const_expr(len(specs) > batch):
            return self.poll(specs[:batch], batch, site, triples) + self.poll(
                specs[batch:], batch, site, triples
            )
        spaces = {a.space for a, _, _ in specs}
        assert len(spaces) == 1, f"one poll batch spans {spaces}"
        scope = _SCOPE[spaces.pop()]
        tag = self.tag16 if const_expr(triples) else self.tag

        def high_tag(w):
            return ((w >> 16) & 0xFFFF) if const_expr(triples) else w

        def load_all():
            words = []
            for a, i, n in specs:
                for q in range(n):
                    w = load64(a.value, fx.Int32(i) + q, scope)
                    words += [fx.Int32(w), fx.Int32(w >> 32)]
            return fx.Vector.from_elements(words, fx.Int32)

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

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

        v, timed_out = _spin(load_all, pending, self.err_base, self.ticks)
        _fail_if(timed_out, self.fail, site)
        outs, e = [], 0
        for _, _, n in specs:
            if const_expr(triples):
                outs.append([(v[e + 2 * q], v[e + 2 * q + 1]) for q in range(n)])
            else:
                outs.append([v[e + 2 * q] for q in range(n)])
            e += 2 * n
        return outs

acquire(space)

After a poll of published pairs, before what they announce.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/sync.py 
262
263
264
def acquire(self, space):
    """After a poll of published pairs, before what they announce."""
    sync_acquire(_SCOPE[space])

poll(specs, batch=POLL_MAX, site=0, triples=False)

Pairs specs = [(Addr, pair index, npairs in {1, 2})], all of one space, re-loaded together until every tag is this launch's. Returns the value words; triples: of triple pairs (put_triple), both words.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/sync.py 
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
def poll(self, specs, batch=POLL_MAX, site=0, triples=False):
    """Pairs ``specs`` = [(Addr, pair index, npairs in {1, 2})], all of one
    space, re-loaded together until every tag is this launch's. Returns the
    value words; ``triples``: of triple pairs (``put_triple``), both
    words."""
    if const_expr(len(specs) == 0):
        return []
    if const_expr(len(specs) > batch):
        return self.poll(specs[:batch], batch, site, triples) + self.poll(
            specs[batch:], batch, site, triples
        )
    spaces = {a.space for a, _, _ in specs}
    assert len(spaces) == 1, f"one poll batch spans {spaces}"
    scope = _SCOPE[spaces.pop()]
    tag = self.tag16 if const_expr(triples) else self.tag

    def high_tag(w):
        return ((w >> 16) & 0xFFFF) if const_expr(triples) else w

    def load_all():
        words = []
        for a, i, n in specs:
            for q in range(n):
                w = load64(a.value, fx.Int32(i) + q, scope)
                words += [fx.Int32(w), fx.Int32(w >> 32)]
        return fx.Vector.from_elements(words, fx.Int32)

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

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

    v, timed_out = _spin(load_all, pending, self.err_base, self.ticks)
    _fail_if(timed_out, self.fail, site)
    outs, e = [], 0
    for _, _, n in specs:
        if const_expr(triples):
            outs.append([(v[e + 2 * q], v[e + 2 * q + 1]) for q in range(n)])
        else:
            outs.append([v[e + 2 * q] for q in range(n)])
        e += 2 * n
    return outs

publish(a, i, v)

Pair i of a := (v, tag) after a synchronizing release: what happened before it (the CTA's stores past a cta_sync(drain=True), pairs it acquired) happens before an acquire of the pair.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/sync.py 
255
256
257
258
259
260
def publish(self, a, i, v):
    """Pair i of ``a`` := (v, tag) after a synchronizing release: what
    happened before it (the CTA's stores past a ``cta_sync(drain=True)``,
    pairs it acquired) happens before an acquire of the pair."""
    sync_release(_SCOPE[a.space])
    self.put(a, i, v)

put(a, i, v)

Pair i of a := (v, tag), monotonic: a value no other data hangs on.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/sync.py 
246
247
248
249
def put(self, a, i, v):
    """Pair i of ``a`` := (v, tag), monotonic: a value no other data hangs on."""
    bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v)
    store64(a.value, i, pack(bits, self.tag), _SCOPE[a.space])

put_release_all(addrs, i, v)

Pair i of every addrs := (v, tag) after one release fence.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/sync.py 
266
267
268
269
270
def put_release_all(self, addrs, i, v):
    """Pair i of every ``addrs`` := (v, tag) after one release fence."""
    release(_SCOPE[addrs[0].space])
    for a in addrs:
        self.put(a, i, v)

put_triple(a, i, bits48)

Pair i of a := (48 value bits | tag16 << 48), monotonic.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/sync.py 
284
285
286
287
def put_triple(self, a, i, bits48):
    """Pair i of ``a`` := (48 value bits | tag16 << 48), monotonic."""
    word = (bits48 & ((1 << 48) - 1)) | (fx.Int64(self.tag16) << 48)
    store64(a.value, i, word, _SCOPE[a.space])

_ptr(addr, i, bits=64)

The global pointer to element i (of bits) of address addr.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/sync.py 
55
56
57
58
59
60
61
62
63
def _ptr(addr, i, bits=64):
    """The global pointer to element i (of ``bits``) of address ``addr``."""
    at = addr + fx.Int64(i) * (bits // 8)
    return fx.inttoptr(
        fx.PointerType.get(
            T.i64 if bits == 64 else T.i32, fx.AddressSpace.Global, bits // 8
        ),
        at,
    )

_sync_fence(ordering, scope)

A fence that synchronizes without making anything available or visible.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/sync.py 
145
146
147
148
149
150
151
def _sync_fence(ordering, scope):
    """A fence that synchronizes without making anything available or
    visible."""
    fence = _llvm.FenceOp(ordering, syncscope=scope)
    fence.attributes["llvm.mmra"] = ir.Attribute.parse(
        '#llvm.mmra_tag<"amdgcn-av":"none">'
    )

cta_sync(drain=False)

gpu.barrier: workgroup release fence, s_barrier, acquire fence (HIP's __syncthreads). drain: every wave's stores complete first.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/sync.py 
195
196
197
198
199
200
def cta_sync(drain=False):
    """``gpu.barrier``: workgroup release fence, ``s_barrier``, acquire fence
    (HIP's ``__syncthreads``). ``drain``: every wave's stores complete first."""
    if const_expr(drain):
        rocdl.s_waitcnt(vmcnt=0)
    gpu.barrier()

load128(addr, i, scope)

Dwords i .. i + 3 (i a multiple of 4): one load-visible at scope.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/sync.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
def load128(addr, i, scope):
    """Dwords i .. i + 3 (i a multiple of 4): one load-visible at ``scope``."""
    ptr = _llvm.IntToPtrOp(
        ir.Type.parse("!llvm.ptr<1>"), as_ir_value(addr + fx.Int64(i) * 4)
    ).result
    md = _llvm.MetadataAsValueOp(
        ir.Attribute.parse(f'#llvm.md_node<#llvm.md_string<"{scope}">>'),
        results=[ir.Type.parse("!llvm.metadata")],
    ).result
    return fx.Vector(
        _llvm.CallIntrinsicOp(
            T.vec(4, T.i32), "llvm.amdgcn.av.load.b128.p1", [ptr, md], [], []
        ).results[0]
    )

load_dev(addr, i, n=1)

Dwords i .. i + n of this GPU's scratch (n in {1, 2, 4}, i aligned), load-visible at agent scope.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/sync.py 
134
135
136
137
138
139
140
141
142
def load_dev(addr, i, n=1):
    """Dwords i .. i + n of this GPU's scratch (n in {1, 2, 4}, i aligned),
    load-visible at agent scope."""
    if const_expr(n == 4):
        return load128(addr, i, "agent")
    if const_expr(n == 2):
        w = load64(addr + fx.Int64(i) * 4, 0, "agent")
        return fx.Vector.from_elements([fx.Int32(w), fx.Int32(w >> 32)], fx.Int32)
    return load32(addr, i, "agent")

store_dev(addr, i, words)

Dwords i .. of this GPU's scratch := words (1 or 2, i aligned), store-available at agent scope.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/sync.py 
125
126
127
128
129
130
131
def store_dev(addr, i, words):
    """Dwords i .. of this GPU's scratch := ``words`` (1 or 2, i aligned),
    store-available at agent scope."""
    if const_expr(len(words) == 1):
        store32(addr, i, words[0], "agent")
    else:
        store64(addr + fx.Int64(i) * 4, 0, pack(words[0], words[1]), "agent")

wave_sync(count, target)

Every lane of the wave: what each lane did before (its acquires) happens before what any lane does after. Each lane releases, adds one to the wave's LDS count, and acquires once it reads target; the read value comes after every lane's add in the counter's order (a release sequence), so it synchronizes with every lane's release.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/sync.py 
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
@traced
def wave_sync(count, target):
    """Every lane of the wave: what each lane did before (its acquires)
    happens before what any lane does after. Each lane releases, adds one to
    the wave's LDS ``count``, and acquires once it reads ``target``; the read
    value comes after every lane's add in the counter's order (a release
    sequence), so it synchronizes with every lane's release."""
    sync_release("wavefront")
    fx.llvm.atomic_add(count, fx.Int32(1), syncscope="wavefront")
    n = fx.generic_load(
        count,
        dtype=fx.Int32,
        memory_order=fx.AtomicOrdering.Monotonic,
        syncscope="wavefront",
    )
    while n < target:
        n = fx.generic_load(
            count,
            dtype=fx.Int32,
            memory_order=fx.AtomicOrdering.Monotonic,
            syncscope="wavefront",
        )
    sync_acquire("wavefront")