vllm.models.muse_spark.nvidia.specialized_decode.cutedsl.gemm.unpooler

The Muse Spark unpooler GEMM for small M: this rank's column block y[M, N] = x[M, K] @ W[N, K]^T of the MoE unpooler (N = hidden / tp, K = 4096), BF16 in and out, FP32 accumulate, 1 <= M <= 16. An optional epilogue also publishes y into the next boundary's Lamport mailbox (zero-copy publish).

Grid (N / 16, S), cluster (1, S, 1) when S > 1. CTA (t, s) owns output features [16 t, 16 t + 16) and the K slice [s K / S, (s + 1) K / S); each of its warps owns K / (S * warps) consecutive K. Operands are loaded straight into registers and fed to mma.sync.m16n8k16 (A = 16 weight rows, B = 8 tokens).

Weights are loaded before griddepcontrol.wait, activations after it. Each warp's FP32 partial goes to cluster rank 0's shared memory, which sums the partials in a fixed order, rounds to BF16 and stores y. The summation order depends only on the launch shape, so runs are bitwise reproducible.

Zero-copy epilogue (zerocopy): rank 0 also multicasts y (-0.0 stored as +0.0) into this rank's slot of the next boundary's LamportMailbox at stage NEXT_STAGE and sets ACTIVE_STAGE, the same bytes and stage word the mailbox's own publisher would write.

Ticket stage (ticket_stage, with zerocopy): the stage is instead the upper 16 bits of the draft boundary's ticket word, and no stage word is stored; the consuming boundary advances the stage itself.

Classes:

• Unpooler –

  One MoE layer's unpooler column block, prepared at build (before CUDA

Functions:

• compile_unpooler –

  The compiled kernel for weight [n, k] and one token tile; called as

• launch_unsupported –

  Why the kernel has no launch for an [n, k] weight with K split over

Unpooler

One MoE layer's unpooler column block, prepared at build (before CUDA graph capture): the [N, K] BF16 weight shard (this rank's rows of the replicated unpooler weight) and the kernels of both token tiles compiled.

With mailbox (the next boundary's LamportMailbox, width N) the epilogue also publishes y into it (zero-copy); without, y is only stored and the boundary publishes it. With ticket (int32 [1], stage << 16: the draft boundary's ticket on that mailbox) the slot's stage is the ticket's, and the mailbox's own stage state is not used. With rows only the token tiles those row counts take are compiled at build (default: every row count); a call with another row count compiles its tile then, which must not happen under CUDA graph capture. run is graph-capturable and allocates nothing.

Methods:

• from_moe –

  The unpooler of a MuseSpark1xMoE: rows [rank N, (rank + 1) N) of its

• run –

  out[m, N] = x[m, K] @ W^T (and its publish with a mailbox).

• unsupported –

  Why the GEMM cannot take weight (an [N, K] shard) at launch

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/gemm/unpooler.py

