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

class GemmRsAr:
    """Own the symmetric workspace for Kimi-K3 GEMM-RS/AR launches.

    All TP ranks must belong to one NVLink domain for multimem instructions.

    Each instance is bound to either RS or AR. A vLLM worker has one static
    sequence-parallel topology, so the process-wide singleton only needs one
    mode. Two independent mode-specific singletons would lift that restriction
    but duplicate the large symmetric workspace. A future mixed-mode design
    should instead use lightweight RS/AR frontends over one shared multicast
    workspace; that is outside this integration's current scope.
    """

    def __init__(self, *, max_M: int, N: int, all_reduce: bool = False) -> None:
        tp_group = get_tp_group()
        group = tp_group.device_group
        rank = tp_group.rank_in_group
        world_size = tp_group.world_size
        device = torch.device("cuda", torch.accelerator.current_device_index())

        assert 1 < world_size <= 16
        assert 128 % world_size == 0
        assert max_M >= 128 and N % 128 == 0

        max_M = (max_M + world_size - 1) // world_size * world_size
        self.rank = rank
        self.world_size = world_size
        self.max_M = max_M
        self.N = N
        self.device = device
        self.all_reduce = all_reduce

        self.partial = symm_mem.empty((max_M, N), dtype=torch.bfloat16, device=device)
        self.partial_handle = symm_mem.rendezvous(self.partial, group)
        if self.partial_handle.multicast_ptr == 0:
            raise RuntimeError("GEMM-RS/AR requires NVLink multicast memory")
        self.partial_mc_ptr = make_ptr(
            BFloat16,
            self.partial_handle.multicast_ptr,
            cute.AddressSpace.gmem,
            assumed_align=32,
        )

        grid_m = (max_M + 127) // 128
        cta_group = 2 if max_M >= 1024 or grid_m % 2 == 0 else 1
        grid_m = (grid_m + cta_group - 1) // cta_group * cta_group
        self.num_sms = torch.cuda.get_device_properties(device).multi_processor_count
        max_flags = grid_m * (N // 128) + self.num_sms
        self.flags = symm_mem.empty(max_flags, dtype=torch.int32, device=device)
        self.flags_handle = symm_mem.rendezvous(self.flags, group)
        if self.flags_handle.multicast_ptr == 0:
            raise RuntimeError("GEMM-RS/AR requires NVLink multicast memory")
        self.flags.zero_()
        self.flags_mc_ptr = make_ptr(
            Int32,
            self.flags_handle.multicast_ptr,
            cute.AddressSpace.gmem,
            assumed_align=16,
        )
        self.peer_flag_ptr = make_ptr(
            Int64,
            self.flags_handle.buffer_ptrs_dev,
            cute.AddressSpace.gmem,
            assumed_align=8,
        )

        torch.accelerator.synchronize(device)
        tp_group.barrier()

    def can_run(self, linear: LinearBase) -> bool:
        # Validate projection-invariant requirements once during model init.
        # only supports BF16 for now
        if not isinstance(linear.quant_method, UnquantizedLinearMethod):
            return False
        w = linear.weight
        if w.ndim != 2:
            return False
        K = w.shape[1]
        return (
            w.shape == (self.N, K)
            and K % 64 == 0
            and w.dtype == torch.bfloat16
            and w.device == self.device
            and w.is_contiguous()
        )

    def should_run(self, x: torch.Tensor) -> bool:
        # Use the same threshold for RS and AR for now. Small-M shapes are
        # supported but faster on the existing LL path.
        return x.shape[0] >= 128

    def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
        assert x.ndim == 2
        M, K = x.shape
        assert 0 < M <= self.max_M
        assert w.shape == (self.N, K) and K % 64 == 0
        assert w.dtype == torch.bfloat16
        assert w.device == self.device
        assert w.is_contiguous()
        assert x.dtype == torch.bfloat16
        assert x.device == self.device
        assert x.is_contiguous()
        N = w.shape[0]
        padded_M = (M + self.world_size - 1) // self.world_size
        padded_M *= self.world_size
        local_M = padded_M // self.world_size

        grid_m = (M + 127) // 128
        # Avoid padding small odd grids; 2-CTA wins consistently for M >= 1024.
        cta_group = 2 if M >= 1024 or grid_m % 2 == 0 else 1
        grid_m = (grid_m + cta_group - 1) // cta_group * cta_group
        BN = 256 if M * K >= 24 * 1024 * 1024 else 128
        assert N % BN == 0

        num_tiles = grid_m * (N // BN)
        num_ctas = min(num_tiles, self.num_sms)
        num_ctas = num_ctas // cta_group * cta_group
        assert self.flags.numel() >= num_tiles + num_ctas

        output = None
        if not self.all_reduce:
            output = torch.empty((local_M, N), dtype=torch.bfloat16, device=self.device)
        compiled = Sm100GemmRsArBF16.compile(
            self.rank,
            self.world_size,
            BN,
            cta_group,
            self.all_reduce,
        )
        compiled(
            x,
            w,
            self.partial[:padded_M],
            self.partial_mc_ptr,
            output,
            self.flags,
            self.flags_mc_ptr,
            self.peer_flag_ptr,
            num_ctas,
        )
        if self.all_reduce:
            # AttnRes may retain output past the next workspace reuse.
            # A future kernel could overlap this copy using an extra warp or
            # the communication warp.
            return self.partial[:M].clone()
        assert output is not None
        return output