`vllm.models.kimi_k3.nvidia.specialized_decode.drafter.attention` ¶

Dense MLA decode of the DSpark draft block with the query tokens folded into heads.

A DSpark draft step runs k query tokens per request (anchor + k - 1 noise tokens) through each dense MLA draft layer. The block is non-causal: every query token of a request attends to exactly the same keys (the whole context plus the block itself, already written to the cache). The stock FLASHINFER_MLA decode has no non-causal flag, so it flattens the block into R * k independent single-token rows, each with its own copy of the request's block table. Every row then streams the request's whole latent cache again: k times the KV bytes, and each CTA's MMA sees only the H local heads of one token (8 at TP8).

Absorbed MLA decode is MQA (one latent KV head), so query heads are independent rows over a shared key set. Folding the k tokens of a request into the head axis, [R * k, H, D] -> [R, 1, k * H, D], gives one decode row per request with k * H heads over the same keys. That is the same attention, exactly: every (token, head) pair still gets its own full softmax over the identical key set, at the same precision (FP8 Q and KV, the stock bmm scales); only the kernel's tiling (which CTA holds which head, the KV split) can differ. The KV is read once per request instead of k times, and the MMA M grows k-fold (32-56 heads at TP8, 256-448 at TP1 / DCP).

The output [R, 1, k * H, L] and LSE [R, k * H] are the stock [R * k, H, L] / [R * k, H] in the same memory order (token-major, then head), so the DCP combine and the W_UV up-projection are unchanged.

Classes:

DraftDcp –

The draft block's MLA decode under DCP (DCP = TP: rank d owns head block d)

Functions:

balanced_decode –

The cute-dsl LSE decode of common (folded_decode's FlashInfer
balanced_tiers –

The split tiers (rows each, longest first) of a cute-dsl LSE draft decode of
cake_supported –

CAKE's Kimi K3 MLA kernel: SM100 family, FP8 query and cache, page 64,
cake_workspace_bytes –

The largest CAKE split-KV workspace any folded draft batch of up to
causal_dcp_decode –

CAUSAL_DCP_DECODE's decode of a causal draft block of rows rows.
causal_folded_decode –

Causal dense MLA decode of R requests x k draft tokens: token j of a
dcp_causal_block_merge –

Add this rank's share of a causal block's keys (global L + i, i < k; rank
dcp_causal_native_decode –

Causal MLA decode of R requests x k draft tokens on this DCP rank, as
draft_forward_mqa –

FlashInferMLAImpl.forward_mqa for a DSpark draft block, folded.
fold_plan –

(query tokens folded into one decode row, FlashInfer MLA backend).
foldable –

Query tokens per request when the decode batch can be folded, else 0.
folded_decode –

Non-causal dense MLA decode of R requests x k query tokens.
reserve_cake_workspace –

Allocate the CAKE workspace for the deployment's largest draft batch now,
tier_workspace_bytes –

The cute-dsl decode's split-KV workspace: FP32 partials + LSE per split

`DraftDcp` ¶

The draft block's MLA decode under DCP (DCP = TP: rank d owns head block d) on the specialized decode's DCP mailboxes, in three launches per layer instead of six (the direct-DCP epoch + query all-gather, the folded decode's split-KV + reducer, the LSE all-to-all push + combine) plus the W_UV bmm:

query every rank multicasts its FP8 query rows [T, H, 576] into head block rank of every rank's DcpBox q mailbox [T, DCPH, 576] (comm.Gather, in place) and polls until all landed decode the stock folded draft decode (folded_decode: the k tokens of a request folded into heads, every DCPH head over the local KV shard, rank-local BF16 rows + LSE), reading the mailbox; a causal block as CAUSAL_DCP_DECODE picks value mla.dcp.DcpValue: re-arms the q mailbox, pushes each row to its head owner, the owner merges its H heads over the DCP sources (exp2(lse - max) / sum in FP32, one BF16 rounding, as the stock combine) and applies W_UV (BF16 MMA, FP32 accumulation) into the attention output. The draft layers have no output gate: the gate operand is +inf (sigmoid exactly 1).

Same attention as the stock DCP decode at every stage: the decode kernel and its query are the stock ones; only the transport and the merge order differ.

Methods:

eligible –

The draft layers this path serves: the stock folded draft decode
forward –

attn(positions, hidden) (a draft layer's MultiHeadLatentAttention) on
latent_major –

W_UV [H, 512, 128] latent-contiguous, as DcpValue's MMA fragments load it
prepare –

Each layer's latent-contiguous W_UV (1 MiB per draft layer), at load.