class Unpooler:
    """One MoE layer's unpooler column block, prepared at build (before CUDA
    graph capture): the [N, K] BF16 weight shard (this rank's rows of the
    replicated unpooler weight) and the kernels of both token tiles compiled.

    With ``mailbox`` (the next boundary's LamportMailbox, width N) the epilogue
    also publishes y into it (zero-copy); without, y is only stored and the
    boundary publishes it. With ``ticket`` (int32 [1], stage << 16: the draft
    boundary's ticket on that mailbox) the slot's stage is the ticket's, and the
    mailbox's own stage state is not used. With ``rows`` only the token tiles
    those row counts take are compiled at build (default: every row count); a
    call with another row count compiles its tile then, which must not happen
    under CUDA graph capture. ``run`` is graph-capturable and allocates
    nothing.
    """

    def __init__(
        self,
        weight: torch.Tensor,
        mailbox: LamportMailbox | None = None,
        *,
        enable_pdl: bool = True,
        shape: UnpoolerShape = UNPOOLER,
        ticket: torch.Tensor | None = None,
        rows: Sequence[int] | None = None,
    ) -> None:
        reason = self.unsupported(weight, shape)
        if reason is not None:
            raise ValueError(reason)
        n, k = map(int, weight.shape)
        if weight.is_cuda and torch.cuda.get_device_capability(weight.device)[0] != 10:
            raise RuntimeError("the unpooler GEMM needs an SM100-family GPU")
        self.n, self.k, self.weight = n, k, weight
        self.mailbox = mailbox
        options: dict[str, Any] = dict(
            splits=shape.splits, kw=shape.kw, accs=shape.accs, enable_pdl=enable_pdl
        )
        if mailbox is not None:
            if mailbox.width != n:
                raise ValueError(f"mailbox width {mailbox.width} != unpooler N {n}")
            options.update(
                zerocopy=True,
                tp=mailbox.tp,
                rank=mailbox.rank,
                capacity_m=mailbox.capacity,
                hidden_local=mailbox.width,
            )
            self.max_m = min(MAX_M, mailbox.capacity)
            self.stage_state = mailbox.stage_state
            self.address = Int64(mailbox.multicast)
            if ticket is not None:
                if (
                    ticket.dtype != torch.int32
                    or tuple(ticket.shape) != (1,)
                    or ticket.device != mailbox.tensor.device
                ):
                    raise ValueError("the ticket must be int32 [1] on the mailbox")
                options.update(ticket_stage=True)
                self.stage_state = ticket
        elif ticket is not None:
            raise ValueError("a ticket stage needs a mailbox")
        else:
            self.max_m = MAX_M
            # Never read: the plain kernel takes no stage.
            self.stage_state = torch.zeros(2, dtype=torch.int32, device=weight.device)
            self.address = Int64(0)
        self.token_tile = {
            m: tile for m, tile in shape.token_tile.items() if 1 <= m <= self.max_m
        }
        if sorted(self.token_tile) != list(range(1, self.max_m + 1)) or any(
            tile not in TOKEN_TILES or tile < m for m, tile in self.token_tile.items()
        ):
            raise ValueError(
                f"token tiles {shape.token_tile} do not cover M 1..{self.max_m}"
            )
        if rows is None:
            rows = tuple(self.token_tile)
        elif any(m not in self.token_tile for m in rows):
            raise ValueError(f"rows {tuple(rows)} outside 1..{self.max_m}")
        self._options = options
        self.kernels = {
            tile: compile_unpooler(n, k, tile, **options)
            for tile in sorted({self.token_tile[m] for m in rows})
        }
        self.weight_view = to_cute(weight.view(-1), 16)
        self.stage_view = to_cute(self.stage_state, 4)

    @staticmethod
    def unsupported(
        weight: torch.Tensor, shape: UnpoolerShape = UNPOOLER
    ) -> str | None:
        """Why the GEMM cannot take ``weight`` (an [N, K] shard) at launch
        shape ``shape``, or None: the constructor's checks of both, which it
        raises as ValueError."""
        if weight.dim() != 2:
            return f"unpooler shard must be 2-D, got {tuple(weight.shape)}"
        n, k = map(int, weight.shape)
        if (
            weight.dtype != torch.bfloat16
            or not weight.is_contiguous()
            or weight.data_ptr() % 16
            or n % TILE_N
        ):
            return (
                f"unpooler shard {tuple(weight.shape)} {weight.dtype}: needs a "
                f"contiguous, 16-byte aligned BF16 [N, K] with N % {TILE_N} == 0"
            )
        return launch_unsupported(n, k, shape.splits, shape.kw, shape.accs)

    @classmethod
    def from_moe(
        cls,
        moe: torch.nn.Module,
        tp: int,
        rank: int,
        mailbox: LamportMailbox | None = None,
        **options: Any,
    ) -> "Unpooler":
        """The unpooler of a MuseSpark1xMoE: rows [rank N, (rank + 1) N) of its
        replicated [hidden, K] weight, N = hidden / tp (as the stock
        _project_local_columns takes them)."""
        weight = moe.unpooler.weight
        n = weight.shape[0] // tp
        return cls(weight.narrow(0, rank * n, n), mailbox, **options)

    def run(self, x: torch.Tensor, out: torch.Tensor) -> torch.Tensor:
        """out[m, N] = x[m, K] @ W^T (and its publish with a mailbox)."""
        m = int(x.shape[0])
        if not 1 <= m <= self.max_m or tuple(x.shape) != (m, self.k):
            raise ValueError(f"x {tuple(x.shape)}: needs [1..{self.max_m}, {self.k}]")
        if tuple(out.shape) != (m, self.n):
            raise ValueError(f"out {tuple(out.shape)} != ({m}, {self.n})")
        for tensor in (x, out):
            if (
                tensor.dtype != torch.bfloat16
                or not tensor.is_contiguous()
                or tensor.data_ptr() % 16
            ):
                raise ValueError("x and out must be contiguous 16-byte aligned BF16")
        tile = self.token_tile[m]
        kernel = self.kernels.get(tile)
        if kernel is None:
            kernel = compile_unpooler(self.n, self.k, tile, **self._options)
            self.kernels[tile] = kernel
        kernel(
            self.weight_view,
            to_cute_dynamic(x.view(-1), 16, divisibility=self.k),
            to_cute_dynamic(out.view(-1), 16, divisibility=self.n),
            self.stage_view,
            self.address,
            Int32(m),
            current_cu_stream(),
        )
        return out

