`vllm.models.deepseek_v41.nvidia.specialized_decode.attention.sparse_mla_cute` ¶

The sparse-MLA decode of small batches, in CuTe DSL.

The problem. In a decode step every query token (DeepSeek-V4's MLA after the query projection: 16 heads a rank at TP4, FP8 E4M3, head dim 512) attends to its own sparse set of KV rows. A row is 512 FP8 bytes and serves as both K and V. The set holds up to 640 columns: the 128 sliding-window rows of one KV pool, then the indexer's top-k rows of the compressed pool, all addressed as storage offsets into their pool, with -1 for an unused column and a per-token count of active columns. TRTLLM-gen splits a token's 640 rows over at most 3 CTAs, so at one or two requests (6 or 12 tokens with DSpark) its kernel runs about 10 µs on a GPU of 152 SMs; most of that is the latency of one CTA's serial work.

What this module computes. sparse_mla_decode writes, for each token t and head h, the BF16 attention output over the token's active columns c (c < lens[t] and indices[t, c] >= 0)::

s_c = bmm1_scale * q[t, h] . kv[c]
p_c = exp(s_c - m)
out[t, h] = bmm2_scale * sum_c p_c kv[c] / (sum_c p_c + exp(sink_h - m))

where kv[c] is row indices[t, c] of the sliding-window pool for the first 128 columns and of the compressed pool for the others, m is the maximum of the scores and the sink, and the sink term is omitted without sinks. A token with no active column gets zeros. This is what trtllm_batch_decode_sparse_mla_dsv4 computes for these shapes.

How. One CTA of 128 threads takes one token and one 128-column chunk of its sparse set (chunk 0 = the sliding-window columns); the n = ceil(width / 128) chunks of a token form a thread block cluster, so a 640-wide layer runs 5 CTAs a token in parallel::

gather the chunk's KV rows and the query (cp.async, one group per 128-byte K
block) into 128B-swizzled tiles
  └► S^T = KV x Q^T on tcgen05 (M 128 rows, N 16 heads, K 512, FP8 -> FP32 TMEM)
       └► chunk softmax per head in base 2: max m_c, sum l_c, and
          P^T = 256 * exp2(S - m_c) as FP8 (scaled: small weights stay normal)
            └► O^T = V^T x P^T on tcgen05, V the same tile read MN-major
               (four MMAs of 128 head dims) ...
            └► ... while the cluster exchanges (m_c, l_c) per head (st.async)
                 └► each CTA scales its partial O by exp2(m_c - M) / L
                    (M, L over the cluster's chunks and the sink)
                      └► block j of the 512 head dims goes to CTA j % n, which
                         sums the n partials and stores BF16

The partial blocks travel by bulk copies between the CTAs' shared memory into the receiving CTA's KV tile, whose MMAs are done by then: a cluster barrier after each CTA's second MMA separates the two uses. A last cluster barrier keeps every CTA resident until the bulk copies reading its shared memory are complete.

Synchronization. The kernel reads the sinks (weights) before its grid dependency wait and everything else after it, and lets its dependents launch once its S MMAs are issued (programmatic dependent launch). Masked columns gather zeros and score -inf.

Scope. Two CTAs fit an SM (about 78 KB of shared memory each), so a launch is one wave only while every token's cluster can be resident: _max_tokens allows 2 x the device's resident clusters of n CTAs at one CTA an SM (on a 152-SM GB300, 8 requests with DSpark at width 640). Larger batches keep TRTLLM-gen, which is faster once the clusters run in waves.

Classes:

SparseMlaDecodeCute –

The sparse-MLA decode (module docstring) of tokens whose sparse sets

Functions:

bulk_copy_to_peer –

Bulk-copy size bytes of this CTA's shared memory at source to
copy16_shared_to_global –

Copy 16 bytes from this CTA's shared memory to global memory.
cp_async_16 –

Copy 16 bytes from global to this CTA's shared memory (cp.async).
ld_global_f32 –

A float from read-only global memory.
ld_global_s32 –

A 32-bit integer from global memory an earlier grid wrote (side-effecting,
mma_f8 –

One tcgen05 FP8 MMA of this CTA: D (TMEM) = A x B (+ D if
pack_e4m3x4 –

Four floats as E4M3 bytes, a in the lowest byte (round to nearest,
precompile –

Compile the kernels of DeepSeek-V4's two widths (128 sliding-window
sparse_mla_decode –

Launch SparseMlaDecodeCute on the current stream (see
st_async_peer_b32 –

Store a word at smem of cluster CTA peer (the same offset as a
st_shared_bf16 –

Store value rounded to BF16 at a shared address.
st_shared_v4 –

Store four 32-bit words at a 16-byte aligned shared address.
st_shared_zero16 –

Store 16 zero bytes at a 16-byte aligned shared address.
supports_sparse_mla_cute –

Whether sparse_mla_decode takes these arguments: 16 FP8 query heads

`SparseMlaDecodeCute` ¶

The sparse-MLA decode (module docstring) of tokens whose sparse sets are at most chunks x 128 columns wide; has_sinks: add each head's attention sink to the softmax denominator.

Methods:

__call__ –

Launch on stream: one cluster of n CTAs a token. Pointers