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

Decode top-k of the dense indexers straight from their BF16 logits, in CuTe DSL.

The problem. In a decode step each dense indexer layer scores every compressed token of a row (SM100's paged MQA logits are BF16) and keeps the k largest: the token top-512, and in the candidate source also the top-2048 of the rows' 8-token block maxima. cooperative_topk reads FP32, so the logits used to be upcast first, and it builds a radix histogram of every value with shared-memory atomics. Real indexer logits crowd a few histogram bins (a row of 150K logits holds about 200 distinct BF16 values), where those atomics serialize.

What this module computes. DecodeTopK writes, for each row, the indices of its k largest elements in any order (exact: every index whose value exceeds the k-th largest, then enough of those equal to it), or 0 .. n-1 and -1 padding when the row has at most k elements. An element is one BF16 logit (token mode), or the NaN-propagating maximum of an aligned group of 8 logits (block mode, the candidate blocks), with the row's newest block pinned at +inf and blocks scoring -inf written as -1. Values are compared through an order-preserving 16-bit key of their BF16 bits, so ties are exact.

How. One thread block cluster of C CTAs per row: 8 CTAs of 512 threads while every row's cluster can be resident at once, else 4 of 1024 (a second wave of clusters would nearly double the time). Each CTA copies one contiguous chunk of the row into shared memory, as 16-byte vectors (8 logits, or one block); the chunks split the row's valid length, not its width. A thread stages 8 vectors and keeps them in registers too, so a cluster holds 262,144 logits. Wider rows need more: at a 1M context the ratio-2 layers' logits are 524,288 wide, and a second build stages 16 vectors a thread at 8 CTAs, the other 8 in shared memory only. The candidate source layer scores every token (ratio 1), 1,048,576 at a 1M context, so a third build runs 16 CTAs a row (a non-portable cluster size, which CuTe DSL allows on every kernel) with 16 vectors a thread. Rows of the wide builds run in waves when they outnumber the resident clusters; the 16-CTA build takes at most two waves (supports_decode_topk). Two ways to the answer:

Sampled (token rows longer than one vector a thread)::

sample the row ──► every CTA: the same cutoff (the q-th largest of the row's 2048 samples, q aiming at 3k survivors) └► keep keys > cutoff and up to k keys == cutoff, push both to CTA 0 └► CTA 0: exact k-th largest among <= 8192 survivors (two 8-bit radix passes), write k indices
Direct (block mode, short token rows, and a sample that misjudged the row: more than 8192 survivors, or fewer than k at or above the cutoff): the exact k-th largest key from two cluster-wide 8-bit radix histograms, then every CTA writes its own selected indices straight to the output.

Both filter the same way: a thread tests each of its vectors' maximum against the cutoff, the vectors that can hold selected keys go to a CTA queue, and the CTA's threads unpack the queued vectors evenly; a CTA scan gives each thread its output slots and one cluster atomic reserves the CTA's range, so no key costs an atomic. Whole-row histograms spread each warp's atomics over 8 copies (lane l counts into copy l mod 8), since a crowded bin serializes the atomics on its word; the radix selects over the few thousand samples or survivors count into one histogram, where the copies' zeroing and summing cost more than they save.

Synchronization. Cluster barriers separate the phases: every CTA has started and zeroed its counters before any distributed shared memory access; samples are complete before the cutoff; survivors are complete before CTA 0 reads them; histogram contributions are complete before they are read; and CTA 0 waits until every peer is done with its shared memory before it exits. The kernel waits on its predecessor grid before reading inputs and lets its dependents launch once its chunk is in shared memory (programmatic dependent launch).

Classes:

DecodeTopK –

The decode top-k (module docstring) of rows of BF16 logits.

Functions:

block_key –

The key of the NaN-propagating maximum of eight BF16 values.
cluster_add_u64 –

Atomic add to an 8-byte aligned cluster shared word pair; returns the
cluster_for –

CTAs a row: 8 while every row's cluster is resident at once, else 4.
cluster_load –

Load a cluster shared word.
cluster_red –

Atomic add to a cluster shared word.
cluster_store_u16 –

Store the low half of a word to cluster shared memory.
cluster_store_u32 –

Store a 32-bit word to cluster shared memory.
decode_topk –

Launch DecodeTopK on the current stream.
hist_add –

Count each valid lane's bin_ with one shared atomic. copies:
load_u32x4 –

Four 32-bit words at a 16-byte aligned global address.
max_nan_bf16x2 –

Lane-wise maximum of two BF16 pairs, NaN if either is NaN.
order_key –

The order-preserving key of the BF16 value in the low half of bits:
peer_address –

The cluster address of this CTA's shared address in CTA rank.
precompile –

Compile the DecodeTopK variants decode_topk launches on rows of
select_bin –

Warp 0: the bin b of a 256-bin histogram holding the q-th
shared_load_u32x4 –

Four 32-bit words at a 16-byte aligned shared address of this CTA.
shared_red –

Atomic add to a shared word of this CTA.
shared_store_u32x4 –

Store four 32-bit words at a 16-byte aligned shared address of this CTA.
shfl_up –

value of the lane delta below (own value in the lowest lanes).
store_s32 –

Store a 32-bit word to global memory.
supports_decode_topk –

Whether decode_topk takes these BF16 logits (16-byte aligned rows
threads_for –

Threads a CTA: 1024 at 4 CTAs a row, else 512.
vector_choices –

The 16-byte vectors a thread may stage: 8, or 16 at 512 threads (1024
vectors_for –

The fewest staged vectors a thread whose cluster holds a row of

`DecodeTopK` ¶

The decode top-k (module docstring) of rows of BF16 logits.

k: 512, 1024 or 2048; cluster: 4, 8 or 16 CTAs a row (4: 1024 threads, else 512); vector_count: 16-byte vectors a thread stages (8, or 16 at 8 or 16 CTAs: the first 8 also stay in registers), so a row holds at most cluster * threads * vector_count vectors; block_mode: elements are 8-logit block maxima.

Methods:

__call__ –

logits: BF16 rows row_bytes apart (16-byte aligned), each
count_queued –

This thread's keys above and equal to cutoff in its share of
count_vector –

Count a staged vector (words, its first element first) into
direct –

Direct path: the exact cutoff, then every CTA writes its keys
emit –

CTA 0, sampled path: write the row's k indices from the gathered
exact_cutoff –

The row's exact k-th largest key and how many keys equal to it
gather –

Sampled path: push this CTA's keys above cutoff and up to k
kernel –

One cluster a row (module docstring).
queue_vectors –

Queue this CTA's vectors that can hold a key >= cutoff (by
radix_select –

The q-th largest of keys[0:count] (shared Uint16, q <= count)
reserve –

This thread's first slots for its gt / eq keys in the
scan2 –

Exclusive CTA prefix sums of a and b over the threads, and
vector_top –

The largest key of a staged vector whose first element is
vector_words –

The four words of the chunk's vector vector in shared memory.
word_key –

The key of element e of a vector's words (element: its