from_moe(moe, tp, rank, mailbox=None, **options) classmethod

The unpooler of a MuseSpark1xMoE: rows [rank N, (rank + 1) N) of its replicated [hidden, K] weight, N = hidden / tp (as the stock _project_local_columns takes them).

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/gemm/unpooler.py

@classmethod
def from_moe(
    cls,
    moe: torch.nn.Module,
    tp: int,
    rank: int,
    mailbox: LamportMailbox | None = None,
    **options: Any,
) -> "Unpooler":
    """The unpooler of a MuseSpark1xMoE: rows [rank N, (rank + 1) N) of its
    replicated [hidden, K] weight, N = hidden / tp (as the stock
    _project_local_columns takes them)."""
    weight = moe.unpooler.weight
    n = weight.shape[0] // tp
    return cls(weight.narrow(0, rank * n, n), mailbox, **options)

run(x, out)

out[m, N] = x[m, K] @ W^T (and its publish with a mailbox).

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/gemm/unpooler.py

def run(self, x: torch.Tensor, out: torch.Tensor) -> torch.Tensor:
    """out[m, N] = x[m, K] @ W^T (and its publish with a mailbox)."""
    m = int(x.shape[0])
    if not 1 <= m <= self.max_m or tuple(x.shape) != (m, self.k):
        raise ValueError(f"x {tuple(x.shape)}: needs [1..{self.max_m}, {self.k}]")
    if tuple(out.shape) != (m, self.n):
        raise ValueError(f"out {tuple(out.shape)} != ({m}, {self.n})")
    for tensor in (x, out):
        if (
            tensor.dtype != torch.bfloat16
            or not tensor.is_contiguous()
            or tensor.data_ptr() % 16
        ):
            raise ValueError("x and out must be contiguous 16-byte aligned BF16")
    tile = self.token_tile[m]
    kernel = self.kernels.get(tile)
    if kernel is None:
        kernel = compile_unpooler(self.n, self.k, tile, **self._options)
        self.kernels[tile] = kernel
    kernel(
        self.weight_view,
        to_cute_dynamic(x.view(-1), 16, divisibility=self.k),
        to_cute_dynamic(out.view(-1), 16, divisibility=self.n),
        self.stage_view,
        self.address,
        Int32(m),
        current_cu_stream(),
    )
    return out

unsupported(weight, shape=UNPOOLER) staticmethod

Why the GEMM cannot take weight (an [N, K] shard) at launch shape shape, or None: the constructor's checks of both, which it raises as ValueError.

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/gemm/unpooler.py

