vllm.models.kimi_k3.nvidia.specialized_decode.cutedsl.dsa.overlap

Query overlap of a K3.1 DSA layer. The indexer path (wq_b -> DsaIndex -> scoring and top-k, under DCP the top-k merge and the local slots) needs only the normalized q_c rows of the query, so the rest of the query stage leaves the critical path:

QueryNorm q_a RMSNorm -> the q_c rows, on the main stream ahead of wq_b (bit-identical to MlaQuery's: its scalar form below tiled_from rows, its tiled kernel's packed FP32 pairs from there) MlaQuery Q_B, W_UK absorption, the query rows, the MLA KV insert and, under DCP, the q mailbox publish, on a side stream (store_qc=False, poll=False) QueryWait DCP: the q mailbox poll (until every rank's heads landed) after the MlaQuery on the side stream, a few CTAs instead of MlaQuery's grid

The main stream joins the side stream just before the sparse core.

Classes:

• QueryNorm –

  The q_a RMSNorm of batch rows of the packed projection -> q_c [B, 1536]

• QueryWait –

  DCP: poll the gathered q mailbox (batch tokens of gathered = DCP x H

Functions:

• overlap_for –

  The query overlap mode of batch rows: policy is a mode or a

• side_stream –

  The DSA layers' query side stream of a device.

QueryNorm

The q_a RMSNorm of batch rows of the packed projection -> q_c [B, 1536] BF16, exactly as MlaQuery (heads) computes it for the same batch.

Methods:

• kernel –

  One row per CTA: 8 columns per thread, MlaQuery.kernel's FMA order.

• tiled –

  One row per warp: 48 columns per lane, MlaQuery.tiled's packed pairs.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/dsa/overlap.py

class QueryNorm:
    """The q_a RMSNorm of ``batch`` rows of the packed projection -> q_c [B, 1536]
    BF16, exactly as MlaQuery (``heads``) computes it for the same batch."""

    def __init__(self, heads: int, batch: int):
        self.batch = batch
        self.wide = batch >= tiled_from(heads)
        self.grid = -(-batch // 8) if self.wide else batch

    @cute.jit
    def __call__(self, *args):
        launch(self.tiled if self.wide else self.kernel, args, self.grid, 256)

    @cute.kernel
    def kernel(self, args):
        """One row per CTA: 8 columns per thread, MlaQuery.kernel's FMA order."""
        qkv, w_qa, qc, eps = args
        tid, _, _ = cute.arch.thread_idx()
        row, _, _ = cute.arch.block_idx()
        lane, warp = tid % 32, tid // 32
        smem = cutlass.utils.SmemAllocator()
        ss = alloc(smem, Float32, 8)
        weights = load_bf16x8(w_qa.iterator + (tid % 192) * 8)
        cute.arch.griddepcontrol_wait()
        cute.arch.griddepcontrol_launch_dependents()
        xv = cute.make_rmem_tensor(8, Float32)
        total = Float32(0)
        if tid < 192:
            values = load_bf16x8(at(qkv, row, tid * 8), mutable=True)
            for i in cutlass.range_constexpr(8):
                xv[i] = values[i]
                total = fma(values[i], values[i], total)
        total = warp_sum(total)
        if lane == 0:
            ss[warp] = total
        cute.arch.sync_threads()
        if tid < 192:
            total = Float32(0)
            for w in cutlass.range_constexpr(8):
                total += ss[w]
            rr = query_norm_rsqrt(total, eps)
            words = [
                bf16x2(xv[i] * rr * weights[i], xv[i + 1] * rr * weights[i + 1])
                for i in range(0, 8, 2)
            ]
            store_words(at(qc, row, tid * 8), tuple(words))

    @cute.kernel
    def tiled(self, args):
        """One row per warp: 48 columns per lane, MlaQuery.tiled's packed pairs."""
        qkv, w_qa, qc, eps = args
        tid, _, _ = cute.arch.thread_idx()
        tile, _, _ = cute.arch.block_idx()
        lane, warp = tid % 32, tid // 32
        row = tile * 8 + warp
        gamma = []
        for i in cutlass.range_constexpr(6):
            gamma += list(load_bf16x8(w_qa.iterator + 256 * i + 8 * lane))
        cute.arch.griddepcontrol_wait()
        cute.arch.griddepcontrol_launch_dependents()
        if row < self.batch:
            values: list = []
            for i in cutlass.range_constexpr(6):
                source = at(qkv, row, 256 * i + 8 * lane)
                for word in load_words(source, count=4, mutable=True):
                    values += list(unpack_bf16(word))
            zero = Float32(0), Float32(0)
            chains = [zero, zero]  # (0, 1), (2, 3): element e feeds chain e % 4
            for e in cutlass.range_constexpr(0, 48, 2):
                x = values[e], values[e + 1]
                chains[e % 4 // 2] = fma2(x, x, chains[e % 4 // 2])
            (c0, c1), (c2, c3) = chains
            total = warp_sum((c0 + c1) + (c2 + c3))
            rr = query_norm_rsqrt(total, eps)
            for i in cutlass.range_constexpr(6):
                words = []
                for e in cutlass.range_constexpr(8 * i, 8 * i + 8, 2):
                    x = values[e], values[e + 1]
                    y = mul2(mul2(x, (rr, rr)), (gamma[e], gamma[e + 1]))
                    words.append(bf16x2(*y))
                store_words(at(qc, row, 256 * i + 8 * lane), tuple(words))

kernel(args)

One row per CTA: 8 columns per thread, MlaQuery.kernel's FMA order.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/dsa/overlap.py

@cute.kernel
def kernel(self, args):
    """One row per CTA: 8 columns per thread, MlaQuery.kernel's FMA order."""
    qkv, w_qa, qc, eps = args
    tid, _, _ = cute.arch.thread_idx()
    row, _, _ = cute.arch.block_idx()
    lane, warp = tid % 32, tid // 32
    smem = cutlass.utils.SmemAllocator()
    ss = alloc(smem, Float32, 8)
    weights = load_bf16x8(w_qa.iterator + (tid % 192) * 8)
    cute.arch.griddepcontrol_wait()
    cute.arch.griddepcontrol_launch_dependents()
    xv = cute.make_rmem_tensor(8, Float32)
    total = Float32(0)
    if tid < 192:
        values = load_bf16x8(at(qkv, row, tid * 8), mutable=True)
        for i in cutlass.range_constexpr(8):
            xv[i] = values[i]
            total = fma(values[i], values[i], total)
    total = warp_sum(total)
    if lane == 0:
        ss[warp] = total
    cute.arch.sync_threads()
    if tid < 192:
        total = Float32(0)
        for w in cutlass.range_constexpr(8):
            total += ss[w]
        rr = query_norm_rsqrt(total, eps)
        words = [
            bf16x2(xv[i] * rr * weights[i], xv[i + 1] * rr * weights[i + 1])
            for i in range(0, 8, 2)
        ]
        store_words(at(qc, row, tid * 8), tuple(words))

tiled(args)

One row per warp: 48 columns per lane, MlaQuery.tiled's packed pairs.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/dsa/overlap.py

@cute.kernel
def tiled(self, args):
    """One row per warp: 48 columns per lane, MlaQuery.tiled's packed pairs."""
    qkv, w_qa, qc, eps = args
    tid, _, _ = cute.arch.thread_idx()
    tile, _, _ = cute.arch.block_idx()
    lane, warp = tid % 32, tid // 32
    row = tile * 8 + warp
    gamma = []
    for i in cutlass.range_constexpr(6):
        gamma += list(load_bf16x8(w_qa.iterator + 256 * i + 8 * lane))
    cute.arch.griddepcontrol_wait()
    cute.arch.griddepcontrol_launch_dependents()
    if row < self.batch:
        values: list = []
        for i in cutlass.range_constexpr(6):
            source = at(qkv, row, 256 * i + 8 * lane)
            for word in load_words(source, count=4, mutable=True):
                values += list(unpack_bf16(word))
        zero = Float32(0), Float32(0)
        chains = [zero, zero]  # (0, 1), (2, 3): element e feeds chain e % 4
        for e in cutlass.range_constexpr(0, 48, 2):
            x = values[e], values[e + 1]
            chains[e % 4 // 2] = fma2(x, x, chains[e % 4 // 2])
        (c0, c1), (c2, c3) = chains
        total = warp_sum((c0 + c1) + (c2 + c3))
        rr = query_norm_rsqrt(total, eps)
        for i in cutlass.range_constexpr(6):
            words = []
            for e in cutlass.range_constexpr(8 * i, 8 * i + 8, 2):
                x = values[e], values[e + 1]
                y = mul2(mul2(x, (rr, rr)), (gamma[e], gamma[e + 1]))
                words.append(bf16x2(*y))
            store_words(at(qc, row, 256 * i + 8 * lane), tuple(words))

QueryWait

DCP: poll the gathered q mailbox (batch tokens of gathered = DCP x H published heads in rows of q_heads, q_bytes per element) until no 16 B chunk is the sentinel. It first waits for the MlaQuery before it, so the KV insert is complete too when the main stream joins.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/dsa/overlap.py

class QueryWait:
    """DCP: poll the gathered q mailbox (``batch`` tokens of ``gathered`` = DCP x H
    published heads in rows of ``q_heads``, ``q_bytes`` per element) until no 16 B
    chunk is the sentinel. It first waits for the MlaQuery before it, so the KV
    insert is complete too when the main stream joins."""

    def __init__(self, batch, gathered, q_heads, q_bytes):
        self.row = 36 * q_bytes  # 16 B chunks per head row
        self.token, self.q_heads = gathered * self.row, q_heads
        self.chunks = batch * self.token
        self.ctas = min(32, -(-self.chunks // (256 * 8)))

    @cute.jit
    def __call__(self, *args):
        launch(self.kernel, args, self.ctas, 256)

    @cute.kernel
    def kernel(self, args):
        (mailbox,) = args
        tid, _, _ = cute.arch.thread_idx()
        cta, _, _ = cute.arch.block_idx()
        cute.arch.griddepcontrol_wait()
        cute.arch.griddepcontrol_launch_dependents()
        threads: cutlass.Constexpr = self.ctas * 256
        token, row, chunks = self.token, self.row, self.chunks
        first = cta * 256 + tid
        steps: cutlass.Constexpr = -(-chunks // threads)
        for i0 in cutlass.range_constexpr(0, steps, 8):
            n: cutlass.Constexpr = min(8, steps - i0)
            picks = [min(first + (i0 + i) * threads, chunks - 1) for i in range(n)]
            sources = [
                mailbox.iterator + (c // token * self.q_heads * row + c % token) * 16
                for c in picks
            ]
            loads = [load_volatile(source, count=4) for source in sources]
            for i in cutlass.range_constexpr(n):  # the loads above are in flight
                words = loads[i]
                while not published(words, sentinel=Q_SENT):
                    words = load_volatile(sources[i], count=4)

overlap_for(batch, policy=None, dcp=False, wide=False, split=False)

The query overlap mode of batch rows: policy is a mode or a "rows:mode,..." table (default QUERY_OVERLAP of the topology, dcp; wide: the split query of TP1, WIDE_MODES only; split: the split query of fewer heads, under DCP its own table).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/dsa/overlap.py

def overlap_for(batch, policy=None, dcp=False, wide=False, split=False):
    """The query overlap mode of ``batch`` rows: ``policy`` is a mode or a
    "rows:mode,..." table (default QUERY_OVERLAP of the topology, ``dcp``;
    ``wide``: the split query of TP1, WIDE_MODES only; ``split``: the split query
    of fewer heads, under DCP its own table)."""
    topology = "tp1" if wide else "dcp" if dcp else "tp"
    if split and dcp and not wide:
        topology = "dcp_split"
    policy = policy or QUERY_OVERLAP[topology]
    mode = policy if policy in MODES else batch_table(policy).get(batch, "fork")
    modes = WIDE_MODES if wide else MODES
    if mode not in modes:
        raise ValueError(f"Kimi K3.1 DSA query overlap {mode!r}: expected {modes}")
    return mode

side_stream(device_index) cached

The DSA layers' query side stream of a device.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/dsa/overlap.py

@cache
def side_stream(device_index):
    """The DSA layers' query side stream of a device."""
    return torch.cuda.Stream(device=device_index)