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

Lamport mailboxes in symmetric memory.

A mailbox is [LAMPORT_GENERATIONS, tp, capacity, width] BF16 on every rank of a TP group, mapped for NVLink multicast. A publisher multicasts its [m, width] rows into slot [stage, rank] of every rank's copy; readers poll their own copy in 16-byte fragments until no fragment holds BF16 negative zero (the sentinel; publishers store -0.0 as +0.0), then re-arm what they read. A publisher can also store rows into one rank's copy alone through peers (the TP tail's route rows); poll_at reads slots at any fixed offsets.

Every spin is bounded: a poll that sees no data for SPIN_TIMEOUT_NS traps, so a rank that diverges fails the launch instead of hanging the group.

Classes:

• LamportMailbox –

  One TP group's [3, tp, capacity, width] BF16 mailbox and stage state.

Functions:

• compile_shared_only_publish –

  Compile the publisher that multicasts a local [m, hidden] BF16 tensor

• fetch –

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

• poll –

  count fragments at base + i * stride (elements, stride at run

• poll_at –

  The fragments at base + offsets[i] (elements, compile-time), polled as

• poll_owned_fragment –

  Column all-gather read: wait for this token's 8 global columns

• rearm –

  Store the sentinel over one 16-byte fragment.

• spin_deadline –

  Call once per poll that found no data: arms the deadline on the first

LamportMailbox

One TP group's [3, tp, capacity, width] BF16 mailbox and stage state.

Construction is collective over group and raises on every rank when the buffer has no multicast mapping. peers are the ranks' buffer addresses in rank order and multicast the multicast address of the whole buffer.

Methods:

• compile_publish –

  The shared-only publisher of this mailbox (call before capture); one

• publish –

  Multicast rows [m, width] BF16 into slot [NEXT_STAGE, rank].

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/common/mailbox.py

class LamportMailbox:
    """One TP group's ``[3, tp, capacity, width]`` BF16 mailbox and stage state.

    Construction is collective over ``group`` and raises on every rank when the
    buffer has no multicast mapping. ``peers`` are the ranks' buffer addresses
    in rank order and ``multicast`` the multicast address of the whole buffer.
    """

    def __init__(
        self,
        group: dist.ProcessGroup,
        *,
        tp: int,
        rank: int,
        capacity: int,
        width: int,
        device: torch.device,
    ) -> None:
        if dist.get_world_size(group) != tp or dist.get_rank(group) != rank:
            raise ValueError("process group does not match tp and rank")
        if width % VEC_BF16:
            raise ValueError("width must be a multiple of 8 BF16 elements")
        self.tp, self.rank, self.capacity, self.width = tp, rank, capacity, width
        self.tensor = symm_mem.empty(
            (LAMPORT_GENERATIONS, tp, capacity, width),
            dtype=torch.bfloat16,
            device=device,
        )
        handle = symm_mem.rendezvous(self.tensor, group)
        self.multicast = int(handle.multicast_ptr or 0)
        if not self.multicast or self.multicast % 16:
            raise RuntimeError(
                "Muse Spark specialized decode needs NVLink multicast memory"
            )
        self.peers = tuple(int(pointer) for pointer in handle.buffer_ptrs)
        # Every slot starts as the sentinel, BF16 negative zero.
        self.tensor.view(torch.int16).fill_(-32768)
        self.stage_state = torch.zeros(2, dtype=torch.int32, device=device)
        torch.accelerator.synchronize(device)
        dist.barrier(group=group)
        self.view = to_cute(self.tensor.flatten(), 16)
        self.stage_view = to_cute(self.stage_state, 4)

    def compile_publish(
        self,
        *,
        threads: int = 128,
        release_before_store: bool = False,
        enable_pdl: bool = True,
    ):
        """The shared-only publisher of this mailbox (call before capture); one
        compile per launch shape, shared by every mailbox of that shape."""
        return _cached_publish(
            hidden=self.width,
            tp=self.tp,
            rank=self.rank,
            capacity_m=self.capacity,
            elements_per_thread=VEC_BF16,
            threads=threads,
            release_before_store=release_before_store,
            enable_pdl=enable_pdl,
        )

    def publish(self, publisher, rows: torch.Tensor) -> None:
        """Multicast ``rows`` [m, width] BF16 into slot [NEXT_STAGE, rank]."""
        publisher(
            to_cute_dynamic(rows.flatten(), 16, divisibility=self.width),
            self.stage_view,
            Int64(self.multicast),
            Int32(rows.shape[0]),
            current_cu_stream(),
        )

compile_publish(*, threads=128, release_before_store=False, enable_pdl=True)

The shared-only publisher of this mailbox (call before capture); one compile per launch shape, shared by every mailbox of that shape.

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/common/mailbox.py

def compile_publish(
    self,
    *,
    threads: int = 128,
    release_before_store: bool = False,
    enable_pdl: bool = True,
):
    """The shared-only publisher of this mailbox (call before capture); one
    compile per launch shape, shared by every mailbox of that shape."""
    return _cached_publish(
        hidden=self.width,
        tp=self.tp,
        rank=self.rank,
        capacity_m=self.capacity,
        elements_per_thread=VEC_BF16,
        threads=threads,
        release_before_store=release_before_store,
        enable_pdl=enable_pdl,
    )

publish(publisher, rows)

Multicast rows [m, width] BF16 into slot [NEXT_STAGE, rank].

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/common/mailbox.py

def publish(self, publisher, rows: torch.Tensor) -> None:
    """Multicast ``rows`` [m, width] BF16 into slot [NEXT_STAGE, rank]."""
    publisher(
        to_cute_dynamic(rows.flatten(), 16, divisibility=self.width),
        self.stage_view,
        Int64(self.multicast),
        Int32(rows.shape[0]),
        current_cu_stream(),
    )

_poll(base, stride, offsets, timeout_ns)

Lamport-poll the 16-byte fragments at base + offsets[i] * stride (elements) into a [len(offsets), 4] u32 register tensor. All loads issue before any check and only dirty fragments re-poll; traps after timeout_ns without data.

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/common/mailbox.py

@cute.jit
def _poll(
    base: cute.Pointer,
    stride: Int64,
    offsets: cutlass.Constexpr[tuple],
    timeout_ns: cutlass.Constexpr[int],
):
    """Lamport-poll the 16-byte fragments at base + offsets[i] * stride
    (elements) into a [len(offsets), 4] u32 register tensor. All loads issue
    before any check and only dirty fragments re-poll; traps after
    ``timeout_ns`` without data."""
    count = len(offsets)
    values = cute.make_rmem_tensor(cute.make_layout((count, 4)), Uint32)
    pending = Uint32(0)
    for i in cutlass.range_constexpr(count):
        pending = pending | fetch(values, i, base + Int64(offsets[i]) * stride)
    deadline = Int64(0)
    while pending != Uint32(0):
        deadline = spin_deadline(deadline, timeout_ns)
        again = Uint32(0)
        for i in cutlass.range_constexpr(count):
            if ((pending >> Uint32(i)) & Uint32(1)) != Uint32(0):
                again = again | fetch(values, i, base + Int64(offsets[i]) * stride)
        pending = again
    return values

compile_shared_only_publish(*, hidden, tp, rank, capacity_m, elements_per_thread, threads, release_before_store, enable_pdl)

Compile the publisher that multicasts a local [m, hidden] BF16 tensor into slot [NEXT_STAGE, rank] and then sets ACTIVE_STAGE.

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/common/mailbox.py

def compile_shared_only_publish(
    *,
    hidden: int,
    tp: int,
    rank: int,
    capacity_m: int,
    elements_per_thread: int,
    threads: int,
    release_before_store: bool,
    enable_pdl: bool,
):
    """Compile the publisher that multicasts a local ``[m, hidden]`` BF16 tensor
    into slot [NEXT_STAGE, rank] and then sets ACTIVE_STAGE."""
    device_kernel = _SharedOnlyPublishDeviceKernel(
        hidden=hidden,
        tp=tp,
        rank=rank,
        capacity_m=capacity_m,
        elements_per_thread=elements_per_thread,
        threads=threads,
        release_before_store=release_before_store,
        enable_pdl=enable_pdl,
    )
    return cute.compile(
        device_kernel,
        make_fake_dynamic_compact_tensor(BFloat16, alignment=16, divisibility=hidden),
        make_fake_compact_tensor(Int32, (2,), assumed_align=4),
        Int64(0),
        Int32(capacity_m),
        current_cu_stream(),
    )

fetch(values, i, pointer)

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

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/common/mailbox.py

@cute.jit
def fetch(values: cute.Tensor, i: cutlass.Constexpr[int], pointer: cute.Pointer):
    """Load fragment ``i`` into ``values[i]``; returns its dirty flag at bit ``i``."""
    got = load_global_u32x4(pointer, volatile=True)
    for k in cutlass.range_constexpr(4):
        values[i, k] = got[k]
    dirty = Uint32(0)
    if fragment_has_negative_zero(got):
        dirty = Uint32(1 << i)
    return dirty

poll(base, stride, count, timeout_ns=SPIN_TIMEOUT_NS)

count fragments at base + i * stride (elements, stride at run time), polled as _poll does: a [count, 4] u32 register tensor.

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/common/mailbox.py

def poll(
    base: cute.Pointer,
    stride: Int64,
    count: int,
    timeout_ns: int = SPIN_TIMEOUT_NS,
):
    """``count`` fragments at base + i * stride (elements, ``stride`` at run
    time), polled as ``_poll`` does: a [count, 4] u32 register tensor."""
    return _poll(base, stride, tuple(range(count)), timeout_ns)

poll_at(base, offsets, timeout_ns=SPIN_TIMEOUT_NS)

The fragments at base + offsets[i] (elements, compile-time), polled as _poll does: a [len(offsets), 4] u32 register tensor.

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/common/mailbox.py

def poll_at(
    base: cute.Pointer,
    offsets: tuple,
    timeout_ns: int = SPIN_TIMEOUT_NS,
):
    """The fragments at base + offsets[i] (elements, compile-time), polled as
    ``_poll`` does: a [len(offsets), 4] u32 register tensor."""
    return _poll(base, Int64(1), offsets, timeout_ns)

poll_owned_fragment(mailbox, stage_state, packed, token, fragment, tp, capacity_m, hidden_local, release_dependents, timeout_ns=SPIN_TIMEOUT_NS)

Column all-gather read: wait for this token's 8 global columns fragment from the one rank that owns them, store the BF16x8 fragment into packed and re-arm the slot; one thread per launch advances the stage. Releases PDL dependents after the poll when release_dependents (a kernel that triggers earlier passes False).

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/common/mailbox.py

@cute.jit
def poll_owned_fragment(
    mailbox: cute.Tensor,
    stage_state: cute.Tensor,
    packed: cute.Tensor,
    token: Int32,
    fragment: Int32,
    tp: cutlass.Constexpr[int],
    capacity_m: cutlass.Constexpr[int],
    hidden_local: cutlass.Constexpr[int],
    release_dependents: cutlass.Constexpr[bool],
    timeout_ns: cutlass.Constexpr[int] = SPIN_TIMEOUT_NS,
) -> None:
    """Column all-gather read: wait for this token's 8 global columns
    ``fragment`` from the one rank that owns them, store the BF16x8 fragment
    into ``packed`` and re-arm the slot; one thread per launch advances the
    stage. Releases PDL dependents after the poll when ``release_dependents``
    (a kernel that triggers earlier passes False)."""
    local_fragments = hidden_local // VEC_BF16
    active_stage = load_volatile_u32(stage_state.iterator + ACTIVE_STAGE)
    source_rank = fragment // local_fragments
    local_fragment = fragment % local_fragments
    source_element = (
        (Int64(active_stage) * tp + Int64(source_rank)) * capacity_m + Int64(token)
    ) * hidden_local + Int64(local_fragment) * VEC_BF16
    source_pointer = cute.make_ptr(
        BFloat16,
        (mailbox.iterator + source_element).llvm_ptr,
        cute.AddressSpace.gmem,
        assumed_align=16,
    )
    deadline = Int64(0)
    # Always true; a runtime value keeps the poll a device loop.
    dirty = fragment >= Int32(0)
    while dirty:
        value = load_global_u32x4(source_pointer, volatile=True)
        dirty = fragment_has_negative_zero(value)
        packed.store(value)
        if dirty:
            deadline = spin_deadline(deadline, timeout_ns)

    if token == 0 and fragment == 0:
        store_global_u32(
            stage_state.iterator + NEXT_STAGE,
            (active_stage + Uint32(1)) % Uint32(LAMPORT_GENERATIONS),
        )
    if cutlass.const_expr(release_dependents):
        cute.arch.griddepcontrol_launch_dependents()
    store_lamport_sentinel_u32x4(Int64(source_pointer.toint()))

rearm(pointer)

Store the sentinel over one 16-byte fragment.

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/common/mailbox.py

@cute.jit
def rearm(pointer: cute.Pointer) -> None:
    """Store the sentinel over one 16-byte fragment."""
    store_lamport_sentinel_u32x4(Int64(pointer.toint()))

spin_deadline(deadline, timeout_ns)

Call once per poll that found no data: arms the deadline on the first call (zero means unarmed) and traps once it has passed.

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/common/mailbox.py

@cute.jit
def spin_deadline(deadline: Int64, timeout_ns: cutlass.Constexpr[int]) -> Int64:
    """Call once per poll that found no data: arms the deadline on the first
    call (zero means unarmed) and traps once it has passed."""
    now = global_timer_ns()
    if deadline == Int64(0):
        deadline = now + Int64(timeout_ns)
    elif now > deadline:
        trap()
    return deadline