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:

• PreparedKda –

Functions:

• core –

  The KDA core for a step shape: KdaWS, KdaStream, KdaWide or KdaSplit by the

• decode_conv_weight –

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

• deferred_supported –

  Whether the KDA metadata of the specialized decode defers a layer's RecoverSSM

• replays –

  Whether a deferred RecoverSSM verify step runs KdaFused (else: flush + core).

• stream_ctas –

  Persistent CTAs (two per SM) with an equal item count each.

• stream_layout –

  (conv dim-first, window) when KdaStream can cp.async this step's operands

PreparedKda

Methods:

• prepare_caches –

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

• process_weights_after_loading –

  Refill the derived conv weight and FP32 norm after a weight reload.

• run –

  accepted [>= B] (verify steps of q tokens): slots is [>= B, q]

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py

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, {}
        # RecoverSSM: verify steps leave the state checkpoint and write the records.
        records = getattr(attention, "recoverssm_records", None)
        self.records: tuple = (records.correction, records.key_gate) if records else ()
        self.backend = getattr(attention, "recoverssm_backend", None)
        self.lower_bound = attention.gate_lower_bound
        self.fused = False  # set by prepare_caches (deferred_supported)
        self.flushes = set()  # deferred steps that fold the replays first (keys)

    @torch.no_grad()
    def process_weights_after_loading(self, attention):
        """Refill the derived conv weight and FP32 norm after a weight reload."""
        conv_w, *_, norm = self.weights
        conv_w.copy_(decode_conv_weight(attention.conv1d.weight))
        norm.copy_(attention.o_norm.weight)

    def _operands(self, caches, projected, slots, out, accepted=None, deferred=()):
        """Kernel tensors in launch order: 13 before the scalars, then F_B (and the
        verify step's accepted counts and RecoverSSM records; fused: the record rows
        and pending 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, *self.records, *deferred)
        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."""
        self.flushes = set()
        self.fused = bool(self.records) and deferred_supported(
            self.backend, caches[1], self.lower_bound
        )
        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)
            record = tokens > 1 and bool(self.records)
            # RecoverSSM verify steps have one state slot (the checkpoint).
            shape = (requests, 1 if record else tokens) if tokens > 1 else (batch,)
            slots = torch.empty(shape, dtype=torch.int32, **meta)
            accepted = accepted if tokens > 1 else None
            fused = record and self.fused
            deferred: tuple = ()
            if fused and not replays(self.heads, requests, True):
                self.flushes.add(key)
            if fused:
                rows = torch.empty(requests, dtype=torch.int32, **meta)
                count = self.records[0].shape[0]  # one pending count per record slot
                deferred = (rows, torch.empty(count, dtype=torch.int32, **meta))
            operands = self._operands(caches, projected, slots, out, accepted, deferred)
            self.kernels[key] = executable(
                tuple(map(layout, operands)),
                *(self.heads, mode, self.has_bias, self.device.index, tokens),
                record,
                fused,
            )
        return self

    def run(self, projected, out, *, caches, slots, accepted=None, q=1, deferred=None):
        """``accepted`` [>= B] (verify steps of ``q`` tokens): ``slots`` is [>= B, q]
        contiguous, or [>= B, 1] (the checkpoints) under RecoverSSM. ``deferred``
        (the KDA metadata's ``recoverssm_deferred``): the record rows [>= B] and
        pending counts of a deferred RecoverSSM step."""
        batch = projected.shape[0]
        key, rows = batch, batch
        extra: tuple = ()
        if accepted is not None:
            rows = batch // q
            key, accepted = (rows, q), accepted[:rows]
            if slots.stride() != (slots.shape[1], 1):  # e.g. a block-table slice
                slots = slots.contiguous()
            if self.records and self.fused:
                assert deferred is not None, "a deferred KDA core needs its metadata"
                extra = (deferred.record_rows[:rows], deferred.pending)
                if key in self.flushes:
                    deferred.flush()  # this core verifies from committed checkpoints
        tensors = self._operands(caches, projected, slots[:rows], out, accepted, extra)
        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, deferred=())

