vllm.models.kimi_k3.nvidia.specialized_decode.cutedsl.moe.experts_runtime

Persistent routed experts for B<=WIDE_ROWS (B>8: grouped units of <=8 tokens per expert; B>64: the wide grouped engine, always pregrouped, over its own wide workspace) over one shared core and model-owned workspace. The entry wait completes the predecessor (incl. workspace rearm/counters) before shared state is read; output rows keep the caller's top16 order.

Classes:

• PersistentExperts –

Functions:

• compile_experts –

  Only code and geometry are cached; no owner tensors or descriptors. knobs:

• compile_group –

  The pregrouped experts' unit-table kernel: (ids, workspace, batch); early:

• compile_plan –

  Plan writer for callers without the routing kernel's copy role in front.

• dynamic –

  Static B<=8: FC2 tile-pair jobs claimed from a grid-wide counter

• early_group –

  Pregrouped batches whose unit-table kernel releases the experts at its entry.

• pregrouped –

  B>8 experts read the unit table a one-CTA kernel builds after routing, beside

• scratch –

  Model-owned expert buffers: the workspace and the grouped plan (MXFP8 x, written

• static –

  Streamed B<=8 with TP8 FC2 pairs: the MMA warp walks the producer's static

• stealing –

  B>8 experts balance the FC2(B) tail by claims.

• streamed –

  B<=8 on the incremental producer: ~1-1.5% at B4-8, neutral at B1-2, which keep

PersistentExperts

Methods:

• group –

  Pregrouped batches prepared early: the unit-table kernel (run it on the

• prepare –

  publish: (multicast, parity, rank) of the WaveMoeTail latent slot to

• run –

  ahead: the next weights' [PREFETCH_ROWS, 2] (address, bytes) table, for

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py

class PersistentExperts:
    def __init__(self, w13, s13, w2, s2, *, beta=4.0, linear_beta=25.0):
        experts, rows, _ = w13.shape
        self.tp, self.grid = (8, 136) if rows == 768 else (16, 152)
        self.device, self.experts = w13.device, experts
        self.weights = (w13, s13, w2, s2)
        # TP8: B3-4 on 148 CTAs; B5-8 publish on 128-136 CTAs, leaving SMs for the
        # one-wave tail; grouped B9-64 on 152.
        grids = "3-4:148,5-6:128,7:132,9-64:152" if self.tp == 8 else ""
        grids += f",65-{WIDE_ROWS}:152"  # wide: every SM
        self.grids = batch_table(grids)
        scalars = (beta, linear_beta, 1 / beta, 1 / linear_beta)
        self.scalars = tuple(map(c.Float32, scalars))
        # The experts kernel uses maps 0-1 (L2 promotion on w13); the TP8 B1 cluster
        # experts use maps 2-4 (w13, w2 two-chunk and one-chunk boxes, promoted); TP8
        # FC2 pairs (B4-64) use map 5 (w2 chunk 2 of 256 rows).
        w2_tile = (192, 3, 64) if self.tp == 8 else (128, 2, 32)
        specs = [(w13, rows, 1792, 28, 64, 2, True), (w2, 3584, *w2_tile, 2, False)]
        if self.tp == 8:
            specs += [specs[0]] + [(w2, 3584, 192, 3, 64, box, True) for box in (2, 1)]
        dtype = cuda.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B
        swizzle = cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B
        l2 = cuda.CUtensorMapL2promotion
        l2 = (l2.CU_TENSOR_MAP_L2_PROMOTION_NONE, l2.CU_TENSOR_MAP_L2_PROMOTION_L2_256B)
        maps = [
            (weight.data_ptr(), dtype, (128, experts * n, chunks), (row_bytes, stride))
            + ((128, 128, box), swizzle, l2[promote])
            for weight, n, row_bytes, chunks, stride, box, promote in specs
        ]
        if self.tp == 8:
            shape = (128, experts * 3584, 3), (192, 64), (128, 256, 1)
            maps.append((w2.data_ptr(), dtype, *shape, swizzle, l2[False]))
        self.maps = tensor_maps(maps, self.device)
        scales = (tensor(s.view(torch.uint8)) for s in (s13, s2))
        self.views = (to_cute(self.maps[:256], 64), *scales)
        self.all_maps = to_cute(self.maps, 64)  # TP8 FC2 pairs: maps 0-5
        self.prepared, self.ahead = {}, {}

    def prepare(
        self, batch, buffers, out, publish=(), planned=False, twoshot=False,
        early=False, ahead=False, listed=False,
    ):  # fmt: skip
        """``publish``: (multicast, parity, rank) of the WaveMoeTail latent slot to
        publish (``twoshot``: the ``common.twoshot`` push box); ``planned``: B>8 calls
        find the plan written (by the routing kernel), else run writes it from x first;
        ``early``: the caller runs ``group`` (pregrouped batches) before ``run``;
        ``listed`` (wide): the routing kernel appended the routing lists
        (``buffers["lists"]``, passed to ``run`` as ``ids``) and wrote the plan; no
        unit-table kernel runs."""
        mode = "twoshot" if twoshot else bool(publish)
        grouped = pregrouped(self.tp, batch)
        listed = listed and batch > 64
        args: tuple = (batch, self.grid_for(batch), self.device.index, mode)
        ahead = ahead and (batch > 8 or streamed(batch))  # the stream() producer
        args += (streamed(batch), grouped, stealing(batch), listed, ahead)
        args += (static(batch), dynamic(batch))
        code = compile_experts(self.tp, self.experts, *args)
        # CTAs wait on other clusters' partials: the whole grid must be resident.
        check_residency(code, self.device.index, args[1], 384, "Kimi K3 MoE experts")
        ws, maps = tensor(buffers["workspace"]), self.views[0]
        plan: tuple | None = None
        capacity = 1 << (batch - 1).bit_length()
        if fc2_pairs(self.tp, capacity, batch > 8 or streamed(batch)):
            maps = self.all_maps
        group = None
        if batch > 8:
            need = plan_bytes(plan_rows(capacity))
            if buffers["plan"].numel() < need:
                raise ValueError(f"B{batch} needs the wide expert scratch")
            plan = (tensor(buffers["plan"]), compile_plan(batch, self.device.index))
            plan = plan[:1] if planned else plan
        if grouped and not listed:
            release = early_group(self.tp, batch)
            grouping = compile_group(self.tp, self.experts, *args[:3], release)
            group = grouping, ws, early
        self.prepared[batch] = (out, tensor(out), ws, code, publish, plan, maps, group)
        self.ahead[batch] = ahead

    def group(self, ids):
        """Pregrouped batches prepared ``early``: the unit-table kernel (run it on the
        stream after routing; ``run`` then skips it)."""
        group = self.prepared[ids.shape[0]][7]
        if group is not None and group[2]:
            self.launch_group(ids, group)

    def launch_group(self, ids, group):
        code, ws, _ = group
        code(tensor(ids), ws, c.Int32(ids.shape[0]), stream(self.device))

    def grid_for(self, batch):
        return self.grids.get(batch, self.grid)

    def run(self, x, ids, weights=None, ahead=None):
        """``ahead``: the next weights' [PREFETCH_ROWS, 2] (address, bytes) table, for
        batches prepared ``ahead``."""
        batch = x.shape[0]
        out, output, workspace, code, box, plan, maps, group = self.prepared[batch]
        output = tensor(weights) if box else output
        source = to_cute(x, 16)
        if plan is not None:
            if len(plan) > 1:
                plan[1](source, plan[0], stream(self.device))
            source = plan[0]
        if group is not None and not group[2]:
            self.launch_group(ids, group)
        args = (source, tensor(ids), maps, *self.views[1:], workspace, output)
        tail = (to_cute(ahead, 16),) if self.ahead.get(batch) else ()
        code(*args, c.Int32(batch), *self.scalars, *box, *tail, stream(self.device))
        return out

group(ids)

Pregrouped batches prepared early: the unit-table kernel (run it on the stream after routing; run then skips it).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py

def group(self, ids):
    """Pregrouped batches prepared ``early``: the unit-table kernel (run it on the
    stream after routing; ``run`` then skips it)."""
    group = self.prepared[ids.shape[0]][7]
    if group is not None and group[2]:
        self.launch_group(ids, group)

prepare(batch, buffers, out, publish=(), planned=False, twoshot=False, early=False, ahead=False, listed=False)

publish: (multicast, parity, rank) of the WaveMoeTail latent slot to publish (twoshot: the common.twoshot push box); planned: B>8 calls find the plan written (by the routing kernel), else run writes it from x first; early: the caller runs group (pregrouped batches) before run; listed (wide): the routing kernel appended the routing lists (buffers["lists"], passed to run as ids) and wrote the plan; no unit-table kernel runs.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py

def prepare(
    self, batch, buffers, out, publish=(), planned=False, twoshot=False,
    early=False, ahead=False, listed=False,
):  # fmt: skip
    """``publish``: (multicast, parity, rank) of the WaveMoeTail latent slot to
    publish (``twoshot``: the ``common.twoshot`` push box); ``planned``: B>8 calls
    find the plan written (by the routing kernel), else run writes it from x first;
    ``early``: the caller runs ``group`` (pregrouped batches) before ``run``;
    ``listed`` (wide): the routing kernel appended the routing lists
    (``buffers["lists"]``, passed to ``run`` as ``ids``) and wrote the plan; no
    unit-table kernel runs."""
    mode = "twoshot" if twoshot else bool(publish)
    grouped = pregrouped(self.tp, batch)
    listed = listed and batch > 64
    args: tuple = (batch, self.grid_for(batch), self.device.index, mode)
    ahead = ahead and (batch > 8 or streamed(batch))  # the stream() producer
    args += (streamed(batch), grouped, stealing(batch), listed, ahead)
    args += (static(batch), dynamic(batch))
    code = compile_experts(self.tp, self.experts, *args)
    # CTAs wait on other clusters' partials: the whole grid must be resident.
    check_residency(code, self.device.index, args[1], 384, "Kimi K3 MoE experts")
    ws, maps = tensor(buffers["workspace"]), self.views[0]
    plan: tuple | None = None
    capacity = 1 << (batch - 1).bit_length()
    if fc2_pairs(self.tp, capacity, batch > 8 or streamed(batch)):
        maps = self.all_maps
    group = None
    if batch > 8:
        need = plan_bytes(plan_rows(capacity))
        if buffers["plan"].numel() < need:
            raise ValueError(f"B{batch} needs the wide expert scratch")
        plan = (tensor(buffers["plan"]), compile_plan(batch, self.device.index))
        plan = plan[:1] if planned else plan
    if grouped and not listed:
        release = early_group(self.tp, batch)
        grouping = compile_group(self.tp, self.experts, *args[:3], release)
        group = grouping, ws, early
    self.prepared[batch] = (out, tensor(out), ws, code, publish, plan, maps, group)
    self.ahead[batch] = ahead

run(x, ids, weights=None, ahead=None)

ahead: the next weights' [PREFETCH_ROWS, 2] (address, bytes) table, for batches prepared ahead.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py

def run(self, x, ids, weights=None, ahead=None):
    """``ahead``: the next weights' [PREFETCH_ROWS, 2] (address, bytes) table, for
    batches prepared ``ahead``."""
    batch = x.shape[0]
    out, output, workspace, code, box, plan, maps, group = self.prepared[batch]
    output = tensor(weights) if box else output
    source = to_cute(x, 16)
    if plan is not None:
        if len(plan) > 1:
            plan[1](source, plan[0], stream(self.device))
        source = plan[0]
    if group is not None and not group[2]:
        self.launch_group(ids, group)
    args = (source, tensor(ids), maps, *self.views[1:], workspace, output)
    tail = (to_cute(ahead, 16),) if self.ahead.get(batch) else ()
    code(*args, c.Int32(batch), *self.scalars, *box, *tail, stream(self.device))
    return out

compile_experts(tp, experts, batch, grid, device, publish=False, *knobs) cached

Only code and geometry are cached; no owner tensors or descriptors. knobs: (B<=8 stream, pregrouped, steal, listed, ahead, static, dynamic).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py

@functools.cache
def compile_experts(tp, experts, batch, grid, device, publish=False, *knobs):
    """Only code and geometry are cached; no owner tensors or descriptors. ``knobs``:
    (B<=8 stream, pregrouped, steal, listed, ahead, static, dynamic)."""
    defaults: tuple = (streamed(batch), False, False, False, False, static(batch))
    defaults += (dynamic(batch),)
    stream, grouped, steal, listed, ahead, walk, claims = knobs + defaults[len(knobs) :]
    capacity = 1 << (batch - 1).bit_length()
    moe = ExpertsKernel(
        tp, capacity, experts, grid, publish, stream, grouped, steal, listed, ahead,
        walk, claims,
    )  # fmt: skip
    flat = lambda n, dtype=c.Uint8, align=16: fake((n,), (1,), dtype, align)
    box = (c.Int64(0), c.Int32(0), c.Int32(0)) if publish else ()
    if publish == "twoshot":  # receive slots, owners' bases, parity x TP + rank
        box = (fake((3 * tp, 64, 3584 // tp)), fake((tp,), None, c.Int64), c.Int32(0))
    x, ids = fake((batch, 3584), (3584, 1)), flat(batch * 16, c.Int32)
    ids = flat(LIST_WORDS, c.Int32) if moe.listed else ids  # the routing lists
    x = flat(plan_bytes(moe.plan_rows)) if moe.grouped else x  # grouped: MXFP8 plan
    scales = (flat(experts * n) for n in (moe.w13_scale_bytes, moe.w2_scale_bytes))
    out = flat(batch * 16 * (1 if publish else 3584), c.BFloat16)
    maps = flat(768 if moe.paired else 256, align=64)
    tensors = (x, ids, maps, *scales, flat(moe.workspace_bytes), out)
    scalars = (c.Int32(batch), *map(c.Float32, (4, 25, 0.25, 0.04)), *box)
    if ahead:  # the next weights' (address, bytes) table
        scalars += (fake((PREFETCH_ROWS, 2), (2, 1), c.Int64),)
    with torch.accelerator.device_index(device):
        return cute.compile(moe, *tensors, *scalars, make_fake_stream())

compile_group(tp, experts, batch, grid, device, early=False) cached

The pregrouped experts' unit-table kernel: (ids, workspace, batch); early: it releases the experts at entry.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py

@functools.cache
def compile_group(tp, experts, batch, grid, device, early=False):
    """The pregrouped experts' unit-table kernel: (ids, workspace, batch); ``early``:
    it releases the experts at entry."""
    capacity = 1 << (batch - 1).bit_length()
    kernel = ExpertsGroupKernel(tp, capacity, experts, grid, early=early)
    flat = lambda n, dtype=c.Uint8: fake((n,), (1,), dtype, 16)
    args = (flat(batch * 16, c.Int32), flat(kernel.workspace_bytes), c.Int32(batch))
    with torch.accelerator.device_index(device):
        return cute.compile(kernel, *args, make_fake_stream())

compile_plan(batch, device) cached

Plan writer for callers without the routing kernel's copy role in front.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py

@functools.cache
def compile_plan(batch, device):
    """Plan writer for callers without the routing kernel's copy role in front."""
    plan = fake((plan_bytes(plan_rows(batch)),), (1,), c.Uint8)
    args = (fake((batch, 3584)), plan, make_fake_stream())
    with torch.accelerator.device_index(device):
        return cute.compile(ExpertsPlanKernel(batch), *args)

dynamic(batch)

Static B<=8: FC2 tile-pair jobs claimed from a grid-wide counter (ExpertsKernel.dynamic): CTA stream rates vary by a few us per call, so a static FC2 slice leaves the slowest CTA as the kernel's tail.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py

def dynamic(batch):
    """Static B<=8: FC2 tile-pair jobs claimed from a grid-wide counter
    (``ExpertsKernel.dynamic``): CTA stream rates vary by a few us per call, so a
    static FC2 slice leaves the slowest CTA as the kernel's tail."""
    return static(batch) and bool(batch_table(DYNAMIC).get(batch, 0))

early_group(tp, batch)

Pregrouped batches whose unit-table kernel releases the experts at its entry.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py

def early_group(tp, batch):
    """Pregrouped batches whose unit-table kernel releases the experts at its entry."""
    table = batch_table(EARLY[tp])
    return pregrouped(tp, batch) and bool(table.get(batch, 0))

pregrouped(tp, batch)

B>8 experts read the unit table a one-CTA kernel builds after routing, beside the fused shared expert (hidden at B9-32; at B64 its ~2.3 us after routing outlast the shared expert).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py

def pregrouped(tp, batch):
    """B>8 experts read the unit table a one-CTA kernel builds after routing, beside
    the fused shared expert (hidden at B9-32; at B64 its ~2.3 us after routing outlast
    the shared expert)."""
    table = batch_table(PREGROUP[tp])
    return batch > 64 or batch > 8 and bool(table.get(batch, 0))

scratch(tp, device, wide=False)

Model-owned expert buffers: the workspace and the grouped plan (MXFP8 x, written by the routing kernel or compile_plan). wide: the wide engines' own workspace and plan (B > 64), never shared with the B <= 64 engines.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py

def scratch(tp, device, wide=False):
    """Model-owned expert buffers: the workspace and the grouped plan (MXFP8 x, written
    by the routing kernel or ``compile_plan``). ``wide``: the wide engines' own
    workspace and plan (B > 64), never shared with the B <= 64 engines."""
    geometry = ExpertsGeometry(tp, WIDE_ROWS if wide else 8)
    size, u8 = (geometry.workspace_bytes, torch.uint8)
    workspace = torch.full((size,), 255, device=device, dtype=u8)
    plan = torch.empty(plan_bytes(geometry.plan_rows), device=device, dtype=u8)
    buffers = {"workspace": workspace, "plan": plan}
    if wide:  # the listed engines' routing lists (counts start at 0)
        buffers["lists"] = torch.zeros(LIST_WORDS, device=device, dtype=torch.int32)
    return buffers

static(batch)

Streamed B<=8 with TP8 FC2 pairs: the MMA warp walks the producer's static schedule (ExpertsKernel.consume_static). The grouped engines keep the slot metadata: walking their FC1 stages gained 1-1.5% alone but lost 0.25% of the TP8 R4 step.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py

def static(batch):
    """Streamed B<=8 with TP8 FC2 pairs: the MMA warp walks the producer's static
    schedule (``ExpertsKernel.consume_static``). The grouped engines keep the slot
    metadata: walking their FC1 stages gained 1-1.5% alone but lost 0.25% of the
    TP8 R4 step."""
    return bool(batch_table(STATIC).get(batch, 0))

stealing(batch)

B>8 experts balance the FC2(B) tail by claims.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py

def stealing(batch):
    """B>8 experts balance the FC2(B) tail by claims."""
    return batch > 8 and bool(batch_table(STEAL).get(batch, 0))

streamed(batch)

B<=8 on the incremental producer: ~1-1.5% at B4-8, neutral at B1-2, which keep the previous producer.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py

def streamed(batch):
    """B<=8 on the incremental producer: ~1-1.5% at B4-8, neutral at B1-2, which keep
    the previous producer."""
    return bool(batch_table("4-8:1").get(batch, 0))