Skip to content

vllm.models.kimi_k3.nvidia.specialized_decode.cutedsl.common.mailbox

Lamport mailboxes: BF16 rows published in 16-byte fragments; a reader polls until no fragment holds the -0.0 sentinel, then re-arms.

Functions:

  • fetch –

    Load fragment i into values[i]; returns its dirty flag at bit i.

  • poll –

    Lamport-poll count 16-byte fragments at base + i*stride (elements); all

fetch(values, i, pointer)

Load fragment i into values[i]; returns its dirty flag at bit i.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/mailbox.py 
19
20
21
22
23
24
25
@cute.jit
def fetch(values, i: cutlass.Constexpr[int], pointer):
    """Load fragment ``i`` into ``values[i]``; returns its dirty flag at bit ``i``."""
    got = load_words(pointer, count=4, volatile=True)
    for k in cutlass.range_constexpr(4):
        values[i, k] = got[k]
    return Uint32(not published(got)) << i

poll(base, stride, count)

Lamport-poll count 16-byte fragments at base + i*stride (elements); all loads issue before any check and only dirty fragments re-poll.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/mailbox.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
@cute.jit
def poll(base, stride, count: cutlass.Constexpr[int]):
    """Lamport-poll ``count`` 16-byte fragments at base + i*stride (elements); all
    loads issue before any check and only dirty fragments re-poll."""
    values = cute.make_rmem_tensor((count, 4), Uint32)
    pending = Uint32(0)
    for i in cutlass.range_constexpr(count):
        pending = pending | fetch(values, i, base + i * stride)
    while pending != 0:
        again = Uint32(0)
        for i in cutlass.range_constexpr(count):
            if (pending >> i) & 1:
                again = again | fetch(values, i, base + i * stride)
        pending = again
    return values