`vllm.models.kimi_k3.nvidia.specialized_decode.cutedsl.dsa.topk_exchange` ¶

DCP top-k candidate exchange over NVLink and its exact merge (K3.1 DSA decode).

The #63 one-round merge (vllm/models/kimi_k3/nvidia/ops/dcp_topk_merge) all-gathers every rank's local top-N B' rows with NCCL (a 1-CTA ring all-gather of ~19 us at DCP8 plus its launch gap), then selects (one kernel) and rescues saturated rows (a second kernel). Here, instead:

push IndexerLocalTopk's epilogue (indexer.py, exchange) stores its row's B' words into every rank's symmetric TopkExchange slot [parity, source rank, row] as NCCL-LL pairs: each 16-byte store carries two words and the call's epoch twice, (w0, e, w1, e), so a reader checks each 8-byte half's epoch instead of a flag behind a system fence. With LL64 each half is one 64-bit element (the word in the low 32 bits, the epoch in the high) of a v2.b64 store, read with one 64-bit load, both relaxed at system scope: the PTX memory model treats a vector access as one access per element, and two morally strong accesses to the same 8 bytes are single-copy atomic, so a half whose epoch matches holds that call's word. (The 32-bit form stores and loads the word and the epoch as separate elements.) Both forms write the same bytes. merge ExchangeMerge (one 512-thread CTA per row, launched early by PDL) polls the W ranks' LL words of its row, then runs the #63 merge's own select (_make_radix_select: the same keys, threshold and tie rule, so the same ids and rows in the same order), its saturation word and saturated-rank push (the rescue workspace of init_dcp_topk_merge, the same seq2 / flag protocol), and the rescue's second round in the same CTA.

MERGE_SELECT picks, per DCP world size and rows, the "digits" form of the merge's first round instead (the same rows, counts, saturation word and rescue, bit for bit; laid out for latency): warp w of the CTA owns source rank w. It issues every LL load of that rank before it checks a tag, so each rank's words are decoded as they land, and adds the keys to an 11-bit histogram of the score; two digit scans (11 bits over all warps, then the low 5 bits) give the K-th score. Ties at it are ranked by global id: by ballots against the listed ties when there are few, else by a block-wide two-digit select of the tie ids. Only the warp of this rank's own source then writes the rows: the #63 select writes a thread's keys in thread order and this rank's N keys sit on N consecutive flat slots (thread = slot mod 512, N <= 512), so their order is a rotation of the slot order. A row with saturated ranks moves the round-1 keys into the #63 thread order and runs the same rescue code.

The epoch of row t is epoch[t] + 1 (epoch: per rank, bumped by the merge after it read the row); every rank runs the same calls over the same rows, so the epochs agree. The slot parity is the epoch's low bit: a rank can store call n + 1's words while a peer still reads call n's.

Classes:

ExchangeMerge –

Host side of the exchange merge of one (rows, block size).
TopkExchange –

The symmetric LL slots of the exchange over the DCP group: int32

Functions:

build_merge –

ExchangeMerge of rows rows (see the module docstring); digits:
compile_merge –

The compiled ExchangeMerge of rows rows (shared by every layer).
lanes_below –

The lanes l with 2 l < x.
ll_load –

(word, epoch) of the LL half at address.
ll_load_pair –

(w0, epoch, w1, epoch) of the LL store at address.
ll_offset –

Byte offset of B' word j in a row's LL pairs.
ll_pairs –

16-byte LL stores per B' row (two words each).
ll_push –

One LL store: words w0, w1 tagged with epoch.
load_ll_half –

(word, epoch) of one LL half: one 64-bit relaxed system-scope load.
load_ll_pair –

(w0, epoch, w1, epoch) of one LL store: one 16-byte relaxed system-scope
merge_select –

The first round of the merge of rows rows: "digits" where
reserve –

The device's TopkExchange of at least rows rows (collective on
store_ll_pairs –

Halves (w0, epoch), (w1, epoch) as the two 64-bit elements of one 16-byte