Skip to content

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

Prepared first dense MLP, sharing projection and SiTU decode arithmetic.

Classes:

  • PreparedDenseMLP –

    Dense MLP; TP16 pads down K 2112 -> 2176 and the activation zero-fills it.

PreparedDenseMLP

Dense MLP; TP16 pads down K 2112 -> 2176 and the activation zero-fills it.

Methods:

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/dense.py 
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
class PreparedDenseMLP:
    """Dense MLP; TP16 pads down K 2112 -> 2176 and the activation zero-fills it."""

    def __init__(self, mlp, batches, reduce=None):
        gate_up, down = mlp.gate_up_proj.weight, mlp.down_proj.weight
        self.device, self.reduce = gate_up.device, reduce
        self.hidden = hidden = gate_up.shape[0] // 2
        self.padded = padded = (hidden + 127) // 128 * 128
        linear = -1 if mlp.act_fn.linear_beta is None else mlp.act_fn.linear_beta
        self.scalars = tuple(map(cutlass.Float32, (mlp.act_fn.beta, linear)))
        if padded != hidden:
            down = torch.nn.functional.pad(down.detach(), (0, padded - hidden))
        self.gates = gate_up.new_empty((64, 2 * hidden))
        self.hidden_rows = gate_up.new_empty((64, padded))
        tile = {b: next(n for n in (8, 16, 32, 64) if b <= n) for b in batches}
        up_tiles = {8: (64, 8, 4, 8), 16: (64, 16, 2, 8), 32: (64, 32, 2, 8)}
        up_tiles[64] = (64, 64, 2, 5)
        down_tiles = {8: (128, 8, 1, 6), 16: (128, 16, 1, 6), 32: (128, 32, 1, 5)}
        down_tiles[64] = (64, 64, 1, 5)
        policy = lambda tiles: {b: tiles[n] for b, n in tile.items()}
        self.gate = PreparedProjection(gate_up, policy(up_tiles))
        self.down = PreparedProjection(down, policy(down_tiles))
        self.gate.prepare(), self.down.prepare()
        self.activation = _compile_activation(hidden, padded, architecture(self.device))

    @torch.no_grad()
    def process_weights_after_loading(self, mlp):
        """After a weight reload: the padded TP16 down copy is ours, refill it in
        place (padding columns stay zero). Unpadded, it is the backbone weight."""
        if self.padded != self.hidden:
            self.down.weight[:, : self.hidden].copy_(mlp.down_proj.weight)

    def run(self, x, out):
        gate, hidden = self.gates[: x.shape[0]], self.hidden_rows[: x.shape[0]]
        self.gate.run(x, gate)
        views = (dynamic_view(gate), dynamic_view(hidden))
        self.activation(*views, *self.scalars, stream(self.device))
        self.down.run(hidden, out)
        if self.reduce is not None:
            reduced = self.reduce(out)
            if reduced.data_ptr() != out.data_ptr():
                out.copy_(reduced)
        return out

process_weights_after_loading(mlp)

After a weight reload: the padded TP16 down copy is ours, refill it in place (padding columns stay zero). Unpadded, it is the backbone weight.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/dense.py 
76
77
78
79
80
81
@torch.no_grad()
def process_weights_after_loading(self, mlp):
    """After a weight reload: the padded TP16 down copy is ours, refill it in
    place (padding columns stay zero). Unpadded, it is the backbone weight."""
    if self.padded != self.hidden:
        self.down.weight[:, : self.hidden].copy_(mlp.down_proj.weight)