Skip to content

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

Prepared KDA core. One packed projection supplies Q/K/V, the output gate, F_A and beta; per-batch policies select the gate schedule and the projection row pitch.

Classes:

Functions:

  • decode_conv_weight –

    The [3, width, D] (Q, K, V) decode layout of the [3 * D, 1, width] conv

PreparedKda

Methods:

  • prepare_caches –

    Policy keys: a batch, or (requests, tokens) for verify steps.

  • run –

    accepted [>= B] (verify steps): slots is [>= B, q] contiguous.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py 
 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
class PreparedKda:
    def __init__(self, attention, policies, conv_dim_first):
        bias, dt = attention.conv1d.bias, attention.dt_bias
        self.heads = attention.A_log.numel()
        self.dim, self.has_bias = self.heads * 128, bias is not None
        conv_w = decode_conv_weight(attention.conv1d.weight)
        bias, fb = dt if bias is None else bias, attention.f_b_proj.weight
        norm = attention.o_norm.weight.detach().float()  # BF16 slows the fused kernel
        self.weights = (conv_w, bias, fb, attention.A_log, dt, norm)
        self.device, self.policies = dt.device, policies
        scalars = (attention.gate_lower_bound, 128**-0.5, attention.o_norm.eps)
        self.scalars = tuple(map(cutlass.Float32, scalars))
        self.conv_dim_first, self.kernels = conv_dim_first, {}

    def _operands(self, caches, projected, slots, out, accepted=None):
        """Kernel tensors in launch order: 13 before the scalars, then F_B (and the
        verify step's accepted counts)."""
        conv, state = caches[:2]
        conv = conv if self.conv_dim_first else conv.transpose(-1, -2)
        w, bias, fb, a_log, dt, norm = self.weights
        d, h = self.dim, self.heads
        p = lambda start, width: projected[:, start : start + width]
        ops = (p(0, 3 * d), w, bias, conv, p(4 * d, 128), p(4 * d + 128, h), a_log, dt)
        tail = (fb,) if accepted is None else (fb, accepted)
        return ops + (slots, state, out, p(3 * d, d), norm) + tail

    def prepare_caches(self, caches):
        """Policy keys: a batch, or (requests, tokens) for verify steps."""
        for key, (pitch, mode) in self.policies.items():
            requests, tokens = key if isinstance(key, tuple) else (key, 1)
            batch, meta = requests * tokens, dict(device="meta")
            projected = torch.empty((batch, pitch), dtype=torch.bfloat16, **meta)
            out = torch.empty(batch, self.dim, dtype=torch.bfloat16, **meta)
            accepted = torch.empty(requests, dtype=torch.int32, **meta)
            shape = (requests, tokens) if tokens > 1 else (batch,)
            slots = torch.empty(shape, dtype=torch.int32, **meta)
            accepted = accepted if tokens > 1 else None
            operands = self._operands(caches, projected, slots, out, accepted)
            self.kernels[key] = executable(
                tuple(map(layout, operands)),
                *(self.heads, mode, self.has_bias, self.device.index, tokens),
            )
        return self

    def run(self, projected, out, *, caches, slots, accepted=None):
        """``accepted`` [>= B] (verify steps): ``slots`` is [>= B, q] contiguous."""
        batch = projected.shape[0]
        key, rows = batch, batch
        if accepted is not None:
            rows = batch // slots.shape[1]
            key, accepted = (rows, slots.shape[1]), accepted[:rows]
            if slots.stride() != (slots.shape[1], 1):  # e.g. a block-table slice
                slots = slots.contiguous()
        tensors = self._operands(caches, projected, slots[:rows], out, accepted)
        views = tuple(to_cute(t, 16) for t in tensors)
        self.kernels[key](*views[:13], *self.scalars, *views[13:], stream(self.device))
        return out

_operands(caches, projected, slots, out, accepted=None)

Kernel tensors in launch order: 13 before the scalars, then F_B (and the verify step's accepted counts).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py 
70
71
72
73
74
75
76
77
78
79
80
def _operands(self, caches, projected, slots, out, accepted=None):
    """Kernel tensors in launch order: 13 before the scalars, then F_B (and the
    verify step's accepted counts)."""
    conv, state = caches[:2]
    conv = conv if self.conv_dim_first else conv.transpose(-1, -2)
    w, bias, fb, a_log, dt, norm = self.weights
    d, h = self.dim, self.heads
    p = lambda start, width: projected[:, start : start + width]
    ops = (p(0, 3 * d), w, bias, conv, p(4 * d, 128), p(4 * d + 128, h), a_log, dt)
    tail = (fb,) if accepted is None else (fb, accepted)
    return ops + (slots, state, out, p(3 * d, d), norm) + tail

prepare_caches(caches)

Policy keys: a batch, or (requests, tokens) for verify steps.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py 
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def prepare_caches(self, caches):
    """Policy keys: a batch, or (requests, tokens) for verify steps."""
    for key, (pitch, mode) in self.policies.items():
        requests, tokens = key if isinstance(key, tuple) else (key, 1)
        batch, meta = requests * tokens, dict(device="meta")
        projected = torch.empty((batch, pitch), dtype=torch.bfloat16, **meta)
        out = torch.empty(batch, self.dim, dtype=torch.bfloat16, **meta)
        accepted = torch.empty(requests, dtype=torch.int32, **meta)
        shape = (requests, tokens) if tokens > 1 else (batch,)
        slots = torch.empty(shape, dtype=torch.int32, **meta)
        accepted = accepted if tokens > 1 else None
        operands = self._operands(caches, projected, slots, out, accepted)
        self.kernels[key] = executable(
            tuple(map(layout, operands)),
            *(self.heads, mode, self.has_bias, self.device.index, tokens),
        )
    return self

run(projected, out, *, caches, slots, accepted=None)

accepted [>= B] (verify steps): slots is [>= B, q] contiguous.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
def run(self, projected, out, *, caches, slots, accepted=None):
    """``accepted`` [>= B] (verify steps): ``slots`` is [>= B, q] contiguous."""
    batch = projected.shape[0]
    key, rows = batch, batch
    if accepted is not None:
        rows = batch // slots.shape[1]
        key, accepted = (rows, slots.shape[1]), accepted[:rows]
        if slots.stride() != (slots.shape[1], 1):  # e.g. a block-table slice
            slots = slots.contiguous()
    tensors = self._operands(caches, projected, slots[:rows], out, accepted)
    views = tuple(to_cute(t, 16) for t in tensors)
    self.kernels[key](*views[:13], *self.scalars, *views[13:], stream(self.device))
    return out

decode_conv_weight(weight)

The [3, width, D] (Q, K, V) decode layout of the [3 * D, 1, width] conv weight.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py 
49
50
51
52
53
def decode_conv_weight(weight):
    """The [3, width, D] (Q, K, V) decode layout of the [3 * D, 1, width] conv
    weight."""
    width = weight.shape[-1]
    return weight.detach().view(3, -1, width).transpose(1, 2).contiguous()