Skip to content

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

Prepared MoE provider: one model owner serializes its scratch, no tensor pointers are rebound at execution, and the side stream joins before the tail.

Classes:

MoEScratch

Methods:

  • join –

    The main stream waits for the urgent side stream (once per step).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/prepared.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
class MoEScratch:
    def __init__(self, tp, device, batches):
        self.tp, self.device, self.batches = tp, torch.device(device), tuple(batches)
        cap, h = max(batches), 6144 // tp
        with torch.accelerator.device_index(self.device.index):
            self.stream = torch.cuda.Stream(device=self.device)
            # ``SharedFront.gated`` batches: their tail Lamport-polls the shared rows
            # instead of joining the side stream, so the shared expert runs at the
            # greatest priority (its CTAs take freed SMs first) and never trails.
            top = torch.cuda.Stream.priority_range()[1]
            self.urgent = torch.cuda.Stream(device=self.device, priority=top)
            self.pending = False  # ``urgent`` work the main stream has not joined
            self.begin, self.shared_done = torch.cuda.Event(), torch.cuda.Event()
            self.gate_done = torch.cuda.Event()  # the split shared expert's front half
            f32, i32, bf16 = torch.float32, torch.int32, torch.bfloat16
            empty = lambda s, t=bf16: torch.empty(s, dtype=t, device=self.device)
            self.latent = empty((cap, 3584))
            self.front: tuple = (empty((cap, 896), f32), empty((cap, 2 * h), f32))
            self.front += (empty((cap, 16), i32), empty((cap, 16)), empty((cap, h)))
            self.front += (empty((cap, 7168)), self.latent)
            self.expanded = empty((cap * 16, 3584))
            self.experts = expert_scratch(tp, self.device)
            # The experts' plan (MXFP8 x), written by the routing kernel at B>8.
            self.plan = self.experts["plan"]
            self.shared = shared_scratch(self.device)  # fused shared-expert counters
        self.views = {}
        for b in batches:
            self.views[b] = (self.latent[:b], self.expanded[: b * 16], self.experts)

    def join(self):
        """The main stream waits for the urgent side stream (once per step)."""
        if self.pending:
            torch.cuda.current_stream(self.device).wait_stream(self.urgent)
            self.pending = False

join()

The main stream waits for the urgent side stream (once per step).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/prepared.py 
43
44
45
46
47
def join(self):
    """The main stream waits for the urgent side stream (once per step)."""
    if self.pending:
        torch.cuda.current_stream(self.device).wait_stream(self.urgent)
        self.pending = False

PreparedMoE

