`vllm.models.kimi_k3.nvidia.specialized_decode.cutedsl.moe.ep.comm` ¶

Expert-parallel dispatch and combine of the latent MoE over NVLink symmetric memory (DEP: data-parallel attention at TP1, the 896 routed experts split over N ranks).

Every rank owns one symmetric buffer (EpLayout):

receive rows, source-major (row = source x stride(T) + token): MXFP8 latent (3584 B), its E8M0 block scales (112 B, the linear layout of flashinfer.mxfp8_quantize), the token's 16 expert ids and FP32 routing weights, and a source word (one 128-byte line per row, -1 when empty);
combine slots [T, N, 3584] BF16: the per-destination-rank partial of every own token;
control words: arrival counters and the dispatch / combine flags (one line each).

EpDispatch (one CTA per token) quantizes the BF16 latent to MXFP8 bit-identically to mxfp8_quantize (two threads per 32-block, 16-byte stores) and writes it, with ids, weights and source, into its static row on every rank holding one of its experts. Empty rows keep non-local ids (weight 0), so the expert GEMM over the N T rows skips them. The last CTA of the rank raises its dispatch flag on every rank and waits for all N; on completion the rank's receive rows are complete.

EpCombine (G CTAs) sends every received row's BF16 output (the routing-weighted sum over this rank's experts) to its source's slot [token, rank] and re-arms the row; the last CTA raises the combine flag on every rank. The CTAs with own tokens then wait for all N flags and reduce: the FP32 sum of the destination partials in rank order, rounded once to BF16 (the routed latent), then vLLM's RMSNorm of it (FP32 sum of squares, rsqrt(mean + eps), x * rstd * w rounded once).

Measured on GB300 NVL72 (16 ranks over 4 nodes), these rules decide the latency: no two ranks write one 128-byte line (row blocks padded to 8 rows, one line per flag and per source word: 200 -> 40 us per dispatch + combine at T=5); a rank reaches its own buffer through its local mapping, never its peer mapping; flags are polled by one lane each (a local poll costs a round trip). Every thread that stored remotely fences at system scope before the CTA's arrival is counted (a single fence after the CTA barrier lost partials at 32 ranks), and every reader fences after the flags.

Flags are single-reader: a writer sets 1, the reader re-arms 0. No rank can raise the next layer's flag before the reader re-armed it: the next dispatch from a source needs the source's combine to finish, which needs this rank's combine flag, raised after this rank's dispatch re-armed; the next combine flag needs this rank's next dispatch, issued after its combine re-armed.

Classes:

EpBuffers –

One rank's symmetric EP buffer, its peers' base addresses and tensor views of
EpCombine –

(out [N T, 3584] BF16 per received row, ids [T, 16] own routing, norm weight
EpDispatch –

(x [T, 3584] BF16, ids [T, 16] int32, weights [T, 16] FP32, peers [N] int64) ->
EpLayout –

Byte offsets of one rank's symmetric EP buffer for ranks ranks and at most

Functions:

last_arrival –

Warp-uniform: whether this CTA is the total-th arrival on counter
raise_flags –

Lane d sets this rank's flag on rank d (all ranks in parallel).
rms_rstd –

rsqrtf(total / 3584 + eps) of vLLM's rms_norm (IEEE divide).
stride –

Receive rows per source rank: a multiple of 8, so no 128-byte line of the
wait_flags –

Warp: lane s polls flag s (relaxed, all in flight together) until every flag

`EpBuffers` ¶

One rank's symmetric EP buffer, its peers' base addresses and tensor views of the receive rows (the expert kernel's inputs).