vllm.models.kimi_k3.nvidia.ops.cute_dsl.latent_moe_tail ¶

class AdaptiveUpProjectionKernel:
    """Dispatch static-M Skinny or dynamic-M WGMMA into one mailbox."""

    def __init__(
        self,
        *,
        group: dist.ProcessGroup,
        rank: int,
        tp_size: int,
        latent_dim: int,
        hidden_dim: int,
        max_m: int,
        skinny_max_m: int,
        mma_tiler_mn: tuple[int, int],
        cluster_shape_mn: tuple[int, int],
        b_prime_stages: int,
    ) -> None:
        if hidden_dim % tp_size:
            raise ValueError("hidden_dim must be divisible by TP size")
        if not 0 <= skinny_max_m <= min(8, max_m):
            raise ValueError("skinny_max_m must be in [0, min(8, max_m)]")
        self.rank = rank
        self.tp_size = tp_size
        self.latent_dim = latent_dim
        self.hidden_dim = hidden_dim
        self.shard_dim = hidden_dim // tp_size
        self.max_m = max_m
        self.skinny_max_m = skinny_max_m
        self.mma_tiler_mn = mma_tiler_mn
        self.cluster_shape_mn = cluster_shape_mn
        self.b_prime_stages = b_prime_stages
        device = torch.device("cuda", torch.accelerator.current_device_index())
        self._device = device
        self._dynamic: Any | None = None
        self._skinny_by_m: dict[int, FusedAddMulticastSkinnyGemmKernel] = {}
        validate_configuration(
            latent_dim=latent_dim,
            shard_dim=self.shard_dim,
            mma_tiler_mn=mma_tiler_mn,
            cluster_shape_mn=cluster_shape_mn,
            b_prime_stages=b_prime_stages,
        )
        if skinny_max_m and self.latent_dim % (224 * 8):
            raise ValueError(
                "Skinny up-projection requires latent_dim divisible by 1792."
            )

        self._mailbox = symm_mem.empty(
            (1, max_m, hidden_dim),
            dtype=torch.bfloat16,
            device=device,
        )
        self._mailbox_symm_mem = symm_mem.rendezvous(self._mailbox, group)
        self._mailbox.view(torch.int32).fill_(-0x80000000)
        multicast_ptr = self._mailbox_symm_mem.multicast_ptr
        if multicast_ptr is None or multicast_ptr == 0:
            raise RuntimeError("mailbox NVLS multicast mapping is unavailable")
        self._mailbox_multicast_ptr = int(multicast_ptr)

        cluster_size = math.prod(cluster_shape_mn)
        self._max_active_clusters = utils.HardwareInfo().get_max_active_clusters(
            cluster_size
        )
        self._mailbox_c = _as_cute(self._mailbox)

    def compile_dynamic(self) -> None:
        if self._dynamic is not None:
            return
        device = self._device
        with torch.accelerator.device_index(device.index):
            compile_latent = torch.empty(
                (1, self.max_m, self.latent_dim),
                dtype=torch.bfloat16,
                device=device,
            )
            compile_weight = torch.empty(
                (self.shard_dim, self.latent_dim),
                dtype=torch.bfloat16,
                device=device,
            )
            compile_shared = torch.empty(
                (self.max_m, self.hidden_dim),
                dtype=torch.bfloat16,
                device=device,
            )[
                :,
                self.rank * self.shard_dim : (self.rank + 1) * self.shard_dim,
            ]
            compile_latent_c = _as_cute(compile_latent, dynamic_m=True)
            compile_weight_c = _as_cute(compile_weight.unsqueeze(0))
            compile_shared_c = _as_cute(compile_shared)

            self._dynamic = compile_kernel(
                (self.max_m, self.shard_dim, self.latent_dim, 1),
                compile_latent_c,
                compile_weight_c,
                self._mailbox_c,
                compile_shared_c,
                self.hidden_dim,
                self.shard_dim,
                self.mma_tiler_mn,
                self.cluster_shape_mn,
                self._max_active_clusters,
                self.b_prime_stages,
            )

    def compile_skinny(self, m: int) -> None:
        if not 1 <= m <= self.skinny_max_m:
            raise ValueError(
                f"Skinny up-projection requires M in [1, {self.skinny_max_m}]."
            )
        if m in self._skinny_by_m:
            return
        with torch.accelerator.device_index(self._device.index):
            self._skinny_by_m[m] = FusedAddMulticastSkinnyGemmKernel(
                rank=self.rank,
                tp_size=self.tp_size,
                latent_dim=self.latent_dim,
                hidden_dim=self.hidden_dim,
                num_rows=m,
            )

    def ensure_compiled(self, m: int) -> None:
        if not 1 <= m <= self.max_m:
            raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]")
        if m <= self.skinny_max_m:
            self.compile_skinny(m)
        else:
            self.compile_dynamic()

    def __call__(
        self,
        latent: torch.Tensor,
        weight: torch.Tensor,
        shared_shard: torch.Tensor,
    ) -> torch.Tensor:
        if latent.ndim != 2:
            raise ValueError("latent must be rank-2")
        m = latent.shape[0]
        device = self._mailbox.device
        expected = (
            (latent, (m, self.latent_dim), "latent"),
            (
                weight,
                (self.shard_dim, self.latent_dim),
                "weight",
            ),
            (
                shared_shard,
                (self.max_m, self.shard_dim),
                "shared_shard",
            ),
        )
        for tensor, shape, name in expected:
            if (
                tensor.shape != shape
                or tensor.dtype != torch.bfloat16
                or tensor.device != device
            ):
                raise ValueError(f"{name} must be CUDA torch.bfloat16 {list(shape)}")
        if (
            not latent.is_contiguous()
            or not weight.is_contiguous()
            or shared_shard.stride() != (self.hidden_dim, 1)
        ):
            raise ValueError("up-projection inputs have unsupported strides")
        if not 1 <= m <= self.max_m:
            raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]")

        if m <= self.skinny_max_m:
            skinny = self._skinny_by_m.get(m)
            if skinny is None:
                raise RuntimeError(
                    f"Skinny up-projection M={m} was not compiled before launch."
                )
            return skinny(
                latent,
                weight,
                shared_shard,
                self._mailbox,
                self._mailbox_multicast_ptr,
            )

        if self._dynamic is None:
            raise RuntimeError("Dynamic up-projection was not compiled before launch.")
        with torch.accelerator.device_index(device.index):
            stream = cuda.CUstream(torch.cuda.current_stream(device).cuda_stream)
            self._dynamic(
                _as_cute(latent.unsqueeze(0), dynamic_m=True),
                _as_cute(weight.unsqueeze(0)),
                self._mailbox_c,
                _as_cute(shared_shard),
                cutlass.Int64(m),
                cutlass.Int64(
                    self._mailbox_multicast_ptr + self.rank * self.shard_dim * 2
                ),
                stream,
            )
        return self._mailbox