@staticmethod
def unsupported(
    weight: torch.Tensor, shape: UnpoolerShape = UNPOOLER
) -> str | None:
    """Why the GEMM cannot take ``weight`` (an [N, K] shard) at launch
    shape ``shape``, or None: the constructor's checks of both, which it
    raises as ValueError."""
    if weight.dim() != 2:
        return f"unpooler shard must be 2-D, got {tuple(weight.shape)}"
    n, k = map(int, weight.shape)
    if (
        weight.dtype != torch.bfloat16
        or not weight.is_contiguous()
        or weight.data_ptr() % 16
        or n % TILE_N
    ):
        return (
            f"unpooler shard {tuple(weight.shape)} {weight.dtype}: needs a "
            f"contiguous, 16-byte aligned BF16 [N, K] with N % {TILE_N} == 0"
        )
    return launch_unsupported(n, k, shape.splits, shape.kw, shape.accs)

compile_unpooler(n, k, token_tile, *, splits, kw, accs, enable_pdl=True, zerocopy=False, tp=1, rank=0, capacity_m=1, hidden_local=0, ticket_stage=False) cached

The compiled kernel for weight [n, k] and one token tile; called as fn(weight, x, out, stage_state, mailbox_multicast_address, m, stream) with flat views (Unpooler binds them; stage_state is the ticket word with ticket_stage). Shared by every layer of one shape.

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/gemm/unpooler.py

@cache
def compile_unpooler(
    n: int,
    k: int,
    token_tile: int,
    *,
    splits: int,
    kw: int,
    accs: int,
    enable_pdl: bool = True,
    zerocopy: bool = False,
    tp: int = 1,
    rank: int = 0,
    capacity_m: int = 1,
    hidden_local: int = 0,
    ticket_stage: bool = False,
):
    """The compiled kernel for weight [n, k] and one token tile; called as
    fn(weight, x, out, stage_state, mailbox_multicast_address, m, stream) with
    flat views (``Unpooler`` binds them; ``stage_state`` is the ticket word
    with ``ticket_stage``). Shared by every layer of one shape."""
    assert_not_capturing("the unpooler GEMM")
    kernel = UnpoolerKernel(
        n=n,
        k=k,
        splits=splits,
        token_tile=token_tile,
        kw=kw,
        accs=accs,
        enable_pdl=enable_pdl,
        zerocopy=zerocopy,
        tp=tp,
        rank=rank,
        capacity_m=capacity_m,
        hidden_local=hidden_local,
        ticket_stage=ticket_stage,
    )
    return cute.compile(
        kernel,
        make_fake_compact_tensor(BFloat16, (n * k,), assumed_align=16),
        make_fake_dynamic_compact_tensor(BFloat16, alignment=16, divisibility=k),
        make_fake_dynamic_compact_tensor(BFloat16, alignment=16, divisibility=n),
        make_fake_compact_tensor(Int32, (1 if ticket_stage else 2,), assumed_align=4),
        Int64(0),
        Int32(1),
        current_cu_stream(),
    )

launch_unsupported(n, k, splits, kw, accs)

Why the kernel has no launch for an [n, k] weight with K split over splits CTAs, kw K per warp and accs accumulators, or None.

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/gemm/unpooler.py

def launch_unsupported(n: int, k: int, splits: int, kw: int, accs: int) -> str | None:
    """Why the kernel has no launch for an [n, k] weight with K split over
    ``splits`` CTAs, ``kw`` K per warp and ``accs`` accumulators, or None."""
    if n <= 0 or n % TILE_N or splits not in (1, 2, 4, 8) or accs not in (1, 2, 4):
        return f"unsupported shape or tactic n={n} splits={splits} accs={accs}"
    if k % (splits * kw) or kw % BLOCK_K:
        return f"K={k} does not split into {splits} x warps x {kw}"
    if not 1 <= k // (splits * kw) <= 32:
        return f"{k // (splits * kw)} warps per CTA"
    return None