`vllm.models.muse_spark.nvidia.specialized_decode.cutedsl.attention.qkv_front` ¶

The Muse Spark attention front for decode batches (1 <= M <= 16) in one launch: the QKV + o-gate projection (split-K BF16 GEMM, mma.sync, FP32 accumulate), then the per-head weightless RMS norm, interleaved RoPE, the y store and the paged K/V insert. It writes what the projection GEMM followed by fused_muse_spark_qknorm_rope_kv_insert writes for a BF16 KV cache, up to the GEMM's summation order.

Structure: a thread-block cluster computes one head (64 output columns), its S CTAs splitting K into KSTRIPE-tile cyclic stripes. In each CTA a producer thread streams the W and x boxes into a shared-memory ring with TMA (128B swizzle), four consumer warps run mma.sync on them and an aux warp fetches the epilogue inputs. Each CTA sends its partial of token t to the owner CTA (rank t % S) with st.async; the owner sums the S partials in fixed rank order, rounds to BF16 and runs the epilogue. No global scratch, no atomics, no state across calls: the result depends on the launch shape only.

PDL: the kernel may start while its predecessor still runs. Before griddepcontrol.wait it reads only the weights, which no kernel writes inside a graph; every thread waits before it reads x, positions, slots or cos|sin and before any global store.

Classes:

QkvFrontKernel –

One variant: tiles 8-token column tiles, a stages-deep ring, K

Functions:

compile_qkv_front –

The compiled variant, shared by every layer and shape it serves, with
fixed_rows –

The kernel variant's row count: 8 specializes M == 8, 0 any other M.
hidden_rows_ok –

Whether t is a contiguous, 16-byte aligned bf16 [rows, 8192]
k_split –

The K split (cluster size) for an N-wide projection: one CTA per SM
kv_cache_layout_ok –

Whether a KV cache has the layout the QKV front and attn2 read: bf16
launch –

qkv_front without its checks: for a call unsupported took with
prepare –

Compile the variants that N columns and these row counts take (before
qkv_front –

Out [M, N] = x [M, 8192] @ weight [N, 8192]^T with the q and k heads
ring_stages –

The ring depth at K split split; it does not change the result.
supports_shape –

Whether a variant computes an N-wide projection with these KV heads.
token_tiles –

8-token mma column tiles for M tokens.
unsupported –

Why the kernel cannot take this call (with out, also its output),
variant –

(tiles, stages, splits, fixed_m) of the kernel for N columns, M rows.

`QkvFrontKernel` ¶

One variant: tiles 8-token column tiles, a stages-deep ring, K split splits (the cluster size) and fixed_m (8: M == 8, 0: any M of the tile count).

Methods:

aux –

Slots and cos|sin rows of this CTA's tokens into shared memory, and
consume –

Consumer warp warp: W rows 16 warp .. 16 warp + 15 over all 128 k
finish –

Epilogue of local token u (head columns 2 lane, 2 lane + 1): the
produce –

The producer thread: the weights of stages 1..npre-1 and an L2
send –

This CTA's partial of rows 16 warp + g (+8) for tokens 8 nt + 2 t