class CollectiveKernel:
    """Own and launch the routed AllReduce/RMSNorm plus shared ReduceScatter."""

    def __init__(
        self,
        *,
        group: dist.ProcessGroup,
        rank: int,
        tp_size: int,
        latent_dim: int,
        hidden_dim: int,
        max_m: int,
        max_token_ctas: int,
        rms_eps: float,
        fp32_internal: bool,
    ) -> None:
        validate_shape(
            tp_size=tp_size,
            latent_dim=latent_dim,
            hidden_dim=hidden_dim,
        )
        self.rank = rank
        self.tp_size = tp_size
        self.latent_dim = latent_dim
        self.hidden_dim = hidden_dim
        self.shard_dim = hidden_dim // tp_size
        self.max_m = max_m
        self.max_token_ctas = max_token_ctas
        self.rms_eps = float(rms_eps)
        self.fp32_internal = fp32_internal
        device = torch.device("cuda", torch.accelerator.current_device_index())

        bytes_per_routed_buffer = max_m * tp_size * latent_dim * 2
        routed_bytes = NUM_LAMPORT_BUFFERS * bytes_per_routed_buffer
        self._routed_workspace = symm_mem.empty(
            routed_bytes // 4,
            dtype=torch.float32,
            device=device,
        )
        self._routed_symm_mem = symm_mem.rendezvous(self._routed_workspace, group)
        self._routed_workspace.fill_(-0.0)
        actual_bytes_per_buffer = (
            self._routed_symm_mem.buffer_size // NUM_LAMPORT_BUFFERS // 16 * 16
        )
        if actual_bytes_per_buffer < bytes_per_routed_buffer:
            raise RuntimeError("routed symmetric workspace is too small")
        self._routed_flags = torch.tensor(
            [0, 2, actual_bytes_per_buffer, 0, 0, 0, 0, 0, 0],
            dtype=torch.uint32,
            device=device,
        )
        routed_multicast_ptr = self._routed_symm_mem.multicast_ptr
        if routed_multicast_ptr is None or routed_multicast_ptr == 0:
            raise RuntimeError("routed NVLS multicast mapping is unavailable")
        self._routed_multicast_ptr = int(routed_multicast_ptr)

        self._latent_output = torch.empty(
            (max_m, latent_dim), dtype=torch.bfloat16, device=device
        )
        self._shared_output = torch.empty(
            (max_m, hidden_dim), dtype=torch.bfloat16, device=device
        )
        shard_start = rank * self.shard_dim
        shard_end = shard_start + self.shard_dim
        self._shared_shard = self._shared_output[:, shard_start:shard_end]

        self._shared_workspace = symm_mem.empty(
            (NUM_LAMPORT_BUFFERS, max_m, tp_size, self.shard_dim),
            dtype=torch.bfloat16,
            device=device,
        )
        self._shared_symm_mem = symm_mem.rendezvous(self._shared_workspace, group)
        self._shared_workspace.view(torch.int32).fill_(-0x80000000)
        self._shared_flags = torch.zeros(12, dtype=torch.int32, device=device)
        self._shared_flags[1] = 1
        self._shared_flags[2] = max_m * tp_size * self.shard_dim * 2
        peer_ptrs = [
            self._shared_symm_mem.get_buffer(
                peer,
                self._shared_workspace.shape,
                torch.bfloat16,
            ).data_ptr()
            for peer in range(tp_size)
        ]
        if any(pointer == 0 for pointer in peer_ptrs):
            raise RuntimeError("shared LSA peer mapping is unavailable")
        self._shared_peer_ptrs = torch.tensor(
            peer_ptrs, dtype=torch.int64, device=device
        )

        torch.accelerator.synchronize(device)
        dist.barrier(group=group, device_ids=[device.index])
        for owner in range(tp_size):
            if rank == owner:
                compile_kernel(
                    rank=rank,
                    tp_size=tp_size,
                    latent_dim=latent_dim,
                    hidden_dim=hidden_dim,
                    max_m=max_m,
                    max_token_ctas=max_token_ctas,
                    latent_output=self._latent_output,
                    routed_workspace=self._routed_workspace,
                    routed_flags=self._routed_flags,
                    routed_multicast_ptr=self._routed_multicast_ptr,
                    shared_output=self._shared_output,
                    shared_workspace=self._shared_workspace,
                    shared_flags=self._shared_flags,
                    shared_peer_ptrs=self._shared_peer_ptrs,
                    rms_eps=self.rms_eps,
                    fp32_internal=fp32_internal,
                )
            dist.barrier(group=group, device_ids=[device.index])

    def __call__(
        self,
        latent_source: torch.Tensor,
        shared_source: torch.Tensor,
        gamma: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if latent_source.ndim != 2 or shared_source.ndim != 2:
            raise ValueError("latent_source and shared_source must be rank-2")
        m = latent_source.shape[0]
        device = self._routed_workspace.device
        expected = (
            (latent_source, (m, self.latent_dim), "latent_source"),
            (shared_source, (m, self.hidden_dim), "shared_source"),
            (gamma, (self.latent_dim,), "gamma"),
        )
        for tensor, shape, name in expected:
            if (
                tensor.shape != shape
                or tensor.dtype != torch.bfloat16
                or tensor.device != device
                or not tensor.is_contiguous()
            ):
                raise ValueError(f"{name} must be contiguous CUDA BF16 {list(shape)}")
        if not 1 <= m <= self.max_m:
            raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]")

        with torch.accelerator.device_index(device.index):
            launch(
                latent_source,
                gamma,
                self._latent_output,
                self._routed_workspace,
                self._routed_flags,
                self._routed_multicast_ptr,
                shared_source,
                self._shared_output,
                self._shared_workspace,
                self._shared_flags,
                self._shared_peer_ptrs,
                self.rms_eps,
                rank=self.rank,
                tp_size=self.tp_size,
                latent_dim=self.latent_dim,
                hidden_dim=self.hidden_dim,
                max_m=self.max_m,
                max_token_ctas=self.max_token_ctas,
                fp32_internal=self.fp32_internal,
            )
        return (
            self._latent_output[:m],
            self._shared_shard,
        )

    @property
    def latent_output(self) -> torch.Tensor:
        return self._latent_output

    @property
    def shared_output(self) -> torch.Tensor:
        return self._shared_output

class LamportCopyKernel:
    """Copy a borrowed symmetric mailbox into a fresh local tensor."""

    def __init__(
        self,
        *,
        hidden_dim: int,
        max_m: int,
        ctas: int,
        threads: int,
    ) -> None:
        self.hidden_dim = hidden_dim
        self.max_m = max_m
        self.ctas = ctas
        self.threads = threads
        compile_kernel(
            hidden_dim,
            max_m,
            ctas,
            threads,
            torch.accelerator.current_device_index(),
        )

    def __call__(self, symmetric_mailbox: torch.Tensor, *, m: int) -> torch.Tensor:
        if not symmetric_mailbox.is_cuda:
            raise ValueError("symmetric_mailbox must be a CUDA tensor")
        device = symmetric_mailbox.device
        with torch.accelerator.device_index(device.index):
            output = torch.empty(
                (1, m, self.hidden_dim),
                dtype=torch.bfloat16,
                device=device,
            )
            launch(
                symmetric_mailbox,
                output,
                m=m,
                hidden_dim=self.hidden_dim,
                max_m=self.max_m,
                ctas=self.ctas,
                threads=self.threads,
            )
        return output