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

class LamportCopy:
    """Consume the local physical copy of an NVLS-multicast mailbox."""

    def __init__(self, hidden_dim: int, ctas: int, threads: int):
        self.hidden_dim = hidden_dim
        self.ctas = ctas
        self.threads = threads

    @cute.jit
    def __call__(
        self,
        symmetric_mailbox: cute.Tensor,
        local_output: cute.Tensor,
        m: cutlass.Int32,
        stream: cuda.CUstream,
    ):
        self.kernel(symmetric_mailbox, local_output, m).launch(
            grid=(self.ctas, 1, 1),
            block=(self.threads, 1, 1),
            stream=stream,
            use_pdl=True,
        )

    @cute.kernel
    def kernel(
        self,
        symmetric_mailbox: cute.Tensor,
        local_output: cute.Tensor,
        m: cutlass.Int32,
    ):
        # The CTA may be scheduled early, but mailbox inspection must not pass
        # the producer GEMM's programmatic completion point.
        cute.arch.griddepcontrol_wait()

        tidx, _, _ = cute.arch.thread_idx()
        block, _, _ = cute.arch.block_idx()
        thread = cutlass.Int64(block * self.threads + tidx)
        stride = cutlass.Int64(self.ctas * self.threads)
        fragments = cutlass.Int64(m) * cutlass.Int64(self.hidden_dim // VEC_BF16)

        fragment = thread
        while fragment < fragments:
            element = fragment * VEC_BF16
            source = cute.make_ptr(
                cutlass.BFloat16,
                (symmetric_mailbox.iterator + element).llvm_ptr,
                cute.AddressSpace.gmem,
                assumed_align=16,
            )
            packed = load_global_u32x4(source, volatile=True)
            while fragment_is_dirty(packed):
                packed = load_global_u32x4(source, volatile=True)

            destination = cutlass.Int64((local_output.iterator + element).toint())
            store_global_u32x4(destination, packed, volatile=False)
            fragment = fragment + stride

        # The returned ordinary tensor is complete. A same-stream successor
        # may overlap the mailbox cleanup below.
        cute.arch.griddepcontrol_launch_dependents()

        fragment = thread
        while fragment < fragments:
            element = fragment * VEC_BF16
            source = cute.make_ptr(
                cutlass.BFloat16,
                (symmetric_mailbox.iterator + element).llvm_ptr,
                cute.AddressSpace.gmem,
                assumed_align=16,
            )
            store_lamport_sentinel_128(source)
            fragment = fragment + stride

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