Skip to content

vllm.models.deepseek_v41.nvidia.specialized_decode.logits.logits

Logits of the specialized decode: the low-batch LM head GEMM (lm_head) and vocab-parallel logits gathered with the TP group's custom all-gather.

For decode batches the one-shot custom all-gather is several times faster than the NCCL ring. The kernel falls back to NCCL per call by dtype and size alone, so every TP rank takes the same path. It shares its staging buffers and signals with the group's other custom collectives, so the logits must be gathered on the forward stream after any side stream has rejoined it.