vllm.models.glm5next.nvidia.specialized_decode.logits_gather

All-gather of the vocab-parallel LM head logits straight into [M, vocab].

Stock gathers the [M, vocab / tp] shards with NCCL into [tp, M, vocab / tp] and permutes that into [M, vocab]: ~62 us of all-gather and ~19 us of copy at 64 rows (the verify step of MTP3 at 16 requests, or 64 requests without MTP). Here each rank stores its shard into column block rank of every rank's symmetric [M, vocab] buffer over NVLink, then every CTA meets its counterparts on the other ranks at a flag barrier, and the caller reads the local buffer. Unicast stores to the peers beat one NVLS multicast store per vector (34 vs 39 us at 64 rows, destination out of L2).

Two buffers (and flag sets) alternate: a rank starting gather k + 2 has passed gather k + 1's barrier, which its peers reach only after their streams finished with gather k's logits.

Classes:

• LogitsGather –

  tensor_model_parallel_all_gather of [M, vocab shard] BF16 logits for

LogitsGather

tensor_model_parallel_all_gather of [M, vocab shard] BF16 logits for M <= capacity, returned as a view of a buffer that the next gather but one overwrites. Construction is collective over the TP group.

Methods:

• __call__ –

  Eager calls only: the host alternates the buffers.

Source code in vllm/models/glm5next/nvidia/specialized_decode/logits_gather.py

class LogitsGather:
    """`tensor_model_parallel_all_gather` of [M, vocab shard] BF16 logits for
    M <= capacity, returned as a view of a buffer that the next gather but one
    overwrites. Construction is collective over the TP group."""

    def __init__(self, *, columns: int, capacity: int, device: torch.device):
        group = get_tp_group().device_group
        rank = dist.get_rank(group)
        tp = dist.get_world_size(group)
        device = torch.device(device)
        self.columns, self.capacity = columns, capacity
        ctas = torch.cuda.get_device_properties(device).multi_processor_count
        with torch.accelerator.device_index(device.index):
            self._buffer = symm_mem.empty(
                (BUFFERS, capacity, tp * columns), dtype=torch.bfloat16, device=device
            )
            self._flags = symm_mem.empty(
                BUFFERS * ctas, dtype=torch.int32, device=device
            )
            self._flags.zero_()
            buffer = symm_mem.rendezvous(self._buffer, group)
            flags = symm_mem.rendezvous(self._flags, group)
            self._flags_multicast = int(flags.multicast_ptr or 0)
            if not self._flags_multicast:
                raise RuntimeError("NVLink multicast mapping is unavailable")
            self._handles = (buffer, flags)
            self._kernel = cute.compile(
                _GatherKernel(
                    columns=columns,
                    rank=rank,
                    capacity=capacity,
                    ctas=ctas,
                    buffers=tuple(int(p) for p in buffer.buffer_ptrs),
                ),
                make_fake_dynamic_compact_tensor(
                    BFloat16, alignment=16, divisibility=columns
                ),
                make_fake_compact_tensor(Int32, (BUFFERS * ctas,), assumed_align=4),
                Int64(0),
                Int32(0),
                Int32(capacity),
                current_cu_stream(),
            )
            torch.accelerator.synchronize()
        dist.barrier(group=group)
        self._parity = 0

    def __call__(self, logits: torch.Tensor) -> torch.Tensor:
        """Eager calls only: the host alternates the buffers."""
        m = logits.shape[0]
        assert 1 <= m <= self.capacity and logits.shape[1] == self.columns
        assert logits.dtype == torch.bfloat16 and logits.is_contiguous()
        parity, self._parity = self._parity, self._parity ^ 1
        self._kernel(
            to_cute_dynamic(logits.flatten(), 16, divisibility=self.columns),
            to_cute(self._flags, 4),
            Int64(self._flags_multicast),
            Int32(parity),
            Int32(m),
            current_cu_stream(),
        )
        return self._buffer[parity, :m]

__call__(logits)

Eager calls only: the host alternates the buffers.

Source code in vllm/models/glm5next/nvidia/specialized_decode/logits_gather.py

def __call__(self, logits: torch.Tensor) -> torch.Tensor:
    """Eager calls only: the host alternates the buffers."""
    m = logits.shape[0]
    assert 1 <= m <= self.capacity and logits.shape[1] == self.columns
    assert logits.dtype == torch.bfloat16 and logits.is_contiguous()
    parity, self._parity = self._parity, self._parity ^ 1
    self._kernel(
        to_cute_dynamic(logits.flatten(), 16, divisibility=self.columns),
        to_cute(self._flags, 4),
        Int64(self._flags_multicast),
        Int32(parity),
        Int32(m),
        current_cu_stream(),
    )
    return self._buffer[parity, :m]

_arrive(address)

Publishes this CTA's stores and counts it in every rank's flag.

Source code in vllm/models/glm5next/nvidia/specialized_decode/logits_gather.py

@asm()
def _arrive(address):
    """Publishes this CTA's stores and counts it in every rank's flag."""
    code = "fence.acq_rel.sys; multimem.red.release.sys.global.add.u32 [$0], 1;"
    return code, "l,~{memory}", (Int64(address),)