Kernel tensors in launch order: 13 before the scalars, then F_B (and the verify step's accepted counts and RecoverSSM records; fused: the record rows and pending counts).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py

def _operands(self, caches, projected, slots, out, accepted=None, deferred=()):
    """Kernel tensors in launch order: 13 before the scalars, then F_B (and the
    verify step's accepted counts and RecoverSSM records; fused: the record rows
    and pending 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, *self.records, *deferred)
    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

def prepare_caches(self, caches):
    """Policy keys: a batch, or (requests, tokens) for verify steps."""
    self.flushes = set()
    self.fused = bool(self.records) and deferred_supported(
        self.backend, caches[1], self.lower_bound
    )
    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)
        record = tokens > 1 and bool(self.records)
        # RecoverSSM verify steps have one state slot (the checkpoint).
        shape = (requests, 1 if record else tokens) if tokens > 1 else (batch,)
        slots = torch.empty(shape, dtype=torch.int32, **meta)
        accepted = accepted if tokens > 1 else None
        fused = record and self.fused
        deferred: tuple = ()
        if fused and not replays(self.heads, requests, True):
            self.flushes.add(key)
        if fused:
            rows = torch.empty(requests, dtype=torch.int32, **meta)
            count = self.records[0].shape[0]  # one pending count per record slot
            deferred = (rows, torch.empty(count, dtype=torch.int32, **meta))
        operands = self._operands(caches, projected, slots, out, accepted, deferred)
        self.kernels[key] = executable(
            tuple(map(layout, operands)),
            *(self.heads, mode, self.has_bias, self.device.index, tokens),
            record,
            fused,
        )
    return self

process_weights_after_loading(attention)

Refill the derived conv weight and FP32 norm after a weight reload.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py

@torch.no_grad()
def process_weights_after_loading(self, attention):
    """Refill the derived conv weight and FP32 norm after a weight reload."""
    conv_w, *_, norm = self.weights
    conv_w.copy_(decode_conv_weight(attention.conv1d.weight))
    norm.copy_(attention.o_norm.weight)

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

accepted [>= B] (verify steps of q tokens): slots is [>= B, q] contiguous, or [>= B, 1] (the checkpoints) under RecoverSSM. deferred (the KDA metadata's recoverssm_deferred): the record rows [>= B] and pending counts of a deferred RecoverSSM step.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py

def run(self, projected, out, *, caches, slots, accepted=None, q=1, deferred=None):
    """``accepted`` [>= B] (verify steps of ``q`` tokens): ``slots`` is [>= B, q]
    contiguous, or [>= B, 1] (the checkpoints) under RecoverSSM. ``deferred``
    (the KDA metadata's ``recoverssm_deferred``): the record rows [>= B] and
    pending counts of a deferred RecoverSSM step."""
    batch = projected.shape[0]
    key, rows = batch, batch
    extra: tuple = ()
    if accepted is not None:
        rows = batch // q
        key, accepted = (rows, q), accepted[:rows]
        if slots.stride() != (slots.shape[1], 1):  # e.g. a block-table slice
            slots = slots.contiguous()
        if self.records and self.fused:
            assert deferred is not None, "a deferred KDA core needs its metadata"
            extra = (deferred.record_rows[:rows], deferred.pending)
            if key in self.flushes:
                deferred.flush()  # this core verifies from committed checkpoints
    tensors = self._operands(caches, projected, slots[:rows], out, accepted, extra)
    views = tuple(to_cute(t, 16) for t in tensors)
    self.kernels[key](*views[:13], *self.scalars, *views[13:], stream(self.device))
    return out

core(specs, heads, gate_mode, has_bias, device, tokens, record, fused=False)

The KDA core for a step shape: KdaWS, KdaStream, KdaWide or KdaSplit by the per-kind (request, head) item thresholds above (DEP's 96-head RecoverSSM verify takes KdaStream below KdaWS: TP1 R1 split / wide / stream 14.3 / 13.0 / 12.2, R2 19.2 / 19.6 / 17.8, R6 47.3 / 37.6 / 33.7 us).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py

def core(specs, heads, gate_mode, has_bias, device, tokens, record, fused=False):
    """The KDA core for a step shape: KdaWS, KdaStream, KdaWide or KdaSplit by the
    per-kind (request, head) item thresholds above (DEP's 96-head RecoverSSM
    verify takes KdaStream below KdaWS: TP1 R1 split / wide / stream 14.3 / 13.0 /
    12.2, R2 19.2 / 19.6 / 17.8, R6 47.3 / 37.6 / 33.7 us)."""
    requests = specs[0][0][0] // tokens
    items, kind = heads * requests, kind_of(tokens, record)
    if replays(heads, requests, fused):
        streamed = stream_layout(specs, tokens)
        # q = 2 (k = 1): KdaFusedWS faults (open issue), KdaFused serves it.
        if streamed is not None and items >= FUSED_WS_ITEMS and tokens >= 3:
            conv_size = torch.empty((), dtype=specs[3][2]).element_size()
            rq = specs[15][0][2]  # tokens a record holds (corrections)
            size = fused_ws.shared_bytes(tokens, rq, *streamed, conv_size)
            if size <= fused_ws.SMEM_LIMIT:
                sms = torch.cuda.get_device_properties(device).multi_processor_count
                per = max(2, -(-items // sms))
                ctas = -(-items // per)
                mode = gate_mode % 2
                return KdaFusedWS(heads, tokens, mode, has_bias, *streamed, ctas)
        return KdaFused(heads, tokens, gate_mode % 2, has_bias, fused_split(items))
    sms = torch.cuda.get_device_properties(device).multi_processor_count
    streamed = stream_layout(specs, tokens)
    mode = gate_mode % 2
    small = kind == "record" and WS_SMALL <= items < WS_SMALL_END
    if streamed is not None and (items >= WS_ITEMS[kind] or small):
        conv_size = torch.empty((), dtype=specs[3][2]).element_size()
        if (
            specialized.shared_bytes(tokens, *streamed, conv_size)
            <= specialized.SMEM_LIMIT
        ):
            per = max(2, -(-items // sms))  # items per CTA (one CTA per SM)
            ctas = -(-items // per)
            return KdaWS(heads, tokens, record, mode, has_bias, *streamed, ctas)
    stream = items >= STREAM_ITEMS[kind] or (kind == "record" and heads >= STREAM_HEADS)
    if streamed is not None and stream:
        ctas = stream_ctas(items, device)
        return KdaStream(heads, tokens, record, mode, has_bias, *streamed, ctas)
    if wide(heads, requests, tokens, record):
        return KdaWide(heads, tokens, record, mode, has_bias, 2 * sms)
    table = SCHEDULE if tokens == 1 else SPEC_SCHEDULE
    schedule = batch_table(table.get(heads, ""))
    schedule = schedule.get(requests, (0, 0, 0)) if gate_mode != 2 else (0,) * 3
    if schedule[0] >= 2 and specs[9][1][-2:] != (128, 1):  # staging needs whole rows
        schedule = (1, *schedule[1:4])  # no evict hint without staging
    return KdaSplit(heads, requests, gate_mode, has_bias, schedule, tokens, record)

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

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()

deferred_supported(backend, state, lower_bound)

Whether the KDA metadata of the specialized decode defers a layer's RecoverSSM commit: the layer's RecoverSSM kernels are Triton (the stock FP32-state path) and the specialized FP32 commit can fold its checkpoints (FP32 [slot, H, 128, 128] on SM10x, a bounded gate). recoverssm.context.KDADeferredCommitContext takes the same decision (its commit backend fp32_deferred).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py

def deferred_supported(backend, state, lower_bound):
    """Whether the KDA metadata of the specialized decode defers a layer's RecoverSSM
    commit: the layer's RecoverSSM kernels are Triton (the stock FP32-state path) and
    the specialized FP32 commit can fold its checkpoints (FP32 [slot, H, 128, 128] on
    SM10x, a bounded gate). ``recoverssm.context.KDADeferredCommitContext`` takes
    the same decision (its commit backend ``fp32_deferred``)."""
    from .commit import supported

    return backend == "triton" and supported(state, lower_bound)

replays(heads, requests, fused)

Whether a deferred RecoverSSM verify step runs KdaFused (else: flush + core).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py

def replays(heads, requests, fused):
    """Whether a deferred RecoverSSM verify step runs KdaFused (else: flush + core)."""
    if isinstance(FUSED_ITEMS, int):  # a forced threshold (tests, benchmarks)
        return fused and heads * requests >= FUSED_ITEMS
    items = FUSED_ITEMS.get(heads, FUSED_DEFAULT_ITEMS)
    return fused and heads * requests >= items

stream_ctas(items, device)

Persistent CTAs (two per SM) with an equal item count each.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py

def stream_ctas(items, device):
    """Persistent CTAs (two per SM) with an equal item count each."""
    slots = 2 * torch.cuda.get_device_properties(device).multi_processor_count
    per = -(-items // slots)
    return -(-items // per)

stream_layout(specs, tokens)

(conv dim-first, window) when KdaStream can cp.async this step's operands (16-byte aligned rows and slots), else None.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py

def stream_layout(specs, tokens) -> tuple[bool, int] | None:
    """(conv dim-first, window) when KdaStream can cp.async this step's operands
    (16-byte aligned rows and slots), else None."""
    (_, x_stride, _), conv = specs[0], specs[3]
    shape, stride, dtype = conv
    size = torch.empty((), dtype=dtype).element_size()
    window = shape[-1]
    if stride[-1] == 1 and stride[-2] == window:
        dim_first = True
    elif stride[-2] == 1:
        dim_first = False
    else:
        return None
    aligned = (stride[0] * size) % 16 == 0 and (x_stride[0] * 2) % 16 == 0
    if not dim_first:
        aligned = aligned and (stride[-1] * size) % 16 == 0
    return (dim_first, window) if aligned and window >= tokens + 2 else None