Methods:

  • finish –

    The routed experts, the next layer's prefetch fork, then the tail once the

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/prepared.py 
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
class PreparedMoE:
    def __init__(self, front, experts, tail, state, *, prefetch=None, layer_idx=None):
        self.front, self.experts, self.tail, self.state = front, experts, tail, state
        self.prefetch, self.layer_idx = prefetch, layer_idx
        self.small_batch = None  # optional small-batch chain
        # Split batches whose routed experts launch beside the shared expert's
        # gate/up half instead of after it.
        self.join = batch_table(JOIN[state.tp])
        self.sms = torch.cuda.get_device_properties(state.device).multi_processor_count
        for b, sms in front.split.items():  # the split's down CTAs take the SMs left
            experts.grids[b] = min(experts.grid_for(b), self.sms - sms)
        # Batches whose experts prefetch the next layer's attention weights in-kernel
        # (runtime.prefetch.EXPERTS_AHEAD): their (address, bytes) table.
        self.ahead: dict[int, torch.Tensor] = {}
        for b in state.batches:
            _, expanded, buffers = state.views[b]
            publish, planned = tail.publish.get(b, ()), front.planned(b)
            twoshot = b in tail.twoshots
            table = None if prefetch is None else prefetch.ahead(layer_idx, b)
            args = (b, buffers, expanded, publish, planned, twoshot, True)
            experts.prepare(*args, ahead=table is not None)
            if table is not None and experts.ahead[b]:
                self.ahead[b] = table

    def run(self, x, out, *, consume=None):
        state, b = self.state, x.shape[0]
        if self.small_batch is not None and b in self.small_batch.batches:
            return self.small_batch.run(out, consume)
        main = torch.cuda.current_stream(state.device)
        # The shared expert runs on the side stream beside the routed experts, or, when
        # the persistent expert grid spans every SM, beside the front only. Split: its
        # gate/up half beside the front, its down rows beside the experts.
        before = self.experts.grid_for(b) >= self.sms
        if self.front.merged_batch(b):
            # One kernel publishes the routing logits and leaves h; the down rows then
            # run beside the experts.
            ids, weights, latent = self.front.route(x)
            state.begin.record(main)
            with torch.cuda.stream(state.stream):
                state.stream.wait_event(state.begin)
                shared, beside = self.front.shared_expert(x, self.tail)
                beside()
                state.shared_done.record(state.stream)
            self.experts.group(ids)
            return self.finish(b, out, latent, ids, weights, shared, consume)
        state.begin.record(main)
        side = state.urgent if self.front.gated(b) else state.stream
        with torch.cuda.stream(side):
            side.wait_event(state.begin)
            shared, beside = self.front.shared_expert(x, self.tail)
            if beside is not None:
                state.gate_done.record(side)
                beside()
            state.shared_done.record(side)
        ids, weights, latent = self.front.route(x)
        self.experts.group(ids)  # pregrouped unit table, beside the shared expert
        if before or (beside is not None and self.join.get(b, 1)):
            main.wait_event(state.shared_done if before else state.gate_done)
        return self.finish(b, out, latent, ids, weights, shared, consume, before)

    def finish(self, b, out, latent, ids, weights, shared, consume, before=False):
        """The routed experts, the next layer's prefetch fork, then the tail once the
        shared expert's side stream is done. Early-released batches skip that join:
        their tail Lamport-polls the shared rows the side stream's last kernel pushes,
        which also orders the side stream's reads of x before the next layer rewrites
        it, so the tail keeps its single (programmatic) dependency on the experts and
        launches as their CTAs exit; the step joins the urgent stream at its end
        (``MoEScratch.join``)."""
        main = torch.cuda.current_stream(self.state.device)
        ahead = self.ahead.get(b)
        expanded = self.experts.run(latent, ids, weights, ahead)
        if ahead is None and self.prefetch is not None:
            self.prefetch.fork("next_attention", self.layer_idx, b)
        if self.front.gated(b):
            self.state.pending = True
        elif not before:
            main.wait_event(self.state.shared_done)
        return self.tail.run(expanded, weights, shared, out, consume=consume)

finish(b, out, latent, ids, weights, shared, consume, before=False)

The routed experts, the next layer's prefetch fork, then the tail once the shared expert's side stream is done. Early-released batches skip that join: their tail Lamport-polls the shared rows the side stream's last kernel pushes, which also orders the side stream's reads of x before the next layer rewrites it, so the tail keeps its single (programmatic) dependency on the experts and launches as their CTAs exit; the step joins the urgent stream at its end (MoEScratch.join).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/prepared.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def finish(self, b, out, latent, ids, weights, shared, consume, before=False):
    """The routed experts, the next layer's prefetch fork, then the tail once the
    shared expert's side stream is done. Early-released batches skip that join:
    their tail Lamport-polls the shared rows the side stream's last kernel pushes,
    which also orders the side stream's reads of x before the next layer rewrites
    it, so the tail keeps its single (programmatic) dependency on the experts and
    launches as their CTAs exit; the step joins the urgent stream at its end
    (``MoEScratch.join``)."""
    main = torch.cuda.current_stream(self.state.device)
    ahead = self.ahead.get(b)
    expanded = self.experts.run(latent, ids, weights, ahead)
    if ahead is None and self.prefetch is not None:
        self.prefetch.fork("next_attention", self.layer_idx, b)
    if self.front.gated(b):
        self.state.pending = True
    elif not before:
        main.wait_event(self.state.shared_done)
    return self.tail.run(expanded, weights, shared, out, consume=consume)