# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Fused MoE Triton kernels.""" import functools import json import os from typing import Any import torch import vllm.envs as envs import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm import _custom_ops as ops from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import ( MoEActivation, apply_moe_activation, ) from vllm.model_executor.layers.fused_moe.config import ( FUSED_MOE_UNQUANTIZED_CONFIG, FusedMoEQuantConfig, _get_config_dtype_str, ) from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( moe_align_block_size, ) from vllm.model_executor.layers.fused_moe.utils import ( moe_kernel_quantize_input, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import direct_register_custom_op logger = init_logger(__name__) @triton.jit def write_zeros_to_output( c_ptr, stride_cm, stride_cn, pid_n, N, offs_token, token_mask, BLOCK_SIZE_M, BLOCK_SIZE_N, compute_type, ): accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=compute_type) offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] c_mask = token_mask[:, None] & (offs_cn[None, :] < N) tl.store(c_ptrs, accumulator, mask=c_mask) @triton.jit def fused_moe_kernel_gptq_awq( # Pointers to matrices a_ptr, b_ptr, c_ptr, b_scale_ptr, b_zp_ptr, topk_weights_ptr, sorted_token_ids_ptr, expert_ids_ptr, num_tokens_post_padded_ptr, # Matrix dimensions N: tl.constexpr, K: tl.constexpr, EM, num_valid_tokens, # 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, stride_ak, stride_be, stride_bk, stride_bn, stride_cm, stride_cn, stride_bse, stride_bsk, stride_bsn, stride_bze, stride_bzk, stride_bzn, block_k_diviable: tl.constexpr, group_size: tl.constexpr, # Meta-parameters BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, SPLIT_K: tl.constexpr, MUL_ROUTED_WEIGHT: tl.constexpr, top_k: tl.constexpr, compute_type: tl.constexpr, has_zp: tl.constexpr, use_int4_w4a16: tl.constexpr, use_int8_w8a16: 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 (M, topk, N), where M is the total number of tokens post padding, topk is the number of times each token is repeated, and N is the output feature dimension. - sorted_token_ids: A tensor containing the sorted indices of tokens, repeated topk times and arranged by the expert index they are assigned to. - 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`. The sorting of `sorted_token_ids` by expert index and padding ensures divisibility by BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix multiplication across different blocks processed by the same expert. """ # ----------------------------------------------------------- # 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) num_pid_m = tl.cdiv(EM, 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 % 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 num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr) if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded: return offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) # Cast to int64 to prevent overflow in stride*offset products offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64) token_mask = offs_token < num_valid_tokens off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64) if off_experts == -1: # ----------------------------------------------------------- # Write back zeros to the output when the expert is not # in the current expert parallel rank. write_zeros_to_output( c_ptr, stride_cm, stride_cn, pid_n, N, offs_token, token_mask, BLOCK_SIZE_M, BLOCK_SIZE_N, compute_type, ) return offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N offs_k = tl.arange(0, BLOCK_SIZE_K) a_ptrs = a_ptr + ( offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak ) if use_int4_w4a16: b_ptrs = ( b_ptr + off_experts * stride_be + (offs_k[:, None] // 2) * stride_bk + offs_bn[None, :] * stride_bn ) b_shifter = (offs_k[:, None] % 2) * 4 elif use_int8_w8a16: b_ptrs = ( b_ptr + off_experts * stride_be + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn ) if not has_zp and use_int4_w4a16: b_zp_num = 8 if not has_zp and use_int8_w8a16: b_zp_num = 128 elif has_zp and use_int4_w4a16: b_zp_shifter = (offs_bn[None, :] % 2) * 4 # ----------------------------------------------------------- # 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 not block_k_diviable: k_mask = offs_k[:, None] < K - k * BLOCK_SIZE_K k_other = 0.0 else: k_mask = None k_other = None 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) if use_int4_w4a16: b = (b >> b_shifter) & 0xF b_scale_ptrs = ( b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn + ((offs_k[:, None] + BLOCK_SIZE_K * k) // group_size) * stride_bsk ) b_scale = tl.load(b_scale_ptrs, mask=k_mask, other=k_other) b_scale = b_scale.to(tl.float32) if has_zp and use_int4_w4a16: offs_k_true = (offs_k[:, None] + BLOCK_SIZE_K * k) // group_size b_zp_ptrs = ( b_zp_ptr + off_experts * stride_bze + (offs_bn[None, :] // 2) * stride_bzn + offs_k_true * stride_bzk ) b_zp = tl.load(b_zp_ptrs, mask=k_mask, other=k_other) b_zp = (b_zp >> b_zp_shifter) & 0xF b_zp = b_zp.to(tl.float32) elif has_zp and use_int8_w8a16: offs_k_true = (offs_k[:, None] + BLOCK_SIZE_K * k) // group_size b_zp_ptrs = ( b_zp_ptr + off_experts * stride_bze + offs_bn[None, :] * stride_bzn + offs_k_true * stride_bzk ) b_zp = tl.load(b_zp_ptrs, mask=k_mask, other=k_other) b_zp = b_zp.to(tl.float32) # We accumulate along the K dimension. if has_zp: b = ((b.to(tl.float32) - b_zp) * b_scale).to(compute_type) else: b = ((b.to(tl.float32) - b_zp_num) * b_scale).to(compute_type) accumulator = tl.dot(a, b, acc=accumulator) # Advance the ptrs to the next K block. a_ptrs += BLOCK_SIZE_K * stride_ak if use_int4_w4a16: b_ptrs += (BLOCK_SIZE_K // 2) * stride_bk else: b_ptrs += BLOCK_SIZE_K * stride_bk if MUL_ROUTED_WEIGHT: moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0) accumulator = accumulator * moe_weight[:, None] accumulator = accumulator.to(compute_type) # ----------------------------------------------------------- # 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[:, None] + stride_cn * offs_cn[None, :] c_mask = token_mask[:, None] & (offs_cn[None, :] < N) tl.store(c_ptrs, accumulator, mask=c_mask) @triton.jit def fused_moe_kernel( # Pointers to matrices a_ptr, b_ptr, c_ptr, b_bias_ptr, a_scale_ptr, b_scale_ptr, topk_weights_ptr, sorted_token_ids_ptr, expert_ids_ptr, num_tokens_post_padded_ptr, # Matrix dimensions N, K, EM, num_valid_tokens, # 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, stride_ak, stride_be, stride_bk, stride_bn, stride_cm, stride_cn, stride_asm, stride_ask, stride_bse, stride_bsk, stride_bsn, stride_bbe, # bias expert stride stride_bbn, # bias N stride # Block size for block-wise quantization group_n: tl.constexpr, group_k: tl.constexpr, naive_block_assignment: tl.constexpr, # Meta-parameters BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, SPLIT_K: tl.constexpr, MUL_ROUTED_WEIGHT: tl.constexpr, top_k: tl.constexpr, compute_type: tl.constexpr, use_fp8_w8a8: tl.constexpr, use_int8_w8a8: tl.constexpr, use_int8_w8a16: tl.constexpr, per_channel_quant: tl.constexpr, HAS_BIAS: 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 (M, topk, N), where M is the total number of tokens post padding, topk is the number of times each token is repeated, and N is the output feature dimension. - sorted_token_ids: A tensor containing the sorted indices of tokens, repeated topk times and arranged by the expert index they are assigned to. - 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. - naive_block_assignment: A boolean flag indicating whether to use naive token wise block assignment. If True, each block corresponds to a single token. This kernel performs the multiplication of a token by its corresponding expert matrix as determined by `expert_ids`. The sorting of `sorted_token_ids` by expert index and padding ensures divisibility by BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix multiplication across different blocks processed by the same expert. """ # ----------------------------------------------------------- # 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) num_pid_m = tl.cdiv(EM, 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 % 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 offs = tl.arange(0, BLOCK_SIZE_M).to(tl.int64) num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr) if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded: return if not naive_block_assignment: offs_token_id = pid_m * BLOCK_SIZE_M + offs offs_token = tl.load(sorted_token_ids_ptr + offs_token_id) else: offs_token = tl.where( offs == 0, pid_m, # first element = pid_m num_valid_tokens, # remaining elements = constant ) # Cast to int64 to prevent overflow in stride*offset products # (e.g. stride_cm * offs_token can exceed int32 for large token counts) offs_token = offs_token.to(tl.int64) token_mask = offs_token < num_valid_tokens off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64) if off_experts == -1: # ----------------------------------------------------------- # Write back zeros to the output when the expert is not # in the current expert parallel rank. write_zeros_to_output( c_ptr, stride_cm, stride_cn, pid_n, N, offs_token, token_mask, BLOCK_SIZE_M, BLOCK_SIZE_N, compute_type, ) return offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N offs_k = tl.arange(0, BLOCK_SIZE_K) a_ptrs = a_ptr + ( offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak ) b_ptrs = ( b_ptr + off_experts * stride_be + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) ) if use_int8_w8a16: b_scale_ptrs = ( b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn ) b_scale = tl.load(b_scale_ptrs) if use_fp8_w8a8 or use_int8_w8a8: # block-wise if group_k > 0 and group_n > 0: a_scale_ptrs = a_scale_ptr + (offs_token // top_k) * stride_asm offs_bsn = offs_bn // group_n b_scale_ptrs = ( b_scale_ptr + off_experts * stride_bse + offs_bsn * stride_bsn ) # channel-wise elif per_channel_quant: b_scale_ptrs = ( b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn ) b_scale = tl.load(b_scale_ptrs) # Load per-token scale for activations a_scale_ptrs = a_scale_ptr + (offs_token // top_k) * stride_asm a_scale = tl.load(a_scale_ptrs, mask=token_mask, other=0.0)[:, None] # tensor-wise else: a_scale = tl.load(a_scale_ptr) b_scale = tl.load(b_scale_ptr + off_experts) if HAS_BIAS: # bias shape: [num_experts, N] bias_ptrs = b_bias_ptr + off_experts * stride_bbe + offs_bn * stride_bbn bias = tl.load(bias_ptrs, mask=(offs_bn < N), other=0.0) # ----------------------------------------------------------- # 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. 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. if use_int8_w8a16: accumulator = tl.dot(a, b.to(compute_type), acc=accumulator) elif use_fp8_w8a8 or use_int8_w8a8: if group_k > 0 and group_n > 0: 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, :] else: if use_fp8_w8a8: # acc used to enable fp8_fast_accum accumulator = tl.dot(a, b, acc=accumulator) else: accumulator += tl.dot(a, b) else: accumulator += tl.dot(a, b) # Advance the ptrs to the next K block. a_ptrs += BLOCK_SIZE_K * stride_ak b_ptrs += BLOCK_SIZE_K * stride_bk # Dequantization for supported quantization schemes: # - int8_w8a16 # - fp8_w8a8 # - int8_w8a8 # Accumulator and scalings are in float32 to preserve numerical accuracy. if use_int8_w8a16: accumulator = accumulator * b_scale elif (use_fp8_w8a8 or use_int8_w8a8) and not (group_k > 0 and group_n > 0): accumulator = accumulator * a_scale * b_scale # Bias addition: # Bias must be applied after dequantization: # - Since bias is typically not quantized # - Bias should not be scaled by quantization factors if HAS_BIAS: accumulator += bias[None, :] # Router (MoE) weight multiplication: # This multiplication MUST be performed in float32 before any precision # conversion to ensure numerical stability, which is especially critical # on ROCm platforms. if MUL_ROUTED_WEIGHT: moe_weight = tl.load( topk_weights_ptr + offs_token, mask=token_mask, other=0, ) accumulator *= moe_weight[:, None] # Final precision conversion: # Cast once at the end to the desired compute/output dtype. accumulator = accumulator.to(compute_type) # ----------------------------------------------------------- # 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[:, None] + stride_cn * offs_cn[None, :] c_mask = token_mask[:, None] & (offs_cn[None, :] < N) tl.store(c_ptrs, accumulator, mask=c_mask) # NOTE(zyongye): we can remove all the wna16 kernel # once we drop off sm75 support def invoke_fused_moe_wna16_cuda_kernel( A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, B_scale: torch.Tensor | None, B_zp: torch.Tensor | None, topk_weights: torch.Tensor | None, sorted_token_ids: torch.Tensor | None, expert_ids: torch.Tensor, num_tokens_post_padded: torch.Tensor, mul_routed_weight: bool, top_k: int, config: dict[str, Any], block_shape: list[int], ): assert B_scale is not None and B_scale.ndim == 3 assert B_zp is None or B_zp.ndim == 3 assert block_shape is None or block_shape[0] == 0 M = A.size(0) num_tokens = M * top_k bit = 4 config = config.copy() config.update( get_moe_wna16_block_config( config=config, use_moe_wna16_cuda=True, num_valid_tokens=num_tokens, size_k=A.size(1), size_n=B.size(1), num_experts=B.size(1), group_size=block_shape[1], real_top_k=top_k, block_size_m=config["BLOCK_SIZE_M"], ) ) ops.moe_wna16_gemm( A, C, B, B_scale, B_zp, topk_weights if mul_routed_weight else None, sorted_token_ids, expert_ids, num_tokens_post_padded, top_k, config["BLOCK_SIZE_M"], config["BLOCK_SIZE_N"], config["BLOCK_SIZE_K"], bit, ) # NOTE(zyongye): we can remove all the wna16 kernel # once we drop off sm75 support def invoke_fused_moe_wna16_triton_kernel( A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, B_scale: torch.Tensor | None, B_zp: torch.Tensor | None, topk_weights: torch.Tensor | None, sorted_token_ids: torch.Tensor, expert_ids: torch.Tensor, num_tokens_post_padded: torch.Tensor, mul_routed_weight: bool, top_k: int, config: dict[str, Any], compute_type: tl.dtype, use_int8_w8a16: bool, use_int4_w4a16: bool, block_shape: list[int] | None, ): assert B_scale is not None and B_scale.ndim == 3 assert B_zp is None or B_zp.ndim == 3 assert block_shape is not None and block_shape[0] == 0 M = A.size(0) num_tokens = M * top_k EM = sorted_token_ids.size(0) if A.size(0) < config["BLOCK_SIZE_M"]: # optimize for small batch_size. # We assume that top_ids of each token is unique, # so num_valid_experts <= batch_size <= BLOCK_SIZE_M, # and we can skip some invalid blocks. EM = min(sorted_token_ids.size(0), A.size(0) * top_k * config["BLOCK_SIZE_M"]) grid = lambda META: ( triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(B.size(1), META["BLOCK_SIZE_N"]), ) config = config.copy() config.update( get_moe_wna16_block_config( config=config, use_moe_wna16_cuda=False, num_valid_tokens=num_tokens, size_k=A.size(1), size_n=B.size(1), num_experts=B.size(1), group_size=block_shape[1], real_top_k=top_k, block_size_m=config["BLOCK_SIZE_M"], ) ) fused_moe_kernel_gptq_awq[grid]( A, B, C, B_scale, B_zp, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded, B.size(1), A.size(1), EM, num_tokens, A.stride(0), A.stride(1), B.stride(0), B.stride(2), B.stride(1), C.stride(1), C.stride(2), B_scale.stride(0), B_scale.stride(2), B_scale.stride(1), B_zp.stride(0) if B_zp is not None else 0, B_zp.stride(2) if B_zp is not None else 0, B_zp.stride(1) if B_zp is not None else 0, block_k_diviable=A.size(1) % config["BLOCK_SIZE_K"] == 0, group_size=block_shape[1], MUL_ROUTED_WEIGHT=mul_routed_weight, top_k=top_k, compute_type=compute_type, has_zp=B_zp is not None, use_int4_w4a16=use_int4_w4a16, use_int8_w8a16=use_int8_w8a16, **config, ) def invoke_fused_moe_triton_kernel( A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, A_scale: torch.Tensor | None, B_scale: torch.Tensor | None, topk_weights: torch.Tensor | None, sorted_token_ids: torch.Tensor | None, expert_ids: torch.Tensor, num_tokens_post_padded: torch.Tensor, mul_routed_weight: bool, top_k: int, config: dict[str, Any], compute_type: tl.dtype, use_fp8_w8a8: bool, use_int8_w8a8: bool, use_int8_w8a16: bool, use_int4_w4a16: bool, per_channel_quant: bool, block_shape: list[int] | None = None, B_bias: torch.Tensor | None = None, ): assert topk_weights is not None or not mul_routed_weight assert topk_weights is None or topk_weights.stride(1) == 1 assert sorted_token_ids is None or sorted_token_ids.stride(0) == 1 if use_fp8_w8a8 or use_int8_w8a8: assert B_scale is not None assert block_shape is None or triton.cdiv( B.size(-2), block_shape[0] ) == B_scale.size(-2) assert block_shape is None or triton.cdiv( B.size(-1), block_shape[1] ) == B_scale.size(-1) elif use_int8_w8a16 or use_int4_w4a16: assert B_scale is not None assert block_shape is None or block_shape[0] == 0 else: assert A_scale is None assert B_scale is None M = A.size(0) num_tokens = M * top_k if sorted_token_ids is not None: EM = sorted_token_ids.size(0) if A.size(0) < config["BLOCK_SIZE_M"]: # optimize for small batch_size. # We assume that top_ids of each token is unique, # so num_valid_experts <= batch_size <= BLOCK_SIZE_M, # and we can skip some invalid blocks. EM = min( sorted_token_ids.size(0), A.size(0) * top_k * config["BLOCK_SIZE_M"] ) else: EM = num_tokens * config["BLOCK_SIZE_M"] grid = lambda META: ( triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(B.size(1), META["BLOCK_SIZE_N"]), ) HAS_BIAS = B_bias is not None config = config.copy() config["SPLIT_K"] = 1 BLOCK_SIZE_K = config.pop("BLOCK_SIZE_K") if block_shape is not None: BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0], block_shape[1])) fused_moe_kernel[grid]( A, B, C, B_bias, A_scale, B_scale, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded, B.size(1), B.size(2), EM, num_tokens, A.stride(0), A.stride(1), B.stride(0), B.stride(2), B.stride(1), C.stride(1), C.stride(2), A_scale.stride(0) if A_scale is not None and A_scale.ndim == 2 else 0, A_scale.stride(1) if A_scale is not None and A_scale.ndim == 2 else 0, B_scale.stride(0) if B_scale is not None and B_scale.ndim >= 2 else 0, B_scale.stride(2) if B_scale is not None and B_scale.ndim == 3 else 0, B_scale.stride(1) if B_scale is not None and B_scale.ndim >= 2 else 0, B_bias.stride(0) if B_bias is not None else 0, B_bias.stride(1) if B_bias is not None else 0, 0 if block_shape is None else block_shape[0], 0 if block_shape is None else block_shape[1], MUL_ROUTED_WEIGHT=mul_routed_weight, top_k=top_k, compute_type=compute_type, use_fp8_w8a8=use_fp8_w8a8, use_int8_w8a8=use_int8_w8a8, use_int8_w8a16=use_int8_w8a16, per_channel_quant=per_channel_quant, naive_block_assignment=(sorted_token_ids is None), HAS_BIAS=HAS_BIAS, BLOCK_SIZE_K=BLOCK_SIZE_K, **config, ) def dispatch_fused_moe_kernel( A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, A_scale: torch.Tensor | None, B_scale: torch.Tensor | None, B_zp: torch.Tensor | None, topk_weights: torch.Tensor | None, sorted_token_ids: torch.Tensor | None, expert_ids: torch.Tensor, num_tokens_post_padded: torch.Tensor, mul_routed_weight: bool, top_k: int, config: dict[str, Any], compute_type: tl.dtype, use_fp8_w8a8: bool, use_int8_w8a8: bool, use_int8_w8a16: bool, use_int4_w4a16: bool, per_channel_quant: bool, block_shape: list[int] | None = None, B_bias: torch.Tensor | None = None, ) -> None: assert topk_weights is not None or not mul_routed_weight assert topk_weights is None or topk_weights.stride(1) == 1 assert sorted_token_ids is None or sorted_token_ids.stride(0) == 1 M = A.size(0) num_tokens = M * top_k if (use_int8_w8a16 or use_int4_w4a16) and ( block_shape is not None and block_shape[1] > 0 ): assert B_bias is None use_moe_wna16_cuda = should_moe_wna16_use_cuda( num_valid_tokens=num_tokens, group_size=block_shape[1], num_experts=B.size(0), bit=4 if use_int4_w4a16 else 8, ) if use_moe_wna16_cuda: invoke_fused_moe_wna16_cuda_kernel( A, B, C, B_scale, B_zp, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded, mul_routed_weight, top_k, config, block_shape, ) return invoke_fused_moe_wna16_triton_kernel( A, B, C, B_scale, B_zp, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded, mul_routed_weight, top_k, config, compute_type, use_int8_w8a16, use_int4_w4a16, block_shape, ) else: invoke_fused_moe_triton_kernel( A, B, C, A_scale, B_scale, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded, mul_routed_weight, top_k, config, compute_type, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, per_channel_quant, block_shape, B_bias, ) @triton.jit def compute_identity_kernel( top_k: int, hidden_states_ptr: tl.tensor, expert_scales_ptr: tl.tensor, num_tokens: int, output_ptr: tl.tensor, hidden_dim: int, scales_stride: int, BLOCK_SIZE: tl.constexpr, ) -> None: pid = tl.program_id(0) batch_id = pid // (hidden_dim // BLOCK_SIZE) dim_offset = pid % (hidden_dim // BLOCK_SIZE) * BLOCK_SIZE if batch_id >= num_tokens or dim_offset >= hidden_dim: return h = tl.load( hidden_states_ptr + batch_id * hidden_dim + dim_offset + tl.arange(0, BLOCK_SIZE), mask=(dim_offset + tl.arange(0, BLOCK_SIZE)) < hidden_dim, ) result = tl.zeros([BLOCK_SIZE], dtype=tl.float32) for i in range(top_k): scale = tl.load(expert_scales_ptr + batch_id * scales_stride + i) result += h * scale tl.store( output_ptr + batch_id * hidden_dim + dim_offset + tl.arange(0, BLOCK_SIZE), result, mask=(dim_offset + tl.arange(0, BLOCK_SIZE)) < hidden_dim, ) def zero_experts_compute_triton( expert_indices: torch.Tensor, expert_scales: torch.Tensor, num_experts: int, zero_expert_type: str, hidden_states: torch.Tensor, ) -> torch.Tensor: N = expert_indices.numel() top_k = expert_indices.size(-1) grid = lambda meta: (triton.cdiv(N, meta["BLOCK_SIZE"]),) if zero_expert_type == "identity": zero_expert_mask = expert_indices < num_experts zero_expert_scales = expert_scales.clone() zero_expert_scales[zero_expert_mask] = 0.0 normal_expert_mask = expert_indices >= num_experts expert_indices[normal_expert_mask] = 0 expert_scales[normal_expert_mask] = 0.0 output = torch.zeros_like(hidden_states).to(hidden_states.device) hidden_dim = hidden_states.size(-1) num_tokens = hidden_states.size(0) grid = lambda meta: (num_tokens * (hidden_dim // meta["BLOCK_SIZE"]),) compute_identity_kernel[grid]( top_k, hidden_states, zero_expert_scales, num_tokens, output, hidden_dim, zero_expert_scales.stride(0), BLOCK_SIZE=256, ) return output # Adapted from: https://github.com/sgl-project/sglang/pull/2628 def get_config_file_name( E: int, N: int, dtype: str | None, block_shape: list[int] | None = None ) -> str: device_name = current_platform.get_device_name().replace(" ", "_") # Set device_name to H200 if a device from the H200 family is detected if "H200" in device_name.split("_"): device_name = "NVIDIA_H200" dtype_selector = "" if not dtype else f",dtype={dtype}" block_shape_selector = ( "" if not block_shape or not all(block_shape) else f",block_shape={block_shape}" ).replace(" ", "") return f"E={E},N={N},device_name={device_name}{dtype_selector}{block_shape_selector}.json" # noqa: E501 # Adapted from: https://github.com/sgl-project/sglang/pull/2628 @functools.lru_cache def get_moe_configs( E: int, N: int, dtype: str | None, block_n: int | None = None, block_k: int | None = None, ) -> dict[int, Any] | None: """ Return optimized configurations for the fused MoE kernel. The return value will be a dictionary that maps an irregular grid of batch sizes to configurations of the fused_moe kernel. To evaluate the kernel on a given batch size bs, the closest batch size in the grid should be picked and the associated configuration chosen to invoke the kernel. """ # Avoid optimizing for the batch invariant case. Use default config if envs.VLLM_BATCH_INVARIANT: return None # First look up if an optimized configuration is available in the configs # directory block_shape = [block_n, block_k] if block_n and block_k else None json_file_name = get_config_file_name(E, N, dtype, block_shape) config_file_paths = [] # note that we prioritize user defined config user_defined_config_folder = envs.VLLM_TUNED_CONFIG_FOLDER if user_defined_config_folder is not None: user_defined_config_file_path = os.path.join( user_defined_config_folder, json_file_name ) config_file_paths.append(user_defined_config_file_path) default_config_file_path = os.path.join( os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name ) config_file_paths.append(default_config_file_path) for config_file_path in config_file_paths: if os.path.exists(config_file_path): with open(config_file_path) as f: logger.info_once( "Using configuration from %s for MoE layer.", config_file_path, scope="global", ) # If a configuration has been found, return it tuned_config = json.load(f) # Delete triton_version from tuned_config tuned_config.pop("triton_version", None) return {int(key): val for key, val in tuned_config.items()} # If no optimized configuration is available, we will use the default # configuration logger.warning_once( "Using default MoE config. Performance might be sub-optimal! " "Config file not found at %s", ", ".join(config_file_paths), ) return None def _ensure_block_size_k_divisible( size_k: int, block_size_k: int, group_size: int ) -> int: """Ensure block_size_k is a divisor of size_k and divisible by group_size. This ensures BLOCK_SIZE_K compatibility with MoeWNA16 CUDA kernel which requires size_k % BLOCK_SIZE_K == 0 and BLOCK_SIZE_K % group_size == 0. Args: size_k: The size_k dimension that must be divisible by result. block_size_k: Preferred block size (will be adjusted if needed). group_size: The result must be divisible by this. Returns: A valid BLOCK_SIZE_K that divides size_k and is divisible by group_size. """ # Fast path: already valid if size_k % block_size_k == 0 and block_size_k % group_size == 0: return block_size_k # Find the largest value that: # 1. Divides size_k (size_k % candidate == 0) # 2. Is divisible by group_size (candidate % group_size == 0) # 3. Is <= block_size_k (prefer smaller values close to block_size_k) # # Strategy: Search from min(block_size_k, size_k) down to group_size, # stepping by group_size to ensure divisibility by group_size max_search = min(block_size_k, size_k) start = (max_search // group_size) * group_size for candidate in range(start, group_size - 1, -group_size): if size_k % candidate == 0: return candidate # Fallback: if group_size divides size_k, use it # This should always be true with correct group_size configuration if size_k % group_size == 0: return group_size # This should not happen with correct group_size, but ensure divisibility return size_k def get_moe_wna16_block_config( config: dict[str, int], use_moe_wna16_cuda: bool, num_valid_tokens: int, size_k: int, size_n: int, num_experts: int, group_size: int, real_top_k: int, block_size_m: int, ): if "BLOCK_SIZE_N" in config and "BLOCK_SIZE_K" in config: # optimal block config is set return {} if not use_moe_wna16_cuda: # triton moe wna16 kernel if num_valid_tokens // real_top_k == 1: # if bs=1, use a smaller BLOCK_SIZE_N return {"BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64} else: return {"BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32} else: # cuda moe wna16 kernel # set default block_size 128, and increase them when num_blocks # is too large. block_size_n = 128 block_size_k = 128 if block_size_k <= group_size: block_size_k = group_size num_n_blocks = size_k // block_size_k num_k_blocks = size_n // block_size_k num_m_blocks = ( num_valid_tokens + block_size_m - 1 ) / block_size_m + num_experts if num_valid_tokens // real_top_k <= block_size_m: num_m_blocks = min(num_m_blocks, num_valid_tokens) num_blocks = num_m_blocks * num_n_blocks * num_k_blocks if size_k % 256 == 0 and num_blocks >= 256 and block_size_k < 256: block_size_k = 256 num_blocks = num_blocks // (256 // block_size_k) if ( num_m_blocks <= 16 and size_k % (block_size_k * 2) == 0 and size_k % (block_size_k * 2) == 0 and block_size_k <= 512 and num_blocks >= 512 ): block_size_k = block_size_k * 2 num_blocks = num_blocks // 2 if num_blocks > 1024: block_size_n = 256 num_n_blocks = num_n_blocks // 2 num_blocks = num_blocks // 2 if size_n <= 1024 and num_blocks >= 1024: # The kernel performance got much better with BLOCK_SIZE_N=1024 # when num_blocks is large, event when N is small. # Not sure why, maybe it force the CUDA SM process only one block # at the same time. block_size_n = 1024 # Ensure BLOCK_SIZE_K is a divisor of size_k for CUDA kernel compatibility block_size_k = _ensure_block_size_k_divisible(size_k, block_size_k, group_size) return {"BLOCK_SIZE_N": block_size_n, "BLOCK_SIZE_K": block_size_k} def should_moe_wna16_use_cuda( num_valid_tokens: int, group_size: int, num_experts: int, bit: int ): return ( current_platform.is_cuda() and bit == 4 and group_size in [32, 64, 128] and num_valid_tokens / num_experts <= 6 ) def get_default_config( M: int, E: int, N: int, K: int, topk: int, dtype: str | None, block_shape: list[int] | None = None, ) -> dict[str, int]: if envs.VLLM_BATCH_INVARIANT: return { "BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8, "SPLIT_K": 1, } # num_stages can cause triton.runtime.errors.OutOfResources on ROCm. num_stages_rocm = 2 if dtype == "fp8_w8a8" and block_shape is not None: # Block-wise quant: tile sizes are constrained by block_shape. # Use a small M tile for decode-like batches where tokens are # spread thin across experts. Larger batches benefit from # GROUP_SIZE_M > 1 because the per-block scales add memory # traffic that benefits from L2 tile reuse. config = { "BLOCK_SIZE_M": 16 if M <= 64 else 64, "BLOCK_SIZE_N": block_shape[0], "BLOCK_SIZE_K": block_shape[1], "GROUP_SIZE_M": 1 if M <= 16 else 32, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3 if not current_platform.is_rocm() else num_stages_rocm, } elif dtype in ["int4_w4a16", "int8_w8a16"] and block_shape is not None: # moe wna16 kernels # only set BLOCK_SIZE_M # BLOCK_SIZE_N and BLOCK_SIZE_K would be set later bit = 4 if dtype == "int4_w4a16" else 8 use_moe_wna16_cuda = should_moe_wna16_use_cuda(M * topk, block_shape[1], E, bit) if use_moe_wna16_cuda: config = {"BLOCK_SIZE_M": min(16, M), "SPLIT_K": 1} elif M <= 20: config = {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 1, "SPLIT_K": 1} elif M <= 40: config = {"BLOCK_SIZE_M": 32, "GROUP_SIZE_M": 1, "SPLIT_K": 1} else: config = {"BLOCK_SIZE_M": 64, "GROUP_SIZE_M": 1, "SPLIT_K": 1} else: # General defaults for bf16/fp16 and fp8 per-tensor. # Tile sizes scale with batch: small batches are memory-bound # (favor tall-K tiles), large batches are compute-bound (favor # large M/N tiles with more warps). if M <= 32: block_m = 16 elif M <= 96: block_m = 32 elif M <= 512: block_m = 64 else: block_m = 128 block_n = 64 if M <= 64 else 128 # Small batches benefit from longer reduction (larger K tile), # while large batches prefer more output parallelism. # FP8 elements are half-width so larger K tiles are always cheap. block_k = 128 if dtype == "fp8_w8a8" or M <= 64 else 64 # Grouping adjacent M-blocks lets them share weight tiles in L2. # Only helps when there are enough M-blocks per expert to group; # with many experts each one sees few tokens so grouping is useless. tokens_per_expert = M // max(E, 1) group_m = 16 if tokens_per_expert > 128 else 1 # Large batches have enough blocks to saturate the GPU, so we # use more warps per block to increase arithmetic intensity. num_warps = 4 if M <= 128 else 8 if current_platform.is_rocm(): num_stages = num_stages_rocm elif M <= 32: num_stages = 4 else: num_stages = 3 config = { "BLOCK_SIZE_M": block_m, "BLOCK_SIZE_N": block_n, "BLOCK_SIZE_K": block_k, "GROUP_SIZE_M": group_m, "SPLIT_K": 1, "num_warps": num_warps, "num_stages": num_stages, } return config def try_get_optimal_moe_config( w1_shape: tuple[int, ...], w2_shape: tuple[int, ...], top_k: int, dtype: str | None, M: int, block_shape: list[int] | None = None, ) -> dict[str, int]: from vllm.model_executor.layers.fused_moe import get_config override_config = get_config() if override_config: config = override_config else: # First try to load optimal config from the file E, _, N = w2_shape if dtype == "int4_w4a16": N = N * 2 block_n = block_shape[0] if block_shape else 0 block_k = block_shape[1] if block_shape else 0 configs = get_moe_configs(E, N, dtype, block_n, block_k) if configs: # If an optimal configuration map has been found, look up the # optimal config config = configs[min(configs.keys(), key=lambda x: abs(x - M))] else: # Else use the default config config = get_default_config(M, E, N, w1_shape[2], top_k, dtype, block_shape) return config def fused_experts_op( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, activation: str = "silu", apply_router_weight_on_input: bool = False, use_fp8_w8a8: bool = False, use_int8_w8a8: bool = False, use_int8_w8a16: bool = False, use_int4_w4a16: bool = False, ocp_mx_scheme: str | None = None, per_channel_quant: bool = False, global_num_experts: int = -1, expert_map: torch.Tensor | None = None, w1_scale: torch.Tensor | None = None, w2_scale: torch.Tensor | None = None, w1_zp: torch.Tensor | None = None, w2_zp: torch.Tensor | None = None, a1_scale: torch.Tensor | None = None, a2_scale: torch.Tensor | None = None, block_shape: list[int] | None = None, w1_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, ) -> torch.Tensor: return fused_experts_impl( hidden_states, w1, w2, topk_weights, topk_ids, activation, apply_router_weight_on_input, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, ocp_mx_scheme, per_channel_quant, global_num_experts, expert_map, w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale, block_shape, w1_bias, w2_bias, ) def fused_experts_op_fake( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, activation: str = "silu", apply_router_weight_on_input: bool = False, use_fp8_w8a8: bool = False, use_int8_w8a8: bool = False, use_int8_w8a16: bool = False, use_int4_w4a16: bool = False, ocp_mx_scheme: str | None = None, per_channel_quant: bool = False, global_num_experts: int = -1, expert_map: torch.Tensor | None = None, w1_scale: torch.Tensor | None = None, w2_scale: torch.Tensor | None = None, w1_zp: torch.Tensor | None = None, w2_zp: torch.Tensor | None = None, a1_scale: torch.Tensor | None = None, a2_scale: torch.Tensor | None = None, block_shape: list[int] | None = None, w1_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, ) -> torch.Tensor: return torch.empty_like(hidden_states) direct_register_custom_op( op_name="fused_experts", op_func=fused_experts_op, fake_impl=fused_experts_op_fake, ) def _prepare_expert_assignment( topk_ids: torch.Tensor, config: dict[str, Any], num_tokens: int, top_k_num: int, global_num_experts: int, expert_map: torch.Tensor | None, *, use_int8_w8a16: bool = False, use_int4_w4a16: bool = False, block_shape: list[int] | None = None, ignore_invalid_experts: bool = False, ) -> tuple[torch.Tensor | None, torch.Tensor, torch.Tensor]: """Prepare expert assignments for the aligned and low-latency Triton paths.""" # SPARSITY_FACTOR is a heuristic margin ensuring tokens_in_chunk * top_k # activates only a small fraction of total experts # Skips moe_align_block_size and activates the `sorted_token_ids is None` # path of the fused_moe_kernel kernel naive_block_assignment = ( expert_map is None and num_tokens * top_k_num * 4 <= global_num_experts and not ( (use_int8_w8a16 or use_int4_w4a16) and block_shape is not None and block_shape[1] > 0 ) ) if naive_block_assignment: return ( None, topk_ids.view(-1), torch.full( (1,), topk_ids.numel() * config["BLOCK_SIZE_M"], dtype=torch.int32, device=topk_ids.device, ), ) return moe_align_block_size( topk_ids, config["BLOCK_SIZE_M"], global_num_experts, expert_map, ignore_invalid_experts=ignore_invalid_experts, ) def fused_experts( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, activation: MoEActivation = MoEActivation.SILU, apply_router_weight_on_input: bool = False, global_num_experts: int = -1, expert_map: torch.Tensor | None = None, quant_config: FusedMoEQuantConfig | None = None, ) -> torch.Tensor: """Run fused MoE expert computation using Triton kernels.""" if quant_config is None: quant_config = FUSED_MOE_UNQUANTIZED_CONFIG return torch.ops.vllm.fused_experts( hidden_states=hidden_states, w1=w1, w2=w2, topk_weights=topk_weights, topk_ids=topk_ids, activation=activation.value, apply_router_weight_on_input=apply_router_weight_on_input, use_fp8_w8a8=quant_config.use_fp8_w8a8, use_int8_w8a8=quant_config.use_int8_w8a8, use_int8_w8a16=quant_config.use_int8_w8a16, use_int4_w4a16=quant_config.use_int4_w4a16, ocp_mx_scheme=quant_config.ocp_mx_scheme, per_channel_quant=quant_config.per_act_token_quant, global_num_experts=global_num_experts, expert_map=expert_map, w1_scale=quant_config.w1_scale, w2_scale=quant_config.w2_scale, w1_zp=quant_config.w1_zp, w2_zp=quant_config.w2_zp, a1_scale=quant_config.a1_scale, a2_scale=quant_config.a2_scale, block_shape=quant_config.block_shape, w1_bias=quant_config.w1_bias, w2_bias=quant_config.w2_bias, ) def _get_config_quant_dtype( use_fp8_w8a8: bool, use_int8_w8a8: bool, ) -> None | torch.dtype | str: """ Get the quantization type based on the quantization strategy flags. We don't have a quant_config at this point so we need to work backwards. A return type of None means no quantization is required because the input is unquantized or has been quantized prior to calling fused_experts_impl. """ if use_fp8_w8a8: return current_platform.fp8_dtype() if use_int8_w8a8: return torch.int8 return None def fused_experts_impl( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, activation: str = "silu", apply_router_weight_on_input: bool = False, use_fp8_w8a8: bool = False, use_int8_w8a8: bool = False, use_int8_w8a16: bool = False, use_int4_w4a16: bool = False, ocp_mx_scheme: str | None = None, per_channel_quant: bool = False, global_num_experts: int = -1, expert_map: torch.Tensor | None = None, w1_scale: torch.Tensor | None = None, w2_scale: torch.Tensor | None = None, w1_zp: torch.Tensor | None = None, w2_zp: torch.Tensor | None = None, a1_scale: torch.Tensor | None = None, a2_scale: torch.Tensor | None = None, block_shape: list[int] | None = None, w1_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, ) -> torch.Tensor: if ocp_mx_scheme is not None: raise NotImplementedError( f"Using ocp_mx_scheme={ocp_mx_scheme} in functional fused_experts call is " "deprecated. Please use OCP_MXQuantizationEmulationTritonExperts." ) # Convert string activation to enum for internal use activation_enum = MoEActivation.from_str(activation) # Check constraints. if use_int4_w4a16: assert hidden_states.size(1) // 2 == w1.size(2), "Hidden size mismatch" else: assert hidden_states.size(1) == w1.size(2), ( f"Hidden size mismatch {hidden_states.size(1)} != {w1.size(2)}" ) assert topk_weights.size() == topk_ids.size(), "topk shape mismatch" assert hidden_states.is_contiguous(), "Hidden_states must be contiguous" assert w1.stride(-1) == 1, "Stride of last dimension must be 1" assert w2.stride(-1) == 1, "Stride of last dimension must be 1" assert hidden_states.dtype in [torch.float32, torch.float16, torch.bfloat16] num_tokens = hidden_states.size(0) E, N, _ = w1.size() K = w2.size(1) if global_num_experts == -1: global_num_experts = E top_k_num = topk_ids.size(1) M = num_tokens config_dtype = _get_config_dtype_str( use_fp8_w8a8=use_fp8_w8a8, use_int8_w8a16=use_int8_w8a16, use_int4_w4a16=use_int4_w4a16, dtype=hidden_states.dtype, ) # Note: for use_int8_w8a16 or use_int4_w4a16, the activations are # quantized prior to calling fused_experts. quant_dtype = _get_config_quant_dtype( use_fp8_w8a8=use_fp8_w8a8, use_int8_w8a8=use_int8_w8a8, ) get_config_func = functools.partial( try_get_optimal_moe_config, w1.size(), w2.size(), top_k_num, config_dtype, block_shape=block_shape, ) config = get_config_func(M) # We can reuse the memory between these because by the time we need # cache3, we're done with cache1 cache13 = torch.empty( M * top_k_num * max(N, K), device=hidden_states.device, dtype=hidden_states.dtype, ) intermediate_cache1 = cache13[: M * top_k_num * N].view(M, top_k_num, N) intermediate_cache3 = cache13[: M * top_k_num * K].view(M, top_k_num, K) # This needs separate memory since it's used concurrently with cache1 activation_out_dim = mk.FusedMoEExpertsModular.adjust_N_for_activation( N, activation_enum ) intermediate_cache2 = torch.empty( (M * top_k_num, activation_out_dim), device=hidden_states.device, dtype=hidden_states.dtype, ) if hidden_states.dtype == torch.bfloat16: compute_type = tl.bfloat16 elif hidden_states.dtype == torch.float16: compute_type = tl.float16 elif hidden_states.dtype == torch.float32: compute_type = tl.float32 else: raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}") out_hidden_states = torch.empty_like(hidden_states) qhidden_states, a1q_scale = moe_kernel_quantize_input( A=hidden_states, A_scale=a1_scale, quant_dtype=quant_dtype, per_act_token_quant=per_channel_quant, block_shape=block_shape, ) sorted_token_ids, expert_ids, num_tokens_post_padded = _prepare_expert_assignment( topk_ids, config, num_tokens, top_k_num, global_num_experts, expert_map, use_int8_w8a16=use_int8_w8a16, use_int4_w4a16=use_int4_w4a16, block_shape=block_shape, ignore_invalid_experts=True, ) dispatch_fused_moe_kernel( qhidden_states, w1, intermediate_cache1, a1q_scale, w1_scale, w1_zp, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded, apply_router_weight_on_input, top_k_num, config, compute_type=compute_type, use_fp8_w8a8=use_fp8_w8a8, use_int8_w8a8=use_int8_w8a8, use_int8_w8a16=use_int8_w8a16, use_int4_w4a16=use_int4_w4a16, per_channel_quant=per_channel_quant, block_shape=block_shape, B_bias=w1_bias, ) apply_moe_activation( activation_enum, intermediate_cache2, intermediate_cache1.view(-1, N) ) qintermediate_cache2, a2q_scale = moe_kernel_quantize_input( A=intermediate_cache2, A_scale=a2_scale, quant_dtype=quant_dtype, per_act_token_quant=per_channel_quant, block_shape=block_shape, ) if expert_map is not None: intermediate_cache3.zero_() dispatch_fused_moe_kernel( qintermediate_cache2, w2, intermediate_cache3, a2q_scale, w2_scale, w2_zp, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded, not apply_router_weight_on_input, 1, config, compute_type=compute_type, use_fp8_w8a8=use_fp8_w8a8, use_int8_w8a8=use_int8_w8a8, use_int8_w8a16=use_int8_w8a16, use_int4_w4a16=use_int4_w4a16, per_channel_quant=per_channel_quant, block_shape=block_shape, B_bias=w2_bias, ) ops.moe_sum( intermediate_cache3.view(*intermediate_cache3.size()), out_hidden_states, ) return out_hidden_states