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

Muse Spark paged decode attention of the 8-token verify step, fused with the per-head output RMS norm (no weight) and the sigmoid gate, as one CuTe kernel.

head_dim 64, 8 query heads per KV head, BF16 paged cache [pages, HKV, 16, 128] (K in columns 0:64, V in 64:128), 8 query rows per request (the last 8 cache entries of the request), any context length L >= 8, optionally a sliding window. One launch covers the rank's KV heads (grid y); the heads are independent.

Two variants, picked on the host from the deployment's maximum length:

One cluster (max_seq_len <= 2048, no window): a cluster of 16 CTAs of 4 warps owns 16 query-head rows of one request and one KV head. Warp w of CTA r attends pages 4 r + w + 64 i (the next page's registers load while one is attended) with an online softmax: mma.sync m16n8k16 BF16 -> FP32 for Q K^T and P V, the softmax weights rounded to BF16, V transposed in registers by movmatrix (head dims permuted: logical (j, 2 t + e) = 16 t + 2 j + e). The warps merge in shared memory, each CTA pushes its merged rows to the CTA that owns them over distributed shared memory (st.async, mbarrier transaction bytes), and the owner merges the 16 CTAs in rank order, rounds o / l to BF16 and applies the RMS norm and the gate.
Split (longer contexts or a window): a request within 8 rounds of the cluster's warps (512 pages from row 0's first visible key) takes the path above in z slice 0, the other slices exit at once; a longer one spreads its pages over 4 num_splits clusters, each CTA streaming its pages through an 8-page cp.async ring, merges in the cluster over distributed shared memory and across clusters through a global workspace in cluster order (an arrival counter per row: the last arriver merges, then resets the counter to zero). The launch shape depends on the deployment only, so a captured graph serves every step. With window W, row t sees keys [L - 8 + t - W + 1, L - 8 + t].

Every FP32 operation is explicit PTX in the fast-math forms (subnormals flushed, ex2.approx, rcp.approx, rsqrt.approx) in a fixed order, so the result depends neither on the launch's timing nor on the device. Not bitwise to the stock decode + rmsnorm_sigmoid_gate pair: the softmax weights are rounded to BF16 before P V and the partials merge in FP32.

Functions:

attn2 –

``out [8 R, 512 kv_heads] = sigmoid(gate) * rmsnorm(attention(q,
compile_attn2 –

cute.compile one variant of the kernel for kv_heads KV heads,
counters_bytes –

The workspace's uint32 counters, one per row, 16-byte padded.
prepare –

Compile the variant a deployment's layer launches (build time).
split_forms –

(merge form, rolled short-path loop) of the split kernel.
split_kernel –

Whether a deployment of max_seq_len launches the split kernel.
window_pages –

The most pages from row 0's first visible key to a request's last
workspace_bytes –

Workspace bytes of a split launch over rows (request, KV head,

`_Attn2Kernel` ¶

Grid (16 * 4 R, kv_heads, num_splits), cluster (16, 1, 1), 128 threads: cluster b * 4 + rg of KV head blockIdx.y computes rows 16 rg + [0, 16) of request b's 64 (token = row // 8, head = row % 8). The variant, PDL, the KV heads, the z slices and the window are compile-time; shapes, strides and addresses are read at the call.