from dataclasses import dataclass from enum import Enum, auto from typing import Optional import torch class CutlassMoEType(Enum): """ Enum for the different types of cutlass moe operations that are currently supported in SGLang. """ BlockscaledFP8 = auto() BlockscaledFP4 = auto() @dataclass class CutlassMoEParams: """ Parameters for the cutlass moe operation. """ # Type as defined above cutlass_moe_type: CutlassMoEType # Strides for activations, weights and output in logical number of elements. # The activations & output stride is the number of elements to the next row. # The weights stride is the number of elements to the next row per expert. # For example, if the weight is [e, n, k], then the b_stride is a tensor of # shape [e] with each element being k. Similarly for activations, if the # shape is [m, k], then the a_stride has shape [e] with each value k. # Similarly for output, if the output is [m, n], then the c_stride is a # tensor of shape [e] with each element being k. # Note: cutlass_fp4_group_mm is designed to accept the strides of # activations and weights to be the same, so it is passed in as a single # tensor. # ab_strides_13: [e] dtype: int64 [Gemm 1: Activation / Weight strides] # ab_strides_2: [e] dtype: int64 [Gemm 2: Activation / Weight strides] # c_strides_13: [e] dtype: int64 [Gemm 1: Output Strides] # c_strides_2: [e] dtype: int64 [Gemm 2: Output Strides] ab_strides_13: torch.Tensor ab_strides_2: torch.Tensor c_strides_13: torch.Tensor c_strides_2: torch.Tensor # m: Total number of tokens # n: intermediate size per partition # k: hidden size per expert # e: Number of experts # device: Device to run computation on and store tensors m: int intermediate_size_per_partition: int hidden_size: int num_experts: int device: torch.device # Pointers container for calculating offsets of the input activations for each expert # a_ptrs: [e] dtype: int64 a_ptrs: torch.Tensor # Pointers container for calculating offsets of the input weights for each expert # b_ptrs: [e] dtype: int64 b_ptrs: torch.Tensor # Pointers container for calculating offsets of the output activations for each expert # out_ptrs: [e] dtype: int64 out_ptrs: torch.Tensor # Pointers container for calculating offsets of the input scales for each expert # a_scales_ptrs: [e] dtype: int64 # b_scales_ptrs: [e] dtype: int64 a_scales_ptrs: torch.Tensor b_scales_ptrs: torch.Tensor # Pointers for per-expert alpha values alpha_ptrs: torch.Tensor # CUTLASS blockscale layouts for A and B operands layout_sfa: torch.Tensor layout_sfb: torch.Tensor # Offsets that mark at which token index each expert begins its computation # The number of tokens computed with expert E is expert_offsets[E + 1] - expert_offsets[E] # expert_offsets: [e+1] dtype: int32 expert_offsets: torch.Tensor # Problem size: (num_experts, (m,2n,k)) for first GEMM # problem_sizes1: [e, 3] dtype: int32 # Problem size: (num_experts, (m,n,k)) for second GEMM # problem_sizes2: [e, 3] dtype: int32 problem_sizes1: torch.Tensor problem_sizes2: torch.Tensor # Similar to expert_offsets, but for blockscales for FP4 blockscaled Group GEMM blockscale_offsets: Optional[torch.Tensor] = None def __init__( self, cutlass_moe_type: CutlassMoEType, device: torch.device, num_experts: int, intermediate_size_per_partition: int, hidden_size: int, ): self.cutlass_moe_type = cutlass_moe_type self.device = device self.num_experts = num_experts self.intermediate_size_per_partition = intermediate_size_per_partition self.hidden_size = hidden_size self.n = self.intermediate_size_per_partition self.k = self.hidden_size self.e = self.num_experts self.ab_strides_13 = torch.full( (self.e,), self.k, dtype=torch.int64, device=self.device ) self.ab_strides_2 = torch.full( (self.e,), self.n, dtype=torch.int64, device=self.device ) self.c_strides_13 = torch.full( (self.e,), 2 * self.n, dtype=torch.int64, device=self.device ) self.c_strides_2 = torch.full( (self.e,), self.k, dtype=torch.int64, device=self.device ) self.expert_offsets = torch.empty( (self.e + 1,), dtype=torch.int32, device=self.device ) self.problem_sizes1 = torch.empty( (self.e, 3), dtype=torch.int32, device=self.device ) self.problem_sizes2 = torch.empty( (self.e, 3), dtype=torch.int32, device=self.device ) if self.cutlass_moe_type == CutlassMoEType.BlockscaledFP4: self.blockscale_offsets = torch.empty( (self.e + 1,), dtype=torch.int32, device=self.device ) else: self.blockscale_offsets = None self.a_ptrs = torch.empty((self.e,), dtype=torch.int64, device=self.device) self.b_ptrs = torch.empty((self.e,), dtype=torch.int64, device=self.device) self.out_ptrs = torch.empty((self.e,), dtype=torch.int64, device=self.device) self.a_scales_ptrs = torch.empty( (self.e,), dtype=torch.int64, device=self.device ) self.b_scales_ptrs = torch.empty( (self.e,), dtype=torch.int64, device=self.device ) self.alpha_ptrs = torch.empty((self.e,), dtype=torch.int64, device=self.device) self.layout_sfa = torch.empty( (self.e, 5), dtype=torch.int64, device=self.device ) self.layout_sfb = torch.empty( (self.e, 5), dtype=torch.int64, device=self.device ) def to_gemm1_args(self) -> dict: return { "ab_strides": self.ab_strides_13, "c_strides": self.c_strides_13, "problem_sizes": self.problem_sizes1, "expert_offsets": self.expert_offsets[:-1], "blockscale_offsets": self.blockscale_offsets[:-1], "a_ptrs": self.a_ptrs, "b_ptrs": self.b_ptrs, "out_ptrs": self.out_ptrs, "a_scales_ptrs": self.a_scales_ptrs, "b_scales_ptrs": self.b_scales_ptrs, "alpha_ptrs": self.alpha_ptrs, "layout_sfa": self.layout_sfa, "layout_sfb": self.layout_sfb, } def to_gemm2_args(self) -> dict: return { "ab_strides": self.ab_strides_2, "c_strides": self.c_strides_2, "problem_sizes": self.problem_sizes2, "expert_offsets": self.expert_offsets[:-1], "blockscale_offsets": self.blockscale_offsets[:-1], "a_ptrs": self.a_ptrs, "b_ptrs": self.b_ptrs, "out_ptrs": self.out_ptrs, "a_scales_ptrs": self.a_scales_ptrs, "b_scales_ptrs": self.b_scales_ptrs, "alpha_ptrs": self.alpha_ptrs, "layout_sfa": self.layout_sfa, "layout_sfb": self.layout_sfb, }