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:

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 
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 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
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)
        for b in state.batches:
            _, expanded, buffers = state.views[b]
            publish, planned = tail.publish.get(b, ()), front.planned(b)
            twoshot = b in tail.twoshots
            experts.prepare(b, buffers, expanded, publish, planned, twoshot, True)

    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)
        with torch.cuda.stream(state.stream):
            state.stream.wait_event(state.begin)
            shared, beside = self.front.shared_expert(x, self.tail)
            if beside is not None:
                state.gate_done.record(state.stream)
                beside()
            state.shared_done.record(state.stream)
        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."""
        main = torch.cuda.current_stream(self.state.device)
        expanded = self.experts.run(latent, ids, weights)
        if self.prefetch is not None:
            self.prefetch.fork("next_attention", self.layer_idx, b)
        if 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.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/prepared.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
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."""
    main = torch.cuda.current_stream(self.state.device)
    expanded = self.experts.run(latent, ids, weights)
    if self.prefetch is not None:
        self.prefetch.fork("next_attention", self.layer_idx, b)
    if not before:
        main.wait_event(self.state.shared_done)
    return self.tail.run(expanded, weights, shared, out, consume=consume)