vllm.models.inkling.nvidia.ops.lamport ¶

class LamportRSConv:
    """Persistent symmetric buffers for one TP group."""

    def _initialize_lamport_buffers(self) -> None:
        """Arm every slot and collectively verify the exact sentinel bits.

        Do not replace this with a zero-fill: Lamport readiness distinguishes
        the bf16 bit pattern for -0.0 (0x8000) from every published payload.
        The validation is deliberately collective so one rank cannot enter the
        first polling kernel while another rank still has an unarmed buffer.
        """
        if not self.buf_in.is_contiguous() or not self.buf_out.is_contiguous():
            raise RuntimeError("Lamport symmetric buffers must be contiguous")
        if self.buf_in.numel() % 2 or self.buf_out.numel() % 2:
            raise RuntimeError("Lamport buffers must contain whole uint32 pairs")

        # Fill through int16 so each bf16 lane receives the exact -0.0 bits.
        # This covers all three generations and all max-token slots, including
        # slots that a smaller first invocation does not touch.
        self.buf_in.view(torch.int16).fill_(-0x8000)
        self.buf_out.view(torch.int16).fill_(-0x8000)
        torch.accelerator.synchronize(self.device)
        self.tp.barrier()

        # Scan the complete local allocations.  Since every rank performs the
        # scan and participates in MAX, success proves every symmetric backing
        # allocation was armed before any rank is allowed to use generation 0.
        bad = torch.logical_or(
            self.buf_in.view(torch.int16).ne(-0x8000).any(),
            self.buf_out.view(torch.int16).ne(-0x8000).any(),
        ).to(dtype=torch.float32)
        bad = self.tp.all_reduce(bad)
        if int(bad.item()) != 0:
            raise RuntimeError("Lamport sentinel initialization failed on a TP rank")
        self.tp.barrier()

    def _initialize_mnnvl_buffers(self) -> None:
        """Initialize and validate FlashInfer fabric-mapped Lamport storage."""
        self._mnnvl_input_handle.lamport_initialize(self.rank, torch.bfloat16)
        self._mnnvl_output_handle.lamport_initialize(self.rank, torch.bfloat16)
        torch.accelerator.synchronize(self.device)
        self.tp.barrier()

        bad = torch.zeros((), dtype=torch.int32, device=self.device)
        num_pairs = self.num_buffers * self.max_tokens * self.hidden_size // 2
        _validate_lamport_init_kernel[(triton.cdiv(num_pairs, 256),)](
            self.input_peer_ptrs,
            self.output_peer_ptrs,
            bad,
            num_pairs,
            RANK=self.rank,
            BLOCK=256,
            num_warps=4,
        )
        bad = self.tp.all_reduce(bad.to(torch.float32))
        if int(bad.item()) != 0:
            raise RuntimeError("MNNVL Lamport sentinel initialization failed")
        self.tp.barrier()

    def __init__(
        self, hidden_size: int, window_size: int, max_tokens: int = _MAX_TOKENS
    ) -> None:
        import torch.distributed._symmetric_memory as symm_mem

        tp = get_tp_group()
        self.tp = tp
        self.group = tp.device_group
        self.world_size = tp.world_size
        self.rank = tp.rank_in_group
        self.device = torch.device(tp.device)
        if self.world_size not in (2, 4, 8):
            raise ValueError(f"TP world size must be 2, 4, or 8, got {self.world_size}")
        if hidden_size % (2 * self.world_size) != 0:
            raise ValueError("hidden size must produce an even shard on every rank")
        if window_size != 4:
            raise ValueError(f"short-conv window size must be 4, got {window_size}")
        if max_tokens < 1 or max_tokens > _MAX_TOKENS:
            raise ValueError(f"max_tokens must be in [1, {_MAX_TOKENS}]")

        is_cross_node = not all(in_the_same_node_as(tp.cpu_group))
        self.hidden_size = hidden_size
        self.window_size = window_size
        self.max_tokens = max_tokens
        self.shard_size = hidden_size // self.world_size
        # Three generations follow FlashInfer's Lamport layout.  A generation
        # is reused only after two intervening collective calls.
        self.num_buffers = 3
        self.input_generation_bytes = max_tokens * hidden_size * 2
        self.output_generation_bytes = max_tokens * hidden_size * 2
        self.use_pdl = torch.cuda.get_device_capability(self.device)[0] >= 9
        if is_cross_node:
            try:
                from flashinfer.comm.mnnvl import (
                    McastGPUBuffer,
                    TorchDistBackend,
                    is_mnnvl_fabric_supported,
                )
            except ImportError as error:
                raise RuntimeError(
                    "cross-node TP requires FlashInfer MNNVL support"
                ) from error

            local_supported = int(
                is_mnnvl_fabric_supported(torch.accelerator.current_device_index())
            )
            unsupported = torch.tensor(
                1 - local_supported, dtype=torch.float32, device=self.device
            )
            unsupported = tp.all_reduce(unsupported)
            if int(unsupported.item()) != 0:
                raise RuntimeError("cross-node TP is supported only on MNNVL fabric")

            comm_backend = TorchDistBackend(self.group)
            allocation_bytes = self.num_buffers * max_tokens * hidden_size * 2
            self._mnnvl_input_handle = McastGPUBuffer(
                allocation_bytes,
                self.world_size,
                self.rank,
                self.device,
                comm_backend,
            )
            self._mnnvl_output_handle = McastGPUBuffer(
                allocation_bytes,
                self.world_size,
                self.rank,
                self.device,
                comm_backend,
            )
            self.input_peer_ptrs = self._mnnvl_input_handle.get_buffer_ptrs_dev()
            self.output_peer_ptrs = self._mnnvl_output_handle.get_buffer_ptrs_dev()
            self._initialize_mnnvl_buffers()
            logger.info("using FlashInfer fabric-mapped MNNVL Lamport buffers")
        else:
            self.buf_in = symm_mem.empty(
                self.num_buffers,
                max_tokens,
                self.world_size,
                self.shard_size,
                dtype=torch.bfloat16,
                device=self.device,
            )
            self.buf_out = symm_mem.empty(
                self.num_buffers,
                max_tokens,
                hidden_size,
                dtype=torch.bfloat16,
                device=self.device,
            )
            group_name = self.group.group_name
            input_handle = symm_mem.rendezvous(self.buf_in, group_name)
            output_handle = symm_mem.rendezvous(self.buf_out, group_name)
            self.input_peer_ptrs = input_handle.buffer_ptrs_dev
            self.output_peer_ptrs = output_handle.buffer_ptrs_dev
            self._input_handle = input_handle
            self._output_handle = output_handle
            self._initialize_lamport_buffers()
        self.generation = 0

    def usable(self, num_tokens: int) -> bool:
        return 0 < num_tokens <= self.max_tokens

    def rs_sconv_ag_add_norm(
        self,
        input_tensor: torch.Tensor,
        residual: torch.Tensor,
        conv_weight: torch.Tensor,
        norm_weight: torch.Tensor | None,
        eps: float,
        cache: torch.Tensor,
        positions: torch.Tensor,
        block_table: torch.Tensor,
        seq_idx: torch.Tensor,
        slot_mapping: torch.Tensor,
        off_s: int,
        ws: int,
        block_size: int,
    ) -> tuple[torch.Tensor | None, torch.Tensor]:
        """Return ``(normed | None, new_residual)``, both shaped ``[T, 6144]``."""
        tokens, hidden_size = residual.shape
        if not self.usable(tokens):
            raise ValueError(f"num_tokens must be in [1, {self.max_tokens}]")
        if hidden_size != self.hidden_size or residual.dtype != torch.bfloat16:
            raise ValueError("residual must be bf16 [T, 6144]")
        if (
            input_tensor.shape != residual.shape
            or input_tensor.dtype != torch.bfloat16
            or input_tensor.stride(1) != 1
        ):
            raise ValueError("input_tensor must be channel-contiguous bf16 [T, 6144]")
        shard_size = hidden_size // self.world_size
        if conv_weight.shape != (shard_size, self.window_size):
            raise ValueError(
                f"conv_weight must have shape [{shard_size}, {self.window_size}]"
            )
        if (
            conv_weight.dtype != torch.bfloat16
            or conv_weight.stride(0) != self.window_size
        ):
            raise ValueError("conv_weight must be contiguous bf16")
        if norm_weight is not None and (
            norm_weight.shape != (hidden_size,) or norm_weight.dtype != torch.bfloat16
        ):
            raise ValueError("norm_weight must be bf16 [6144] or None")
        if cache.dtype != torch.bfloat16 or cache.ndim != 4:
            raise ValueError("cache must be a 4-D bf16 tensor")
        if shard_size % ws != 0 or cache.shape[1] != shard_size // ws:
            raise ValueError("cache head layout is inconsistent with ws")
        if off_s < 0 or off_s + ws > cache.shape[3]:
            raise ValueError("cache channel offset is out of bounds")

        index = self.generation
        input_offset = index * self.input_generation_bytes // 4
        output_offset = index * self.output_generation_bytes // 4
        normed = torch.empty_like(residual) if norm_weight is not None else None
        residual_out = torch.empty_like(residual)
        phase_splits = 8 if tokens == 1 and shard_size % 8 == 0 else 1
        phase_tile_p2 = triton.next_power_of_2(shard_size // phase_splits)
        phase_grid = (tokens, phase_splits)
        # Wide CTAs win before the grid saturates; smaller CTAs reduce register
        # pressure once high-throughput batches provide enough parallelism.
        if 128 <= tokens <= 2048:
            phase_warps = 16
        elif tokens > 2048:
            phase_warps = 8
        else:
            phase_warps = 4
        gather_warps = 4 if tokens >= 256 else 8
        _publish_input_kernel[phase_grid](
            input_tensor,
            self.input_peer_ptrs,
            input_offset,
            input_tensor.stride(0),
            C=hidden_size,
            CS=shard_size,
            CS_P2=phase_tile_p2,
            SPLITS=phase_splits,
            RANK=self.rank,
            WORLD=self.world_size,
            USE_PDL=self.use_pdl,
            launch_pdl=self.use_pdl,
            num_warps=phase_warps,
        )
        _reduce_insert_kernel[phase_grid](
            self.input_peer_ptrs,
            input_offset,
            cache,
            slot_mapping,
            cache.stride(0),
            cache.stride(1),
            cache.stride(2),
            cache.stride(3),
            block_size,
            C=hidden_size,
            CS=shard_size,
            CS_P2=phase_tile_p2,
            SPLITS=phase_splits,
            HEAD_SIZE=ws,
            CACHE_OFFSET=off_s,
            RANK=self.rank,
            WORLD=self.world_size,
            USE_PDL=self.use_pdl,
            launch_pdl=self.use_pdl,
            num_warps=phase_warps,
        )
        _sconv_publish_kernel[phase_grid](
            self.output_peer_ptrs,
            output_offset,
            residual,
            conv_weight,
            cache,
            positions,
            seq_idx,
            slot_mapping,
            block_table,
            residual.stride(0),
            cache.stride(0),
            cache.stride(1),
            cache.stride(2),
            cache.stride(3),
            block_table.stride(0),
            block_table.shape[1],
            block_size,
            C=hidden_size,
            CS=shard_size,
            CS_P2=phase_tile_p2,
            SPLITS=phase_splits,
            HEAD_SIZE=ws,
            CACHE_OFFSET=off_s,
            RANK=self.rank,
            WORLD=self.world_size,
            WINDOW=self.window_size,
            USE_PDL=self.use_pdl,
            launch_pdl=self.use_pdl,
            num_warps=phase_warps,
        )
        _gather_norm_kernel[(tokens,)](
            self.output_peer_ptrs,
            output_offset,
            norm_weight if norm_weight is not None else residual,
            normed if normed is not None else residual_out,
            residual_out,
            eps,
            residual.stride(0),
            C=hidden_size,
            C_P2=triton.next_power_of_2(hidden_size),
            RANK=self.rank,
            HAS_NORM=norm_weight is not None,
            USE_PDL=self.use_pdl,
            launch_pdl=self.use_pdl,
            num_warps=gather_warps,
        )
        self.generation = (index + 1) % self.num_buffers
        return normed, residual_out