"""Operators enabled by external modules.""" from typing import List, Literal, Tuple # noqa: UP035 import tvm from tvm.relax.frontend import nn from tvm.script import ir as I from tvm.script import tirx as T try: import triton import triton.language as tl except ImportError: triton = None tl = None # We use a wrapper function to avoid type annotation issue of "tl.constexpr" when # triton is not installed. def _get_triton_w8a8_block_fp8_gemm(): # Triton kernel adapted from SGLang project # https://github.com/sgl-project/sglang/blob/v0.4.4/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py# noqa: E501 def _triton_w8a8_block_fp8_gemm( # Pointers to inputs and output A, B, C, As, Bs, # Shape for matmul M, N: tl.constexpr, K: tl.constexpr, # Stride for inputs and output stride_am: tl.constexpr, stride_ak: tl.constexpr, stride_bk: tl.constexpr, stride_bn: tl.constexpr, stride_cm: tl.constexpr, stride_cn: tl.constexpr, stride_As_m: tl.constexpr, stride_As_k: tl.constexpr, stride_Bs_k: tl.constexpr, stride_Bs_n: tl.constexpr, # Block size for block-wise quantization group_n: tl.constexpr, group_k: tl.constexpr, # Meta-parameters BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, ): """Triton-accelerated function used to perform linear operations (dot product) on input tensors `A` and `B` with block-wise quantization, and store the result in output tensor `C`. """ pid = tl.program_id(axis=0).to(tl.int64) num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) num_pid_in_group = GROUP_SIZE_M * num_pid_n group_id = pid // num_pid_in_group first_pid_m = group_id * GROUP_SIZE_M group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) pid_m = first_pid_m + (pid % group_size_m) pid_n = (pid % num_pid_in_group) // group_size_m offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N offs_k = tl.arange(0, BLOCK_SIZE_K) a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) As_ptrs = As + offs_am * stride_As_m offs_bsn = offs_bn // group_n Bs_ptrs = Bs + offs_bsn * stride_Bs_n accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) k_start = k * BLOCK_SIZE_K offs_ks = k_start // group_k a_s = tl.load(As_ptrs + offs_ks * stride_As_k) b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k) accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :] a_ptrs += BLOCK_SIZE_K * stride_ak b_ptrs += BLOCK_SIZE_K * stride_bk if C.dtype.element_ty == tl.bfloat16: c = accumulator.to(tl.bfloat16) elif C.dtype.element_ty == tl.float16: c = accumulator.to(tl.float16) else: c = accumulator.to(tl.float32) offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) tl.store(c_ptrs, c, mask=c_mask) return _triton_w8a8_block_fp8_gemm # We use a wrapper function to avoid type annotation issue of "tl.constexpr" when # triton is not installed. def _get_triton_w8a8_block_fp8_group_gemm(): # Triton kernel adapted from SGLang project # https://github.com/sgl-project/sglang/blob/v0.4.4/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py# noqa: E501 def _triton_w8a8_block_fp8_group_gemm( # Pointers to matrices a_ptr, b_ptr, c_ptr, a_scale_ptr, b_scale_ptr, expert_ids_ptr, indptr_ptr, # Matrix dimensions EM, N: tl.constexpr, K: tl.constexpr, num_experts: tl.constexpr, # The stride variables represent how much to increase the ptr by when # moving by 1 element in a particular dimension. E.g. `stride_am` is # how much to increase `a_ptr` by to get the element one row down # (A has M rows). stride_am: tl.constexpr, stride_ak: tl.constexpr, stride_be: tl.constexpr, stride_bk: tl.constexpr, stride_bn: tl.constexpr, stride_cm: tl.constexpr, stride_cn: tl.constexpr, stride_asm: tl.constexpr, stride_ask: tl.constexpr, stride_bse: tl.constexpr, stride_bsk: tl.constexpr, stride_bsn: tl.constexpr, # Block size for block-wise quantization group_n: tl.constexpr, group_k: tl.constexpr, # Meta-parameters BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, even_Ks: tl.constexpr, ): """ Implements the fused computation for a Mixture of Experts (MOE) using token and expert matrices. Key Parameters: - A: The input tensor representing tokens with shape (*, K), where '*' can be any shape representing batches and K is the feature dimension of each token. - B: The stacked MOE weight tensor with shape (E, N, K), where E is the number of experts, K is the input feature dimension, and N is the output feature dimension. - C: The output cache tensor with shape (*, N), where '*' means the same shape as the input tensor A, and N is the output feature dimension. - expert_ids: A tensor containing the indices of the expert for each block. It determines which expert matrix from B should be used for each block in A. This kernel performs the multiplication of a token by its corresponding expert matrix as determined by `expert_ids`. """ # ----------------------------------------------------------- # Map program ids `pid` to the block of C it should compute. # This is done in a grouped ordering to promote L2 data reuse. pid = tl.program_id(axis=0).to(tl.int64) num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M) + num_experts num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) num_pid_in_group = GROUP_SIZE_M * num_pid_n group_id = pid // num_pid_in_group first_pid_m = group_id * GROUP_SIZE_M group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) pid_n = (pid % num_pid_in_group) // group_size_m # ---------------------------------------------------------- # Create pointers for the first blocks of A and B. # We will advance this pointer as we move in the K direction # and accumulate # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers expert_id = tl.load(expert_ids_ptr + pid_m).to(tl.int64) if expert_id == -1: return token_begin = tl.load(indptr_ptr + expert_id) token_end = tl.load(indptr_ptr + expert_id + 1) start_pid_m = tl.cdiv(token_begin, BLOCK_SIZE_M) + expert_id offs_token_id = ( token_begin + (pid_m - start_pid_m) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) ) token_mask = offs_token_id < token_end offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N offs_k = tl.arange(0, BLOCK_SIZE_K) a_ptrs = a_ptr + offs_token_id[:, None] * stride_am + offs_k[None, :] * stride_ak b_ptrs = ( b_ptr + expert_id * stride_be + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) ) a_scale_ptrs = a_scale_ptr + offs_token_id * stride_asm offs_bsn = offs_bn // group_n b_scale_ptrs = b_scale_ptr + expert_id * stride_bse + offs_bsn * stride_bsn # ----------------------------------------------------------- # Iterate to compute a block of the C matrix. # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block # of fp32 values for higher accuracy. # `accumulator` will be converted back to fp16 after the loop. accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): # Load the next block of A and B, generate a mask by checking the # K dimension. if even_Ks: a = tl.load( a_ptrs, mask=token_mask[:, None], other=0.0, ) b = tl.load(b_ptrs) else: a = tl.load( a_ptrs, mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), other=0.0, ) b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) # We accumulate along the K dimension. k_start = k * BLOCK_SIZE_K offs_ks = k_start // group_k a_scale = tl.load(a_scale_ptrs + offs_ks * stride_ask, mask=token_mask, other=0.0) b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk) accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :] # Advance the ptrs to the next K block. a_ptrs += BLOCK_SIZE_K * stride_ak b_ptrs += BLOCK_SIZE_K * stride_bk if c_ptr.dtype.element_ty == tl.bfloat16: accumulator = accumulator.to(tl.bfloat16) elif c_ptr.dtype.element_ty == tl.float16: accumulator = accumulator.to(tl.float16) else: accumulator = accumulator.to(tl.float32) # ----------------------------------------------------------- # Write back the block of the output offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) c_ptrs = c_ptr + stride_cm * offs_token_id[:, None] + stride_cn * offs_cn[None, :] c_mask = token_mask[:, None] & (offs_cn[None, :] < N) tl.store(c_ptrs, accumulator, mask=c_mask) return _triton_w8a8_block_fp8_group_gemm def get_tir_w8a8_block_fp8_matmul( N: int, K: int, block_n: int, block_k: int, in_dtype: Literal["float8_e4m3fn"], out_dtype: Literal["float16", "bfloat16"], BLOCK_SIZE_M: int, BLOCK_SIZE_N: int, BLOCK_SIZE_K: int, GROUP_SIZE_M: int, num_warps: int, num_stages: int, extern_mods: List[tvm.runtime.Module], # noqa: UP006 ): """Get the TIR function for the w8a8_block_fp8_matmul kernel.""" # NOTE: adding the type annotation of " -> Tuple[Optional[tvm.tirx.PrimFunc], str]" # will cause the failure of the type resolution in mypy. if triton is None: raise RuntimeError("Triton is not installed. Please install it with `pip install triton`.") name_suffix = f"_N{N}_K{K}_block_n{block_n}_block_k{block_k}_in{in_dtype}_out{out_dtype}" kernel_name = f"triton_w8a8_block_fp8_gemm{name_suffix}" tir_name = f"tir_w8a8_block_fp8_matmul{name_suffix}" for ext_mod in extern_mods: if ext_mod.implements_function(kernel_name): return [None, tir_name] triton_kernel = _get_triton_w8a8_block_fp8_gemm() triton_kernel.__name__ = kernel_name @I.ir_module class BlockFP8Matmul: @T.prim_func(private=True, s_tir=True) def tir_w8a8_block_fp8_matmul( var_A: T.handle, var_B: T.handle, var_As: T.handle, var_Bs: T.handle, var_C: T.handle, ): T.func_attr({"op_pattern": 8, "tirx.is_scheduled": 1}) M = T.int32() A = T.match_buffer(var_A, (M, K), dtype=in_dtype) B = T.match_buffer(var_B, (N, K), dtype=in_dtype) As = T.match_buffer(var_As, (M, (K + block_k - 1) // block_k), "float32") Bs = T.match_buffer( var_Bs, ((N + block_n - 1) // block_n, (K + block_k - 1) // block_k), "float32", ) C = T.match_buffer(var_C, (M, N), dtype=out_dtype) with T.sblock("root"): T.reads( A[0:M, 0:K], B[0:N, 0:K], As[0:M, 0 : (K + block_k - 1) // block_k], Bs[ 0 : (N + block_n - 1) // block_n, 0 : (K + block_k - 1) // block_k, ], ) T.writes(C[0:M, 0:N]) T.call_kernel( triton.jit(triton_kernel), (T.ceildiv(M, BLOCK_SIZE_M) * T.ceildiv(N, BLOCK_SIZE_N),), A.data, B.data, C.data, As.data, Bs.data, M, N, K, K, # stride_am 1, # stride_ak 1, # stride_bk K, # stride_bn N, # stride_cm 1, # stride_cn (K + block_k - 1) // block_k, # stride_As_m 1, # stride_As_k 1, # stride_Bs_k (K + block_k - 1) // block_k, # stride_Bs_n block_n, block_k, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M, num_warps=num_warps, num_stages=num_stages, ) new_ext_mods = BlockFP8Matmul.attrs["external_mods"] assert len(new_ext_mods) == 1 extern_mods.append(new_ext_mods[0]) return BlockFP8Matmul["tir_w8a8_block_fp8_matmul"], tir_name def get_tir_w8a8_block_fp8_group_matmul( N: int, K: int, num_experts: int, block_n: int, block_k: int, in_dtype: Literal["float8_e4m3fn"], out_dtype: Literal["float16", "bfloat16"], BLOCK_SIZE_M: int, BLOCK_SIZE_N: int, BLOCK_SIZE_K: int, GROUP_SIZE_M: int, num_warps: int, num_stages: int, extern_mods: List[tvm.runtime.Module], # noqa: UP006 ): """Get the TIR function for the w8a8_block_fp8_group_gemm kernel.""" if triton is None: raise RuntimeError("Triton is not installed. Please install it with `pip install triton`.") name_suffix = ( f"_N{N}_K{K}_num_experts{num_experts}_block_n{block_n}" f"_block_k{block_k}_in{in_dtype}_out{out_dtype}" ) kernel_name = f"triton_w8a8_block_fp8_group_gemm{name_suffix}" tir_name = f"tir_w8a8_block_fp8_group_gemm{name_suffix}" for ext_mod in extern_mods: if ext_mod.implements_function(kernel_name): return [None, tir_name] triton_kernel = _get_triton_w8a8_block_fp8_group_gemm() triton_kernel.__name__ = kernel_name @I.ir_module class BlockFP8GroupMatmul: @T.prim_func(private=True, s_tir=True) def tir_w8a8_block_fp8_group_gemm( var_A: T.handle, var_B: T.handle, var_As: T.handle, var_Bs: T.handle, var_expert_ids: T.handle, var_indptr: T.handle, var_C: T.handle, ): T.func_attr({"op_pattern": 8, "tirx.is_scheduled": 1}) EM = T.int32() A = T.match_buffer(var_A, (EM, K), dtype=in_dtype) B = T.match_buffer(var_B, (num_experts, N, K), dtype=in_dtype) As = T.match_buffer(var_As, (EM, (K + block_k - 1) // block_k), "float32") Bs = T.match_buffer( var_Bs, ( num_experts, (N + block_n - 1) // block_n, (K + block_k - 1) // block_k, ), "float32", ) expert_ids = T.match_buffer( var_expert_ids, ((EM + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + num_experts,), "int32", ) indptr = T.match_buffer(var_indptr, (num_experts + 1,), "int32") C = T.match_buffer(var_C, (EM, N), dtype=out_dtype) with T.sblock("root"): T.reads( A[0:EM, 0:K], B[0:num_experts, 0:N, 0:K], As[0:EM, 0 : (K + block_k - 1) // block_k], Bs[ 0:num_experts, 0 : (N + block_n - 1) // block_n, 0 : (K + block_k - 1) // block_k, ], expert_ids[0 : (EM + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + num_experts], indptr[0 : num_experts + 1], ) T.writes(C[0:EM, 0:N]) T.call_kernel( triton.jit(triton_kernel), ((T.ceildiv(EM, BLOCK_SIZE_M) + num_experts) * T.ceildiv(N, BLOCK_SIZE_N),), A.data, B.data, C.data, As.data, Bs.data, expert_ids.data, indptr.data, EM, N, K, num_experts, K, # stride_am 1, # stride_ak N * K, # stride_be 1, # stride_bk K, # stride_bn N, # stride_cm 1, # stride_cn (K + block_k - 1) // block_k, # stride_asm 1, # stride_ask ((N + block_n - 1) // block_n) * ((K + block_k - 1) // block_k), # stride_bse 1, # stride_bsk (K + block_k - 1) // block_k, # stride_Bs_n block_n, block_k, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M, K % BLOCK_SIZE_K == 0, num_warps=num_warps, num_stages=num_stages, ) new_ext_mods = BlockFP8GroupMatmul.attrs["external_mods"] assert len(new_ext_mods) == 1 extern_mods.append(new_ext_mods[0]) return BlockFP8GroupMatmul["tir_w8a8_block_fp8_group_gemm"], tir_name def _compute_expert_id_per_block( indptr: nn.Tensor, num_experts: int, M: nn.IntExpr, BLOCK_SIZE_M: int, ) -> nn.Tensor: """Compute the expert id for each threadblock (CTA). We assign an expert id to each threadblock, and the threadblock will compute the gemm with regard to the specified expert. Parameters ---------- indptr : nn.Tensor The indptr tensor of group gemm, with shape of [num_experts + 1,]. num_experts : int The number of total experts. M : nn.IntExpr The number of tokens. BLOCK_SIZE_M : int The block size of the threadblock along the batch dimension. Returns ------- expert_ids : nn.Tensor The expert id for each threadblock, with shape of [(M + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + num_experts,]. """ @T.prim_func(s_tir=True) def tir_compute_expert_id_per_block( var_indptr: T.handle, var_expert_ids: T.handle, M: T.int64, ): T.func_attr({"op_pattern": 8, "tirx.is_scheduled": 1}) indptr = T.match_buffer(var_indptr, (num_experts + 1,), "int32") expert_ids = T.match_buffer( var_expert_ids, ((M + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + num_experts,), "int32", ) with T.sblock("root"): for eid in T.thread_binding(0, num_experts, thread="threadIdx.x"): start_block_id: T.int32 = (indptr[eid] + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + eid num_blocks: T.int32 = ( indptr[eid + 1] - indptr[eid] + BLOCK_SIZE_M - 1 ) // BLOCK_SIZE_M start_block_id_next_expert: T.int32 = ( (indptr[eid + 1] + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + eid + 1 ) for block_id in T.serial(num_blocks): expert_ids[start_block_id + block_id] = eid for block_id in T.serial( start_block_id_next_expert - (start_block_id + num_blocks) ): expert_ids[start_block_id + num_blocks + block_id] = -1 assert num_experts <= 1024 return nn.tensor_ir_op( tir_compute_expert_id_per_block, "tir_compute_expert_id_per_block", args=[indptr, M], out=nn.Tensor.placeholder( ((M + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + num_experts,), dtype="int32" ), ) def fp8_groupwise_scaled_gemm( x: nn.Tensor, x_scale: nn.Tensor, weight: nn.Tensor, weight_scale: nn.Tensor, block_size: Tuple[int, int], # noqa: UP006 out_dtype: str, ) -> nn.Tensor: """Triton block-scale fp8 gemm operator. Parameters ---------- x : nn.Tensor The input tensor, with shape of [m, k]. x_scale : nn.Tensor The scale tensor, with shape of [m, k // block_size]. weight : nn.Tensor The weight tensor, with shape of [n, k]. weight_scale : nn.Tensor The scale tensor, with shape of [n // block_size, k // block_size]. block_size : Tuple[int, int] The block size. out_dtype : str The data type of the output tensor. Returns ------- out : nn.Tensor The output tensor, with shape of [m, n] and dtype of `out_dtype`. """ assert x.ndim >= 2 assert weight.ndim == 2 assert x_scale.ndim == x.ndim assert weight_scale.ndim == weight.ndim assert x.shape[-1] == weight.shape[1] assert x.shape[:-1] == x_scale.shape[:-1] assert (x.shape[-1] + block_size[1] - 1) // block_size[1] == x_scale.shape[-1] assert (weight.shape[1] + block_size[1] - 1) // block_size[1] == weight_scale.shape[1] assert (weight.shape[0] + block_size[0] - 1) // block_size[0] == weight_scale.shape[0] if x.dtype != "float8_e4m3fn" or weight.dtype != "float8_e4m3fn": raise ValueError( f"x and weight must be float8_e4m3fn, but got x={x.dtype}, weight={weight.dtype}" ) if x_scale.dtype != "float32" and weight_scale.dtype != "float32": raise ValueError( "x_scale and weight_scale must be float32, but got " f"x_scale={x_scale.dtype}, weight_scale={weight_scale.dtype}" ) if out_dtype not in ["float16", "bfloat16"]: raise ValueError(f"out_dtype must be float16 or bfloat16, but got {out_dtype}") M = x.shape[0] for i in range(1, x.ndim - 1): M *= x.shape[i] N = weight.shape[0] K = x.shape[-1] BLOCK_SIZE_M = 64 BLOCK_SIZE_N = block_size[0] BLOCK_SIZE_K = block_size[1] GROUP_SIZE_M = 32 num_warps = 4 num_stages = 3 x_shape = x.shape if x.ndim > 2: x = x.reshape(M, K) x_scale = x_scale.reshape(M, x_scale.shape[-1]) out = nn.extern( "mlc.triton.w8a8_block_fp8_matmul", args=[ x, weight, x_scale, weight_scale, N, K, block_size[0], block_size[1], BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M, num_warps, num_stages, str(x.dtype), str(out_dtype), ], out=nn.Tensor.placeholder((M, N), dtype=out_dtype), ) return out.reshape(*x_shape[:-1], N) if len(x_shape) > 2 else out def fp8_groupwise_scaled_group_gemm( x: nn.Tensor, x_scale: nn.Tensor, weight: nn.Tensor, weight_scale: nn.Tensor, indptr: nn.Tensor, block_size: Tuple[int, int], # noqa: UP006 out_dtype: str, ): """Triton block-scale fp8 group gemm operator. Parameters ---------- x : nn.Tensor The input tensor, with shape of [m, k]. x_scale : nn.Tensor The scale tensor, with shape of [m, k // block_size]. weight : nn.Tensor The weight tensor, with shape of [num_experts, n, k]. weight_scale : nn.Tensor The scale tensor, with shape of [num_experts, n // block_size, k // block_size]. indptr : nn.Tensor The indptr tensor of group gemm, with shape of [num_experts + 1,]. block_size : Tuple[int, int] The block size. out_dtype : str The data type of the output tensor. Returns ------- out : nn.Tensor The output tensor, with shape of [m, n] and dtype of `out_dtype`. """ assert x.ndim >= 2 assert weight.ndim == 3 assert x_scale.ndim == x.ndim assert weight_scale.ndim == weight.ndim assert x.shape[-1] == weight.shape[2] assert (x.shape[-1] + block_size[1] - 1) // block_size[1] == x_scale.shape[-1] assert (weight.shape[2] + block_size[1] - 1) // block_size[1] == weight_scale.shape[2] assert (weight.shape[1] + block_size[0] - 1) // block_size[0] == weight_scale.shape[1] num_experts = weight.shape[0] M = x.shape[0] for i in range(1, x.ndim - 1): M *= x.shape[i] N = weight.shape[1] K = x.shape[-1] assert weight_scale.shape[0] == num_experts assert indptr.ndim == 1 assert indptr.shape[0] == num_experts + 1 BLOCK_SIZE_M = 64 BLOCK_SIZE_N = block_size[0] BLOCK_SIZE_K = block_size[1] GROUP_SIZE_M = 32 num_warps = 4 num_stages = 3 x_shape = x.shape if x.ndim > 2: x = x.reshape(M, K) x_scale = x_scale.reshape(M, x_scale.shape[-1]) expert_ids = _compute_expert_id_per_block(indptr, num_experts, M, BLOCK_SIZE_M) out = nn.extern( "mlc.triton.w8a8_block_fp8_group_matmul", args=[ x, weight, x_scale, weight_scale, expert_ids, indptr, N, K, num_experts, block_size[0], block_size[1], BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M, num_warps, num_stages, str(x.dtype), str(out_dtype), ], out=nn.Tensor.placeholder((M, N), dtype=out_dtype), ) return out.reshape(*x_shape[:-1], N) if len(x_shape) > 2 else out