diff --git a/ex_engine/moe/__init__.py b/ex_engine/moe/__init__.py index 66c220e8..8f434272 100644 --- a/ex_engine/moe/__init__.py +++ b/ex_engine/moe/__init__.py @@ -1,10 +1,160 @@ -""" -ex_engine.moe — MoE expert computation for BI-V100 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -Ported from: - upstream_ref/ds_vllm/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py - upstream_ref/ds_vllm/vllm/model_executor/layers/fused_moe/activation.py -""" +from contextlib import contextmanager +from typing import Any -from ex_engine.moe.naive_batched_experts import naive_batched_moe_forward -from ex_engine.moe.activation import MoEActivation, apply_moe_activation +from vllm.model_executor.layers.fused_moe.activation import ( + MoEActivation, + activation_without_mul, + apply_moe_activation, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.layer import ( + FusedMoE, + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.fused_moe.modular_kernel import ( + FusedMoEActivationFormat, + FusedMoEExpertsModular, + FusedMoEPrepareAndFinalizeModular, +) +from vllm.model_executor.layers.fused_moe.routed_experts import ( + FusedMoeWeightScaleSupported, + RoutedExperts, +) +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( + FusedMoERouter, +) +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear +from vllm.model_executor.layers.fused_moe.runner.moe_runner import ( + MoERunner, +) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) +from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( + UnquantizedFusedMoEMethod, +) +from vllm.triton_utils import HAS_TRITON + +_config: dict[str, Any] | None = None + + +@contextmanager +def override_config(config): + global _config + old_config = _config + _config = config + yield + _config = old_config + + +def get_config() -> dict[str, Any] | None: + return _config + + +__all__ = [ + "FusedMoE", + "FusedMoERouter", + "FusedMoEConfig", + "FusedMoEQuantConfig", + "FusedMoEParallelConfig", + "FusedMoEMethodBase", + "MoEActivation", + "UnquantizedFusedMoEMethod", + "FusedMoeWeightScaleSupported", + "FusedMoEExpertsModular", + "FusedMoEActivationFormat", + "FusedMoEPrepareAndFinalizeModular", + "GateLinear", + "MoERunner", + "RoutingMethodType", + "RoutedExperts", + "SharedExperts", + "activation_without_mul", + "apply_moe_activation", + "fused_moe_make_expert_params_mapping", + "override_config", + "get_config", +] + +if HAS_TRITON: + # import to register the custom ops + from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import ( + BatchedDeepGemmExperts, + ) + from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( + CutlassBatchedExpertsFp8, + CutlassExpertsFp8, + CutlassExpertsW4A8Fp8, + ) + from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import ( + DeepGemmExperts, + ) + from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import ( + BatchedTritonExperts, + ) + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( + AiterExperts, + ) + from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( + TritonOrDeepGemmExperts, + ) + from vllm.model_executor.layers.fused_moe.experts.triton_moe import ( + TritonExperts, + TritonWNA16Experts, + ) + from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( + XPUExperts, + XPUExpertsFp8, + XPUExpertsMxFp4, + ) + from vllm.model_executor.layers.fused_moe.fused_moe import ( + fused_experts, + get_config_file_name, + ) + from vllm.model_executor.layers.fused_moe.router.fused_topk_router import ( + fused_topk, + ) + from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( + GroupedTopk, + ) + + __all__ += [ + "AiterExperts", + "fused_topk", + "fused_experts", + "get_config_file_name", + "GroupedTopk", + "CutlassExpertsFp8", + "CutlassBatchedExpertsFp8", + "CutlassExpertsW4A8Fp8", + "TritonExperts", + "TritonWNA16Experts", + "BatchedTritonExperts", + "DeepGemmExperts", + "BatchedDeepGemmExperts", + "TritonOrDeepGemmExperts", + "XPUExperts", + "XPUExpertsFp8", + "XPUExpertsBlockFp8", + "XPUExpertsMxFp8", + "XPUExpertsMxFp4", + ] +else: + # Some model classes directly use the custom ops. Add placeholders + # to avoid import errors. + def _raise_exception(method: str): + raise NotImplementedError(f"{method} is not implemented as lack of triton.") + + fused_topk = lambda *args, **kwargs: _raise_exception("fused_topk") + fused_experts = lambda *args, **kwargs: _raise_exception("fused_experts") diff --git a/ex_engine/moe/activation.py b/ex_engine/moe/activation.py index 7216b5a0..b2e67e62 100644 --- a/ex_engine/moe/activation.py +++ b/ex_engine/moe/activation.py @@ -122,32 +122,17 @@ def apply_moe_activation( # Activations with gated multiplication (gate × activation(up)) if activation == MoEActivation.SILU: - # BI-V100: torch.ops._C.silu_and_mul not available - # Use corex_attn_head_rms_norm pattern: try C++ first, fallback to PyTorch - d = output.size(-1) - gate = input[..., :d] - up = input[..., d:] - output.copy_(F.silu(gate) * up) + torch.ops._C.silu_and_mul(output, input) elif activation == MoEActivation.GELU: - d = output.size(-1) - gate = input[..., :d] - up = input[..., d:] - output.copy_(F.gelu(gate) * up) + torch.ops._C.gelu_and_mul(output, input) elif activation == MoEActivation.GELU_TANH: - d = output.size(-1) - gate = input[..., :d] - up = input[..., d:] - output.copy_(F.gelu(gate, approximate="tanh") * up) + torch.ops._C.gelu_tanh_and_mul(output, input) elif activation == MoEActivation.SWIGLUOAI: - d = output.size(-1) - gate = input[..., :d] - up = input[..., d:] - output.copy_(F.silu(gate) * up) + torch.ops._C.swigluoai_and_mul(output, input) elif activation == MoEActivation.SWIGLUSTEP: - d = output.size(-1) - gate = input[..., :d] - up = input[..., d:] - output.copy_(F.silu(gate) * up) + from vllm.model_executor.layers.activation import swiglustep_and_mul_triton + + swiglustep_and_mul_triton(output, input) # Activations without gated multiplication elif activation == MoEActivation.SILU_NO_MUL: diff --git a/ex_engine/moe/config.py b/ex_engine/moe/config.py new file mode 100644 index 00000000..1b063559 --- /dev/null +++ b/ex_engine/moe/config.py @@ -0,0 +1,1407 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass +from enum import IntEnum +from typing import Union + +import torch + +from vllm.config import ParallelConfig, SchedulerConfig +from vllm.config.kernel import MoEBackend +from vllm.distributed import get_dp_group, get_pcp_group, get_tensor_model_parallel_rank +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( + OCP_MX_DTYPES, + OCP_MX_Scheme, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_triton_kernels +from vllm.utils.math_utils import cdiv + +logger = init_logger(__name__) + +if has_triton_kernels(): + try: + from triton_kernels.matmul_ogs import PrecisionConfig + except (ImportError, AttributeError) as e: + logger.error( + "Failed to import Triton kernels. Please make sure your triton " + "version is compatible. Error: %s", + e, + ) + + +def _get_config_dtype_str( + dtype: torch.dtype, + use_fp8_w8a8: bool = False, + use_fp8_w8a16: bool = False, + use_int8_w8a16: bool = False, + use_int4_w4a16: bool = False, + ocp_mx_scheme: str | None = None, +) -> str | None: + """ + Return a string used to construct the filename that contains the + tuning info for a particular quantization scheme. See + try_get_optimal_moe_config in fused_moe.py. + """ + if use_fp8_w8a8: + return "fp8_w8a8" + elif use_fp8_w8a16: + return "fp8_w8a16" + elif use_int8_w8a16: + return "int8_w8a16" + elif use_int4_w4a16: + return "int4_w4a16" + elif ocp_mx_scheme is not None: + # The output of this function is passed to `try_get_optimal_moe_config`, + # and as we only simulate OCP MX execution in fused_moe for now, + # we will NOT look for `*,dtype=w_mxfp4_a_mxfp4.json` for now. + return None + elif dtype == torch.float: + # avoiding cases where kernel fails when float32 MoE + # use fp16/bfloat16 configs + return "float32" + return None + + +def _quant_flags_to_group_shape( + quant_dtype: torch.dtype | str | None, + per_act_token_quant: bool, + per_out_ch_quant: bool, + block_shape: list[int] | None, +) -> tuple[GroupShape | None, GroupShape | None]: + """ + Convert MoE quantization flags into more generic GroupShapes. + """ + a_shape: GroupShape | None + w_shape: GroupShape | None + if block_shape is not None: + assert not per_act_token_quant + assert not per_out_ch_quant + # TODO(bnell): this is not quite right for activations since first + # dim should be 1. + a_shape = GroupShape(row=block_shape[0], col=block_shape[1]) + w_shape = GroupShape(row=block_shape[0], col=block_shape[1]) + else: + w_shape = None + a_shape = None if quant_dtype is None else GroupShape.PER_TENSOR + + if per_act_token_quant: + a_shape = GroupShape.PER_TOKEN + + if per_out_ch_quant: + w_shape = GroupShape.PER_TOKEN + + return a_shape, w_shape + + +# The type of method in top-K routing +# Please keep this in sync with the counterpart defined in https://github.com/flashinfer-ai/flashinfer/blob/main/include/flashinfer/trtllm/fused_moe/runner.h +class RoutingMethodType(IntEnum): + # Default: Softmax -> TopK + Default = (0,) + # Renormalize: TopK -> Softmax + Renormalize = (1,) + # DeepSeekV3: Sigmoid -> RoutingBiasAdd -> Top2 in group -> Top4 groups + # -> Top8 experts from the Top4 groups + DeepSeekV3 = (2,) + # Llama4: Top1 -> Sigmoid + Llama4 = (3,) + # RenormalizeNaive: Softmax -> TopK -> Renormalize + RenormalizeNaive = (4,) + # TopK: TopK (no softmax) + TopK = (5,) + # SigmoidRenorm: Sigmoid -> TopK -> Renormalize (divide by sum of top-K) + SigmoidRenorm = (6,) + # MiniMax2: Sigmoid + Bias -> TopK -> ScaledSumNormalize + # (routeScale=1.0, epsilon=1e-20) + MiniMax2 = (7,) + # Sigmoid: Sigmoid -> TopK (no renormalization) + Sigmoid = (8,) + # Unspecified + Unspecified = (9,) + # other routing types (not passed to FlashInfer kernels) + # Deepseek V4 -> sqrtsoftplus + Bias + Normalize + DeepseekV4 = (100,) + Custom = (101,) + Simulated = (102,) + + +def get_routing_method_type( + scoring_func: str, + top_k: int, + renormalize: bool, + num_expert_group: int | None, + has_e_score_bias: bool, + routed_scaling_factor: float | None = 1.0, +) -> RoutingMethodType: + if scoring_func == "sqrtsoftplus": + # DeepSeek V4 uses sqrtsoftplus routing with optional routing bias + # and top-k renormalization. + if renormalize: + return RoutingMethodType.DeepseekV4 + else: + return RoutingMethodType.Unspecified + + if has_e_score_bias: + if scoring_func == "sigmoid": + if not renormalize: + return RoutingMethodType.Unspecified + if (num_expert_group or 0) > 0: + return RoutingMethodType.DeepSeekV3 + if routed_scaling_factor in (None, 1.0): + return RoutingMethodType.MiniMax2 + return RoutingMethodType.Unspecified + else: + return RoutingMethodType.Unspecified + + if scoring_func == "sigmoid": + if renormalize: + return RoutingMethodType.SigmoidRenorm + return RoutingMethodType.Sigmoid + + if scoring_func == "softmax": + if renormalize: + return RoutingMethodType.RenormalizeNaive + else: + return RoutingMethodType.Default + + return RoutingMethodType.Unspecified + + +@dataclass +class FusedMoEQuantDesc: + """ + A quantization descriptor for fused MoE ops. This class can describe + either activations or weights. + """ + + # The quantized type of this parameters. None means unquantized or + # already quantized. + # TODO (bnell): use scalar_type instead of Union. + dtype: torch.dtype | str | None = None + + # A field that describes the quantization group shape, from quant_utils.py. + # * (-1, -1) for per-tensor quantization + # * (1, -1) for per-row quantization + # * (-1, 1) for per-column quantization + # * (128, 128) for 128x128 deepseek style block quantization + # * (1, 128) for deepseek style activation quantization + # (i.e. per-token-per-group) + shape: GroupShape | None = None + + # Quantization scales. + # TODO(bnell): maybe put PrecisionConfigs in subclass of QuantDesc? + scale: Union[torch.Tensor, "PrecisionConfig", None] = None + + # Quantization alphas or gscales, used for nvfp4 types. + # W4A8 FP8: used for per-channel scales + # TODO(bnell): put some of these in subclasses + alpha_or_gscale: torch.Tensor | None = None + + # Zero points for int4/int8 types + zp: torch.Tensor | None = None + + # Biases for GPT triton MoE + bias: torch.Tensor | None = None + + +# TODO(bnell): have subclasses for specific moe methods? +# e.g. for specific arguments bias, precision, etc. +@dataclass +class FusedMoEQuantConfig: + """ + The FusedMoEQuantConfig contains all the quantization parameters for + a single FusedMoEMethodBase operation. It consists of four + FusedMoEQuantDescs, one for each activation and set of weights. + + Each FusedMoEMethodBase must implement a get_fused_moe_quant_config + method to construct a FusedMoEQuantConfig for use with that class. + + FusedMoEQuant configs are only used for modular kernels, fused_experts + (from fused_moe.py), cutlass_moe_fp[48], rocm_aiter_fused_experts and + triton_kernel_moe_forward. Other MoE methods can ignore the + FusedMoEQuantConfig (for now) and hardcode it to None. + + There are currently some restrictions on what can be expressed: + - Most MoE ops only support similar quantization strategies for + each parameter, e.g. both weights must have the same GroupShape + and both activations must share the same GroupShape. One exception to + this is the cutlass moe which allows per channel quantization on the + outputs. Note: this restrictions are not always rigorously checked. + - Not all fused MoE functions support all the parameters, e.g. zero points, + global scales, alphas and biases are not universally supported. + - Fully general GroupShapes are not allowed. Activations only support + per token, per tensor or K-blocked. + - Weights are not required to have a GroupShape since they have already + been quantized. + + Other notes: + - PrecisionConfigs are specific to GPT OSS Triton. + - As a follow up it would probably make sense to subclass FusedMoEQuantDesc + or FusedMoEQuantConfig for particular FusedMoEMethodBase subclasses + so that only the required quantization parameters are used/stored. + """ + + # TODO(bnell) make sure a1_scales/a2_scales don't interfere with chunking + _a1: FusedMoEQuantDesc + _a2: FusedMoEQuantDesc + _w1: FusedMoEQuantDesc + _w2: FusedMoEQuantDesc + is_scale_swizzled: bool = True + + # MXFP4-specific TRTLLM parameters for SwiGLU activation clamping. + # These correspond to gemm1_alpha, gemm1_beta, gemm1_clamp_limit + # in TrtLlmMxfp4ExpertsBase. + gemm1_alpha: float | None = None + gemm1_beta: float | None = None + gemm1_clamp_limit: float | None = None + + mx_alignment: int = 0 + + def __post_init__(self): + assert not self.per_act_token_quant or self.block_shape is None, ( + "illegal quantization" + ) + + # + # Convenience accessors for various properties. + # + + @property + def quant_dtype(self) -> torch.dtype | str | None: + return self._a1.dtype + + @property + def weight_quant_dtype(self) -> torch.dtype | str | None: + return self._w1.dtype + + @property + def is_quantized(self) -> bool: + return self.quant_dtype is not None + + @property + def is_per_act_token(self) -> bool: + return self._a1.shape == GroupShape.PER_TOKEN + + @property + def per_act_token_quant(self) -> bool: + return self._a1.shape == GroupShape.PER_TOKEN + + @property + def per_out_ch_quant(self) -> bool: + return self._w1.shape == GroupShape.PER_TOKEN + + @property + def is_per_tensor(self) -> bool: + return self._a1.shape == GroupShape.PER_TENSOR + + @property + def block_shape(self) -> list[int] | None: + if ( + self._a1.shape is not None + and self._a1.shape != GroupShape.PER_TENSOR + and self._a1.shape != GroupShape.PER_TOKEN + ): + return [self._a1.shape.row, self._a1.shape.col] + else: + return None + + @property + def is_block_quantized(self) -> bool: + return self.block_shape is not None + + @property + def a1_scale(self) -> torch.Tensor | None: + assert self._a1.scale is None or isinstance(self._a1.scale, torch.Tensor) + return self._a1.scale + + @property + def a1_gscale(self) -> torch.Tensor | None: + return self._a1.alpha_or_gscale + + @property + def a2_scale(self) -> torch.Tensor | None: + assert self._a2.scale is None or isinstance(self._a2.scale, torch.Tensor) + return self._a2.scale + + @property + def a2_gscale(self) -> torch.Tensor | None: + return self._a2.alpha_or_gscale + + @property + def w1_scale(self) -> torch.Tensor | None: + assert self._w1.scale is None or isinstance(self._w1.scale, torch.Tensor) + return self._w1.scale + + @property + def w1_zp(self) -> torch.Tensor | None: + return self._w1.zp + + @property + def w1_bias(self) -> torch.Tensor | None: + return self._w1.bias + + @property + def w1_precision(self) -> "PrecisionConfig | None": + assert self._w1.scale is None or isinstance(self._w1.scale, PrecisionConfig) + return self._w1.scale + + @property + def g1_alphas(self) -> torch.Tensor | None: + return self._w1.alpha_or_gscale + + @property + def w2_scale(self) -> torch.Tensor | None: + assert self._w2.scale is None or isinstance(self._w2.scale, torch.Tensor) + return self._w2.scale + + @property + def w2_zp(self) -> torch.Tensor | None: + return self._w2.zp + + @property + def w2_bias(self) -> torch.Tensor | None: + return self._w2.bias + + @property + def w2_precision(self) -> "PrecisionConfig | None": + assert self._w2.scale is None or isinstance(self._w2.scale, PrecisionConfig) + return self._w2.scale + + @property + def g2_alphas(self) -> torch.Tensor | None: + return self._w2.alpha_or_gscale + + @property + def use_fp8_w8a8(self) -> bool: + return self.quant_dtype == current_platform.fp8_dtype() + + @property + def use_int8_w8a8(self) -> bool: + return self.quant_dtype == torch.int8 + + @property + def use_int8_w8a16(self) -> bool: + return self._a1.dtype is None and self._w1.dtype == torch.int8 + + @property + def use_fp8_w8a16(self) -> bool: + return self._a1.dtype is None and self._w1.dtype == current_platform.fp8_dtype() + + @property + def use_int4_w4a16(self) -> bool: + return self._a1.dtype is None and self._w1.dtype == "int4" + + @property + def use_nvfp4_w4a16(self) -> bool: + return self._a1.dtype is None and self._w1.dtype == "nvfp4" + + @property + def ocp_mx_scheme(self) -> str | None: + if not hasattr(self, "_ocp_mx_scheme"): + if (self._a1.dtype is not None and not isinstance(self._a1.dtype, str)) or ( + self._w1.dtype is not None and not isinstance(self._w1.dtype, str) + ): + self._ocp_mx_scheme = None + else: + ocp_mx_scheme = OCP_MX_Scheme.from_quant_dtype( + self._a1.dtype, self._w1.dtype + ) + + if ocp_mx_scheme is not None: + ocp_mx_scheme = ocp_mx_scheme.value + + self._ocp_mx_scheme = ocp_mx_scheme + + return self._ocp_mx_scheme + + @property + def use_mxfp4_w4a16(self) -> bool: + return self._a1.dtype is None and self._w1.dtype == "mxfp4" + + @property + def use_mxfp4_w4a4(self) -> bool: + return self._a1.dtype == "mxfp4" and self._w1.dtype == "mxfp4" + + @property + def use_nvfp4_w4a4(self) -> bool: + return self.quant_dtype == "nvfp4" + + @property + def use_mxfp4_w4a8(self) -> bool: + return self._a1.dtype == "fp8" and self._w1.dtype == "mxfp4" + + def config_name(self, dtype: torch.dtype) -> str | None: + """ + Return a string used to construct the filename that contains the + tuning info for a particular quantization scheme. See + try_get_optimal_moe_config in fused_moe.py. + """ + return _get_config_dtype_str( + use_fp8_w8a8=self.use_fp8_w8a8, + use_fp8_w8a16=self.use_fp8_w8a16, + use_int8_w8a16=self.use_int8_w8a16, + use_int4_w4a16=self.use_int4_w4a16, + ocp_mx_scheme=self.ocp_mx_scheme, + dtype=dtype, + ) + + def scale_shape( + self, + max_tokens: int, + hidden_dim: int, + ) -> tuple[int, int] | None: + """ + Construct the proper activation scale shape for this + config. + """ + if self.is_quantized: + if self.is_block_quantized: + assert self.block_shape is not None + _, block_k = self.block_shape + k_tiles = cdiv(hidden_dim, block_k) + return (max_tokens, k_tiles) + elif self.is_per_act_token: + return (max_tokens, 1) + else: + return (1, 1) + else: + return None + + def batched_scale_shape( + self, + num_experts: int, + max_tokens: int, + hidden_dim: int, + ) -> tuple[int, int, int] | None: + """ + Construct the proper activation batched scale shape for this + config, e.g. (num experts, *scale_shape). + """ + if self.is_quantized: + scale_shape = self.scale_shape(max_tokens, hidden_dim) + assert scale_shape is not None + return (num_experts, *scale_shape) + else: + return None + + @staticmethod + def make( + quant_dtype: torch.dtype | str | None = None, + per_act_token_quant: bool = False, + per_out_ch_quant: bool = False, + block_shape: list[int] | None = None, + w1_scale: Union[torch.Tensor, "PrecisionConfig", None] = None, + w2_scale: Union[torch.Tensor, "PrecisionConfig", None] = None, + a1_scale: torch.Tensor | None = None, + a2_scale: torch.Tensor | None = None, + g1_alphas: torch.Tensor | None = None, + g2_alphas: torch.Tensor | None = None, + a1_gscale: torch.Tensor | None = None, + a2_gscale: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + w1_zp: torch.Tensor | None = None, + w2_zp: torch.Tensor | None = None, + weight_dtype: torch.dtype | str | None = None, + is_scale_swizzled: bool = True, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, + gemm1_clamp_limit: float | None = None, + ) -> "FusedMoEQuantConfig": + """ + General builder function for a FusedMoEQuantConfig. + - quant_dtype: Optional quantization type. None if activations are + unquantized or quantized prior to calling. Note: "nvfp4", "mxfp4", + "mxfp6_e3m2", "mxfp6_e2m3" are the only valid string values + for quant_dtype. + - per_act_token_quant: Activations have per token quantization. + - per_out_ch_quant: Outputs have per channel quantization. (only + for cutlass). + - block_shape: Optional block size for block-wise quantization. + Incompatible with per_act_token and per_out_ch quant. + - w1_scale: Optional scale to be used for w1. + - w2_scale: Optional scale to be used for w2. + - a1_scale: Optional scale to be used for a1. + - a2_scale: Optional scale to be used for a2. + - g1_alphas: Optional global quantization scales for w1 (for nvfp4). + Optional per-channel scales for w1 (for W4A8 FP8). + Optional dq scale i.e. w_scale * a_scale (for W8A8 fp8). + - g2_alphas: Optional global quantization scales for w2 (for nvfp4). + Optional per-channel scales for w2 (for W4A8 FP8). + Optional dq scale i.e. w_scale * a_scale (for W8A8 fp8). + - a1_gscale: Optional global quantization scales for a1 (1.0 /a2_scale). + - a2_gscale: Optional global quantization scales for a2 (1.0 /a2_scale). + + - w1_bias: Optional biases for w1 (GPT OSS Triton). + - w2_bias: Optional biases for w1 (GPT OSS Triton). + - w1_zp: Optional w1 zero points for int4/int8 quantization. + - w2_zp: Optional w2 zero points for int4/int8 quantization. + - is_scale_swizzled: Whether the activation scale-factor layout is + swizzled. Pass through to the underlying quantization kernel for + dtypes that distinguish layouts (nvfp4, mxfp8). Defaults to True. + - gemm1_alpha: Optional MXFP4 TRTLLM SwiGLU alpha parameter. + - gemm1_beta: Optional MXFP4 TRTLLM SwiGLU beta parameter. + - gemm1_clamp_limit: Optional MXFP4 TRTLLM SwiGLU clamp limit. + """ + assert not isinstance(quant_dtype, str) or quant_dtype in { + "nvfp4", + "mxfp4", + "mxfp6_e3m2", + "mxfp6_e2m3", + "mxfp8", + } + assert not isinstance(weight_dtype, str) or weight_dtype in { + "nvfp4", + "mxfp4", + "mxfp6_e3m2", + "mxfp6_e2m3", + "int4", + "mxfp8", + } + + if weight_dtype is None: + weight_dtype = quant_dtype + + a_shape, w_shape = _quant_flags_to_group_shape( + quant_dtype, per_act_token_quant, per_out_ch_quant, block_shape + ) + quant_config = FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(quant_dtype, a_shape, a1_scale, a1_gscale), + _a2=FusedMoEQuantDesc(quant_dtype, a_shape, a2_scale, a2_gscale), + _w1=FusedMoEQuantDesc( + weight_dtype, w_shape, w1_scale, g1_alphas, w1_zp, w1_bias + ), + _w2=FusedMoEQuantDesc( + weight_dtype, w_shape, w2_scale, g2_alphas, w2_zp, w2_bias + ), + is_scale_swizzled=is_scale_swizzled, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, + ) + assert quant_config.per_act_token_quant == per_act_token_quant + assert quant_config.per_out_ch_quant == per_out_ch_quant + assert quant_config.block_shape == block_shape + return quant_config + + +def fp8_w8a8_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + a1_scale: torch.Tensor | None = None, + a2_scale: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + per_act_token_quant: bool = False, + per_out_ch_quant: bool = False, + block_shape: list[int] | None = None, + a1_gscale: torch.Tensor | None = None, + a2_gscale: torch.Tensor | None = None, + g1_alphas: torch.Tensor | None = None, + g2_alphas: torch.Tensor | None = None, + gemm1_clamp_limit: float | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for fp8 activations and fp8 weights. + """ + return FusedMoEQuantConfig.make( + current_platform.fp8_dtype(), + w1_scale=w1_scale, + g1_alphas=g1_alphas, + w2_scale=w2_scale, + g2_alphas=g2_alphas, + w1_bias=w1_bias, + w2_bias=w2_bias, + a1_scale=a1_scale, + a1_gscale=a1_gscale, + a2_scale=a2_scale, + a2_gscale=a2_gscale, + per_act_token_quant=per_act_token_quant, + per_out_ch_quant=per_out_ch_quant, + block_shape=block_shape, + gemm1_clamp_limit=gemm1_clamp_limit, + ) + + +def int8_w8a8_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + a1_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + per_act_token_quant: bool = False, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for int8 activations and int8 weights. + """ + return FusedMoEQuantConfig.make( + torch.int8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + w1_bias=w1_bias, + w2_bias=w2_bias, + per_act_token_quant=per_act_token_quant, + per_out_ch_quant=False, + block_shape=None, + ) + + +def gptq_marlin_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + weight_bits: int, + group_size: int, + w1_zp: torch.Tensor | None = None, + w2_zp: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, +): + """ + Construct a quant config for gptq marlin quantization. + """ + from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape + + w_shape = None if group_size == -1 else GroupShape(row=1, col=group_size) + + # Activations are NOT quantized for GPTQ (fp16/bf16) + a_shape = w_shape # Same as weight shape for alignment + + # Determine weight dtype + if weight_bits == 4: + weight_dtype = "int4" + elif weight_bits == 8: + weight_dtype = torch.int8 + else: + raise ValueError(f"Unsupported weight_bits: {weight_bits}") + + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(dtype=None, shape=a_shape), + _a2=FusedMoEQuantDesc(dtype=None, shape=a_shape), + _w1=FusedMoEQuantDesc(weight_dtype, w_shape, w1_scale, None, w1_zp, w1_bias), + _w2=FusedMoEQuantDesc(weight_dtype, w_shape, w2_scale, None, w2_zp, w2_bias), + ) + + +def mxfp4_w4a16_moe_quant_config( + w1_scale: Union[torch.Tensor, "PrecisionConfig"], + w2_scale: Union[torch.Tensor, "PrecisionConfig"], + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, + gemm1_clamp_limit: float | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for unquantized activations and mxfp4 weights. + """ + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(), + _a2=FusedMoEQuantDesc(), + _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias), + _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias), + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, + ) + + +def mxfp4_mxfp8_moe_quant_config( + w1_scale: Union[torch.Tensor, "PrecisionConfig"], + w2_scale: Union[torch.Tensor, "PrecisionConfig"], + a1_scale: torch.Tensor | None = None, + a2_scale: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + block_shape: list[int] | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, + gemm1_clamp_limit: float | None = None, + mx_alignment: int = 0, + is_scale_swizzled: bool = True, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for mxfp4 activations and mxfp4 weights. + """ + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc("mxfp8"), + _a2=FusedMoEQuantDesc("mxfp8"), + _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias), + _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias), + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, + mx_alignment=mx_alignment, + is_scale_swizzled=is_scale_swizzled, + ) + + +def mxfp4_w4a8_moe_quant_config( + w1_scale: Union[torch.Tensor, "PrecisionConfig"], + w2_scale: Union[torch.Tensor, "PrecisionConfig"], + a1_scale: torch.Tensor | None = None, + a2_scale: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + block_shape: list[int] | None = None, + gemm1_clamp_limit: float | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for fp8 activations and mxfp4 weights. + """ + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc("fp8", None, a1_scale, None, None, None), + _a2=FusedMoEQuantDesc("fp8", None, a2_scale, None, None, None), + _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias), + _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias), + gemm1_clamp_limit=gemm1_clamp_limit, + ) + + +def ocp_mx_moe_quant_config( + quant_dtype: str, + w1_scale: Union[torch.Tensor, "PrecisionConfig"], + w2_scale: Union[torch.Tensor, "PrecisionConfig"], + weight_dtype: str | None = None, + a1_scale: torch.Tensor | None = None, + a2_scale: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + block_shape: list[int] | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, + gemm1_clamp_limit: float | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for mxfp4 activations and mxfp4 weights. + """ + assert quant_dtype in OCP_MX_DTYPES + return FusedMoEQuantConfig.make( + quant_dtype=quant_dtype, + weight_dtype=weight_dtype, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + w1_bias=w1_bias, + w2_bias=w2_bias, + per_act_token_quant=False, + per_out_ch_quant=False, + block_shape=block_shape, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, + ) + + +def nvfp4_moe_quant_config( + g1_alphas: torch.Tensor, + g2_alphas: torch.Tensor, + a1_gscale: torch.Tensor, + a2_gscale: torch.Tensor, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + is_scale_swizzled: bool = True, + gemm1_clamp_limit: float | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for mxfp4 activations and nvp4 weights. + """ + return FusedMoEQuantConfig.make( + "nvfp4", + w1_scale=w1_scale, + w2_scale=w2_scale, + w1_bias=w1_bias, + w2_bias=w2_bias, + a1_gscale=a1_gscale, + a2_gscale=a2_gscale, + g1_alphas=g1_alphas, + g2_alphas=g2_alphas, + per_act_token_quant=False, + per_out_ch_quant=False, + block_shape=None, + is_scale_swizzled=is_scale_swizzled, + gemm1_clamp_limit=gemm1_clamp_limit, + ) + + +def mxfp4_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for MXFP4 x MXFP4 MoE. + MXFP4 uses block scaling only (E8M0 scales, 32-element groups), with no + separate alphas / global activation scales in this config. + """ + return FusedMoEQuantConfig.make( + "mxfp4", + w1_scale=w1_scale, + w2_scale=w2_scale, + per_act_token_quant=False, + per_out_ch_quant=False, + block_shape=None, + ) + + +def nvfp4_w4a16_moe_quant_config( + g1_alphas: torch.Tensor, + g2_alphas: torch.Tensor, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for 16-but activations and nvp4 weights. + """ + return FusedMoEQuantConfig.make( + quant_dtype=None, + w1_scale=w1_scale, + w2_scale=w2_scale, + g1_alphas=g1_alphas, + g2_alphas=g2_alphas, + weight_dtype="nvfp4", + ) + + +def int4_w4a16_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + w1_zp: torch.Tensor | None = None, + w2_zp: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + block_shape: list[int] | None = None, + a1_gscale: torch.Tensor | None = None, + a2_gscale: torch.Tensor | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for 16-bit float activations and int4 weights. + """ + group_shape = GroupShape(*block_shape) if block_shape is not None else None + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a1_gscale), + _a2=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a2_gscale), + _w1=FusedMoEQuantDesc("int4", group_shape, w1_scale, None, w1_zp, w1_bias), + _w2=FusedMoEQuantDesc("int4", group_shape, w2_scale, None, w2_zp, w2_bias), + ) + + +def fp8_w8a16_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + block_shape: list[int] | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for 16-bit float activations and fp8 weights. + """ + group_shape = GroupShape(*block_shape) if block_shape is not None else None + fp8_dtype = current_platform.fp8_dtype() + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(), + _a2=FusedMoEQuantDesc(), + _w1=FusedMoEQuantDesc( + fp8_dtype, + group_shape, + w1_scale, + None, + None, + w1_bias, + ), + _w2=FusedMoEQuantDesc( + fp8_dtype, + group_shape, + w2_scale, + None, + None, + w2_bias, + ), + ) + + +def int8_w8a16_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + w1_zp: torch.Tensor | None = None, + w2_zp: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + block_shape: list[int] | None = None, + a1_gscale: torch.Tensor | None = None, + a2_gscale: torch.Tensor | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for 16-bit float activations and int8 weights. + """ + group_shape = GroupShape(*block_shape) if block_shape is not None else None + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a1_gscale), + _a2=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a2_gscale), + _w1=FusedMoEQuantDesc(torch.int8, group_shape, w1_scale, None, w1_zp, w1_bias), + _w2=FusedMoEQuantDesc(torch.int8, group_shape, w2_scale, None, w2_zp, w2_bias), + ) + + +def int4_w4afp8_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + g1_alphas: torch.Tensor, + g2_alphas: torch.Tensor, + per_act_token_quant: bool = False, + per_out_ch_quant: bool = False, + block_shape: list[int] | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for fp8 activations and int4 weights. + """ + return FusedMoEQuantConfig.make( + torch.float8_e4m3fn, # quant dtype for activations + w1_scale=w1_scale, + w2_scale=w2_scale, + g1_alphas=g1_alphas, + g2_alphas=g2_alphas, + per_act_token_quant=per_act_token_quant, + per_out_ch_quant=per_out_ch_quant, + block_shape=block_shape, + weight_dtype="int4", # weight dtype for weights + ) + + +def biased_moe_quant_config( + w1_bias: torch.Tensor | None, + w2_bias: torch.Tensor | None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for unquantized activations with biases. + """ + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(), + _a2=FusedMoEQuantDesc(), + _w1=FusedMoEQuantDesc(bias=w1_bias), + _w2=FusedMoEQuantDesc(bias=w2_bias), + ) + + +# A FusedMoEQuantConfig constant for an unquantized MoE op. +FUSED_MOE_UNQUANTIZED_CONFIG: FusedMoEQuantConfig = FusedMoEQuantConfig.make() + + +@dataclass +class FusedMoEParallelConfig: + tp_size: int + pcp_size: int + dp_size: int + ep_size: int + tp_rank: int + pcp_rank: int + dp_rank: int + ep_rank: int + sp_size: int + + use_ep: bool # whether to use EP or not + all2all_backend: str # all2all backend for MoE communication + enable_eplb: bool # whether to enable expert load balancing + + @property + def is_sequence_parallel(self) -> bool: + return self.sp_size > 1 + + @property + def use_all2all_kernels(self): + return self.dp_size > 1 and self.use_ep + + @property + def use_deepep_ht_kernels(self): + return ( + self.use_all2all_kernels + and self.all2all_backend == "deepep_high_throughput" + ) + + @property + def use_deepep_ll_kernels(self): + return self.use_all2all_kernels and self.all2all_backend == "deepep_low_latency" + + @property + def use_fi_nvl_two_sided_kernels(self): + return self.use_all2all_kernels and ( + self.all2all_backend == "flashinfer_all2allv" + or self.all2all_backend == "flashinfer_nvlink_two_sided" + ) + + @property + def use_fi_nvl_one_sided_kernels(self): + return ( + self.use_all2all_kernels + and self.all2all_backend == "flashinfer_nvlink_one_sided" + ) + + @property + def use_batched_activation_format(self): + return self.use_deepep_ll_kernels or self.use_nixl_ep_kernels + + @property + def needs_round_robin_routing_tables(self): + return self.use_deepep_ll_kernels or self.use_nixl_ep_kernels + + @property + def use_ag_rs_all2all_kernels(self): + return ( + self.use_all2all_kernels + and self.all2all_backend == "allgather_reducescatter" + ) + + @property + def use_mori_kernels(self): + return self.use_all2all_kernels and self.all2all_backend in ( + "mori_high_throughput", + "mori_low_latency", + ) + + @property + def use_nixl_ep_kernels(self): + return self.use_all2all_kernels and self.all2all_backend == "nixl_ep" + + @property + def use_deepep_v2_kernels(self): + return self.use_all2all_kernels and self.all2all_backend == "deepep_v2" + + @staticmethod + def flatten_tp_across_dp_and_pcp( + tp_size: int, dp_size: int, dp_rank: int, pcp_size: int, pcp_rank: int + ) -> tuple[int, int]: + tp_rank = 0 if tp_size == 1 else get_tensor_model_parallel_rank() + # There are actually dp_size * pcp_size * tp_size devices. + # Update tp_size and tp_rank so we shard across all devices. + flatten_tp_size = dp_size * pcp_size * tp_size + flatten_tp_rank = dp_rank * pcp_size * tp_size + pcp_rank * tp_size + tp_rank + return flatten_tp_size, flatten_tp_rank + + @staticmethod + def make( + tp_size_: int, + pcp_size_: int, + dp_size_: int, + sp_size_: int, + vllm_parallel_config: ParallelConfig, + ) -> "FusedMoEParallelConfig": + """ + Determine MoE parallel configuration. Based on the input `tp_size_`, + `dp_size_` and vllm's parallel config, determine what + level's of parallelism to use in the fused moe layer. + + Args: + tp_size_ (int): `tp_size` passed into the FusedMoE constructor. + pcp_size_ (int): `pcp_size` passed into the FusedMoE constructor. + dp_size_ (int): `dp_size` passed into the FusedMoE constructor. + vllm_parallel_config (ParallelConfig): vLLM's parallel config + object which contains the `enable_expert_parallel` flag. + + Examples: + When there is no parallelism requested, + i.e. `tp_size_` = `pcp_size_` = `dp_size_` = 1, we simply return the sizes + unaltered and the ranks set to 0. + + Expert Parallelism is considered only when either `dp_size_`, `pcp_size_` or + `tp_size_` is non trivial. + + Note that PCP serves the same function as DP here. + + When TP = 2, DP(PCP) = 1 and EP = False, the configuration on different + devices: + + - device 0 : TP = {2, 0} DP = {1, 0} EP = {1, 0} // + legend : {size, rank} + - device 1 : TP = {2, 1} DP = {1, 0} EP = {1, 0} + - Comment : Tensors are sharded across 2 devices. + + When TP = 1, DP(PCP) = 2 and EP = False, the configuration on different + devices: + + - device 0 : TP = {2, 0} DP = {2, 0} EP = {1, 0} + - device 1 : TP = {2, 1} DP = {2, 1} EP = {1, 0} + - Comment: There are 2 engine instances and the tensors are sharded + across 2 decvices. + + When TP = 2, DP(PCP) = 2 and EP = False, the configuration on different + devices: + + - device 0: TP = {4, 0} DP = {2, 0} EP = {1, 0} + - device 1: TP = {4, 1} DP = {2, 0} EP = {1, 0} + - device 2: TP = {4, 2} DP = {2, 1} EP = {1, 0} + - device 3: TP = {4, 3} DP = {2, 1} EP = {1, 0} + - Comment: There are 2 engine instances and the tensors are sharded + across 4 devices. + + When, TP = 2, DP(PCP) = 1 and EP = True, the configuration on different + devices: + + - device 0: TP = {1, 0} DP = {1, 0} EP = {2, 0} + - device 1: TP = {1, 0} DP = {1, 0} EP = {2, 1} + - Comment: The experts are split between the 2 devices. + + When, TP = 1, DP(PCP) = 2 and EP = True, the configuration on different + devices: + + - device 0: TP = {1, 0} DP = {2, 0} EP = {2, 0} + - device 1: TP = {1, 0} DP = {2, 1} EP = {2, 1} + - Comment: There are 2 engine instances and the experts are split + between the 2 devices. + + When TP = 2, DP(PCP) = 2 and EP = True, the configuration on different + devices: + + - device 0: TP = {1, 0} DP = {2, 0} EP = {4, 0} + - device 1: TP = {1, 0} DP = {2, 0} EP = {4, 1} + - device 2: TP = {1, 0} DP = {2, 1} EP = {4, 2} + - device 3: TP = {1, 0} DP = {2, 1} EP = {4, 3} + - Comment: There are 2 engine instances and the experts are split + between the 4 devices. + """ + use_ep = ( + dp_size_ * pcp_size_ * tp_size_ > 1 + and vllm_parallel_config.enable_expert_parallel + ) + + dp_size = dp_size_ + dp_rank = get_dp_group().rank_in_group if dp_size > 1 else 0 + pcp_size = pcp_size_ + pcp_rank = get_pcp_group().rank_in_group if pcp_size > 1 else 0 + tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp( + tp_size_, dp_size_, dp_rank, pcp_size_, pcp_rank + ) + + if not use_ep: + return FusedMoEParallelConfig( + tp_size=tp_size, + tp_rank=tp_rank, + pcp_size=pcp_size, + pcp_rank=pcp_rank, + dp_size=dp_size, + dp_rank=dp_rank, + ep_size=1, + ep_rank=0, + sp_size=sp_size_, + use_ep=False, + all2all_backend=vllm_parallel_config.all2all_backend, + enable_eplb=vllm_parallel_config.enable_eplb, + ) + # DP + EP / TP + EP / DP + TP + EP + assert use_ep + # In EP, each device owns a set of experts fully. There is no tensor + # parallel update tp_size, tp_rank, ep_size and ep_rank to reflect that. + ep_size = tp_size + ep_rank = tp_rank + return FusedMoEParallelConfig( + tp_size=1, + tp_rank=0, + pcp_size=pcp_size, + pcp_rank=pcp_rank, + dp_size=dp_size, + dp_rank=dp_rank, + ep_size=ep_size, + ep_rank=ep_rank, + sp_size=sp_size_, + use_ep=True, + all2all_backend=vllm_parallel_config.all2all_backend, + enable_eplb=vllm_parallel_config.enable_eplb, + ) + + @classmethod + def make_no_parallel(cls) -> "FusedMoEParallelConfig": + """For usage in CI/CD and testing.""" + return FusedMoEParallelConfig( + tp_size=1, + tp_rank=0, + pcp_size=1, + pcp_rank=0, + dp_size=1, + dp_rank=0, + ep_size=1, + ep_rank=0, + sp_size=1, + use_ep=False, + all2all_backend="allgather_reducescatter", + enable_eplb=False, + ) + + +# Adapted from pplx-kernels tests/all_to_all_utils.py +@dataclass +class FusedMoEConfig: + num_experts: int + experts_per_token: int + hidden_dim: int + intermediate_size: int + num_local_experts: int + num_logical_experts: int + activation: MoEActivation + device: torch.device | str + routing_method: RoutingMethodType + moe_parallel_config: FusedMoEParallelConfig + + # The activation type. + in_dtype: torch.dtype + + # Defaults to in_dtype if not specified. + router_logits_dtype: torch.dtype | None = None + + # Defaults to hidden_dim if not specified. + hidden_dim_unpadded: int | None = None + # Defaults to intermediate_size_per_partition if not specified. + intermediate_size_per_partition_unpadded: int | None = None + + moe_backend: MoEBackend = "auto" + max_num_tokens: int = SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP + has_bias: bool = False + is_lora_enabled: bool = False + + # SwiGLU clamp limit. When set, backends that do not implement the clamp + # are filtered out by `FusedMoEExperts.is_supported_config` so the oracle + # cannot silently select one and drop the clamp. + swiglu_limit: float | None = None + + max_capture_size: int = 0 + + # Set by __post_init__ + intermediate_size_per_partition: int = -1 + rocm_aiter_fmoe_enabled: bool = False + aiter_fmoe_shared_expert_enabled: bool = False + + def __post_init__(self): + from vllm._aiter_ops import rocm_aiter_ops + + tp_size = self.moe_parallel_config.tp_size + assert self.intermediate_size % tp_size == 0 + self.intermediate_size_per_partition = self.intermediate_size // tp_size + + if self.dp_size > 1: + logger.debug_once( + "Using FusedMoEConfig::max_num_tokens=%d", self.max_num_tokens + ) + + assert self.max_num_tokens > 0 + + if self.router_logits_dtype is None: + self.router_logits_dtype = self.in_dtype + + if self.hidden_dim_unpadded is None: + self.hidden_dim_unpadded = self.hidden_dim + if self.intermediate_size_per_partition_unpadded is None: + self.intermediate_size_per_partition_unpadded = ( + self.intermediate_size_per_partition + ) + + if self.is_act_and_mul: + self.rocm_aiter_fmoe_enabled = rocm_aiter_ops.is_fused_moe_enabled() + self.aiter_fmoe_shared_expert_enabled = ( + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + ) + + if self.use_mori_kernels: + assert self.rocm_aiter_fmoe_enabled, ( + "Mori needs to be used with aiter fused_moe for now." + ) + assert not self.aiter_fmoe_shared_expert_enabled, ( + "Mori does not support fusion shared expert now. " + "Turn it off by setting VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0" + ) + + if not self.is_act_and_mul and not ( + current_platform.is_cuda_alike() or current_platform.is_xpu() + ): + raise NotImplementedError( + "is_act_and_mul=False is supported only for CUDA, XPU and ROCm for now" + ) + + @property + def is_act_and_mul(self) -> bool: + return self.activation.is_gated + + @property + def tp_size(self): + return self.moe_parallel_config.tp_size + + @property + def dp_size(self): + return self.moe_parallel_config.dp_size + + @property + def pcp_size(self): + return self.moe_parallel_config.pcp_size + + @property + def ep_size(self): + return self.moe_parallel_config.ep_size + + @property + def sp_size(self): + return self.moe_parallel_config.sp_size + + @property + def is_sequence_parallel(self): + return self.moe_parallel_config.is_sequence_parallel + + @property + def tp_rank(self): + return self.moe_parallel_config.tp_rank + + @property + def dp_rank(self): + return self.moe_parallel_config.dp_rank + + @property + def pcp_rank(self): + return self.moe_parallel_config.pcp_rank + + @property + def ep_rank(self): + return self.moe_parallel_config.ep_rank + + @property + def use_ep(self): + return self.moe_parallel_config.use_ep + + @property + def use_deepep_ht_kernels(self): + return self.moe_parallel_config.use_deepep_ht_kernels + + @property + def use_deepep_ll_kernels(self): + return self.moe_parallel_config.use_deepep_ll_kernels + + @property + def use_mori_kernels(self): + return self.moe_parallel_config.use_mori_kernels + + @property + def use_fi_nvl_two_sided_kernels(self): + return self.moe_parallel_config.use_fi_nvl_two_sided_kernels + + @property + def use_fi_nvl_one_sided_kernels(self): + return self.moe_parallel_config.use_fi_nvl_one_sided_kernels + + @property + def use_ag_rs_all2all_kernels(self): + return self.moe_parallel_config.use_ag_rs_all2all_kernels + + @property + def use_nixl_ep_kernels(self): + return self.moe_parallel_config.use_nixl_ep_kernels + + @property + def use_deepep_v2_kernels(self): + return self.moe_parallel_config.use_deepep_v2_kernels + + @property + def needs_round_robin_routing_tables(self): + return self.moe_parallel_config.needs_round_robin_routing_tables diff --git a/ex_engine/moe/experts/__init__.py b/ex_engine/moe/experts/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ex_engine/moe/experts/fallback.py b/ex_engine/moe/experts/fallback.py new file mode 100644 index 00000000..639b2bf2 --- /dev/null +++ b/ex_engine/moe/experts/fallback.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey + + +class FallbackExperts(mk.FusedMoEExpertsModular, ABC): + """Base class for runtime dispatching of expert implementations.""" + + def __init__( + self, + experts: mk.FusedMoEExpertsModular, + fallback_experts: mk.FusedMoEExpertsModular, + ): + super().__init__( + moe_config=experts.moe_config, quant_config=experts.quant_config + ) + self.fallback_experts = fallback_experts + self.experts = experts + + @staticmethod + def get_clses() -> tuple[ + type[mk.FusedMoEExpertsModular], + type[mk.FusedMoEExpertsModular], + ]: + """ + Get the cls for the experts and fallback experts. + + Subclasses should implement this method, so that + we have a consistent way to call the _supports_* + class methods below. + """ + raise NotImplementedError( + "Subclasses must return the cls for the experts and fallback experts." + ) + + @classmethod + def activation_format( + cls: type["FallbackExperts"], + ) -> mk.FusedMoEActivationFormat: + experts_cls, fallback_cls = cls.get_clses() + assert experts_cls.activation_format() == fallback_cls.activation_format() + return experts_cls.activation_format() + + @classmethod + def _supports_current_device(cls) -> bool: + experts_cls, fallback_cls = cls.get_clses() + return ( + experts_cls._supports_current_device() + and fallback_cls._supports_current_device() + ) + + @classmethod + def _supports_no_act_and_mul(cls) -> bool: + experts_cls, fallback_cls = cls.get_clses() + return ( + experts_cls._supports_no_act_and_mul() + and fallback_cls._supports_no_act_and_mul() + ) + + @classmethod + def _supports_quant_scheme( + cls, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + experts_cls, fallback_cls = cls.get_clses() + return experts_cls._supports_quant_scheme( + weight_key, activation_key + ) and fallback_cls._supports_quant_scheme(weight_key, activation_key) + + @classmethod + def _supports_activation(cls, activation: MoEActivation) -> bool: + experts_cls, fallback_cls = cls.get_clses() + return experts_cls._supports_activation( + activation + ) and fallback_cls._supports_activation(activation) + + @classmethod + def _supports_parallel_config( + cls, moe_parallel_config: FusedMoEParallelConfig + ) -> bool: + experts_cls, fallback_cls = cls.get_clses() + return experts_cls._supports_parallel_config( + moe_parallel_config + ) and fallback_cls._supports_parallel_config(moe_parallel_config) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + e_war = self.experts.finalize_weight_and_reduce_impl() + fbe_war = self.fallback_experts.finalize_weight_and_reduce_impl() + is_dge_war = e_war is not None + is_fbe_war = fbe_war is not None + + if is_dge_war and is_fbe_war: + assert e_war == fbe_war, ( + "Both implementations should agree on WeightAndReduce impls. " + f"Got e_war: {e_war}, and fbe_war: {fbe_war}" + ) + + if e_war is not None: + return e_war + assert fbe_war is not None + return fbe_war + + @abstractmethod + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + raise NotImplementedError + + @abstractmethod + def _select_experts_impl( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + ) -> mk.FusedMoEExpertsModular: + raise NotImplementedError + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + experts = self._select_experts_impl(hidden_states, w1, w2) + experts.apply( + output, + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + activation, + global_num_experts, + expert_map, + a1q_scale, + a2_scale, + workspace13, + workspace2, + expert_tokens_meta, + apply_router_weight_on_input, + ) diff --git a/ex_engine/moe/experts/fused_batched_moe.py b/ex_engine/moe/experts/fused_batched_moe.py new file mode 100644 index 00000000..1f5724ac --- /dev/null +++ b/ex_engine/moe/experts/fused_batched_moe.py @@ -0,0 +1,972 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused batched MoE kernel.""" + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.fused_moe import try_get_optimal_moe_config +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceDelegate, +) +from vllm.model_executor.layers.fused_moe.utils import ( + _resize_cache, + moe_kernel_quantize_input, + normalize_batched_scales_shape, + swiglu_limit_func, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + group_broadcast, + kFp8Dynamic128Sym, + kFp8DynamicTensorSym, + kFp8DynamicTokenSym, + kFp8Static128BlockSym, + kFp8StaticChannelSym, + kFp8StaticTensorSym, +) +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + + +@triton.jit +def moe_mmk( + a_ptrs, + b_ptrs, + K, + expert_id, + a_scale_ptr, + b_scale_ptr, + # 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_ak: tl.int64, + stride_bk: tl.int64, + stride_ase: tl.int64, + stride_asm: tl.int64, + stride_ask: tl.int64, + stride_bse: tl.int64, + stride_bsk: tl.int64, + stride_bsn: tl.int64, + # Offsets and masks + offs_m, + offs_n, + offs_bn, + mask_m, + # Block size for block-wise quantization + group_n: tl.constexpr, + group_k: tl.constexpr, + # Meta-parameters + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + compute_type: tl.constexpr, + use_w8a8: tl.constexpr, + use_w8a16: tl.constexpr, + per_act_token_quant: tl.constexpr, +): + offs_k = tl.arange(0, BLOCK_K) + + if use_w8a16: + b_scale_ptrs = ( + b_scale_ptr + expert_id * stride_bse + offs_n[None, :] * stride_bsn + ) + b_scale = tl.load(b_scale_ptrs) + + if use_w8a8: + # block-wise + if group_k > 0 and group_n > 0: + a_scale_ptrs = a_scale_ptr + offs_m * stride_asm + offs_bsn = offs_bn // group_n + b_scale_ptrs = b_scale_ptr + offs_bsn * stride_bsn + + # per act token + elif per_act_token_quant: + # Load per-token scale for activations + a_scale_ptrs = a_scale_ptr + offs_m * stride_asm + a_scale = tl.load(a_scale_ptrs, mask=mask_m, other=0.0)[:, None] + + b_scale_ptrs = b_scale_ptr + offs_bn[None, :] * stride_bsn + b_scale = tl.load(b_scale_ptrs) + + # tensor-wise + else: + a_scale = tl.load(a_scale_ptr) + b_scale = tl.load(b_scale_ptr) + + # ----------------------------------------------------------- + # 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_M, BLOCK_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_K)): + # Load the next block of A and B, generate a mask by checking the + # K dimension. + a = tl.load( + a_ptrs, + mask=mask_m[:, None] & (offs_k[None, :] < K - k * BLOCK_K), + other=0.0, + ) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0) + # We accumulate along the K dimension. + if use_w8a16: + accumulator = tl.dot(a, b.to(compute_type), acc=accumulator) + elif use_w8a8: + if group_k > 0 and group_n > 0: + k_start = k * BLOCK_K + offs_ks = k_start // group_k + a_scale = tl.load( + a_scale_ptrs + offs_ks * stride_ask, mask=mask_m, 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: + # acc used to enable fp8_fast_accum + accumulator = tl.dot(a, b, acc=accumulator) + else: + accumulator += tl.dot(a, b) + + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_K * stride_ak + b_ptrs += BLOCK_K * stride_bk + + if use_w8a16: + accumulator = (accumulator * b_scale).to(compute_type) + elif use_w8a8: + if group_k > 0 and group_n > 0: + accumulator = accumulator.to(compute_type) + else: + accumulator = (accumulator * a_scale * b_scale).to(compute_type) + else: + accumulator = accumulator.to(compute_type) + + return accumulator + + +@triton.jit +def expert_triton_kernel( + a_ptr, # [max_tokens, K] + b_ptr, # [K, N] + c_ptr, # [max_tokens, N] + expert_id, + compute_type: tl.constexpr, + # Dimensions + M, + N, + K, + # Quantization data + a_scale_ptr, + b_scale_ptr, + b_zp_ptr, + # strides + stride_am: tl.int64, + stride_ak: tl.int64, + stride_bk: tl.int64, + stride_bn: tl.int64, + stride_cm: tl.int64, + stride_cn: tl.int64, + stride_ase: tl.int64, + stride_asm: tl.int64, + stride_ask: tl.int64, + stride_bse: tl.int64, + stride_bsk: tl.int64, + stride_bsn: tl.int64, + # offsets + offs_bn, + # Blockwise quantization data + group_n, + group_k, + # Quantization schemes + use_fp8_w8a8: tl.constexpr, + use_int8_w8a16: tl.constexpr, + per_act_token_quant: tl.constexpr, + # Kernel config + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + offs_m = tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) % N + offs_k = tl.arange(0, BLOCK_K) + mask_m = offs_m < M + + # Make grids of a + b pointers + a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak + b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn + + accumulator = moe_mmk( + a_ptrs, + b_ptrs, + K, + expert_id, + a_scale_ptr, + b_scale_ptr, + # 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_ak, + stride_bk, + stride_ase, + stride_asm, + stride_ask, + stride_bse, + stride_bsk, + stride_bsn, + # Offsets and masks + offs_m, + offs_n, + offs_bn, + mask_m, + # Block size for block-wise quantization + group_n, + group_k, + # Meta-parameters + BLOCK_M, + BLOCK_N, + BLOCK_K, + compute_type, + use_fp8_w8a8, + use_int8_w8a16, + per_act_token_quant, + ) + + # store in C + offs_cn = tl.arange(0, BLOCK_N) + c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_cn[None, :] * stride_cn + c_mask = mask_m[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + + +@triton.jit +def batched_triton_kernel( + a_ptr, # [E, max_num_tokens, K] + b_ptr, # [E, K, N] + c_ptr, # [E, max_num_tokens, N] + expert_num_tokens, # [E] + compute_type: tl.constexpr, + # Dimensions + max_num_tokens, + K, + N, + # Quantization data + a_scale_ptr, + b_scale_ptr, + b_zp_ptr, + # 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_ae: tl.int64, + stride_am: tl.int64, + stride_ak: tl.int64, + stride_be: tl.int64, + stride_bk: tl.int64, + stride_bn: tl.int64, + stride_ce: tl.int64, + stride_cm: tl.int64, + stride_cn: tl.int64, + stride_ase: tl.int64, + stride_asm: tl.int64, + stride_ask: tl.int64, + stride_bse: tl.int64, + stride_bsk: tl.int64, + stride_bsn: tl.int64, + # Blockwise quantization data + group_n: tl.constexpr, + group_k: tl.constexpr, + # Quantization schemes + use_fp8_w8a8: tl.constexpr, + use_int8_w8a16: tl.constexpr, + per_act_token_quant: tl.constexpr, + # Kernel config + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + expert_id = tl.program_id(axis=0) + e_num_tokens = tl.load(expert_num_tokens + expert_id) + if e_num_tokens == 0: + # Early exit + return + + # axis 1 is M_blocks * N_blocks + pid_mn = tl.program_id(axis=1) + # num_pid_m = tl.cdiv(max_num_tokens, BLOCK_M) + num_pid_n = tl.cdiv(N, BLOCK_N) + pid_m = pid_mn // num_pid_n + pid_n = pid_mn % num_pid_n + + cta_m_start = pid_m * BLOCK_M + cta_n_start = pid_n * BLOCK_N + if cta_m_start >= e_num_tokens: + # Early exit + return + + cta_m_size = min(BLOCK_M, e_num_tokens - cta_m_start) + cta_n_size = min(BLOCK_N, N - cta_n_start) + + a_ptr = a_ptr + expert_id * stride_ae + cta_m_start * stride_am + b_ptr = b_ptr + expert_id * stride_be + cta_n_start * stride_bn + c_ptr = ( + c_ptr + + expert_id * stride_ce + + cta_m_start * stride_cm + + cta_n_start * stride_cn + ) + + offs_bn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N).to(tl.int64)) % N + + if use_fp8_w8a8: + a_scale_ptr = a_scale_ptr + expert_id * stride_ase + b_scale_ptr = b_scale_ptr + expert_id * stride_bse + + # block-wise + if group_k > 0 and group_n > 0 or per_act_token_quant: + a_scale_ptr = a_scale_ptr + cta_m_start * stride_asm + + expert_triton_kernel( + a_ptr, + b_ptr, + c_ptr, + expert_id, + compute_type, + cta_m_size, # M + cta_n_size, # N + K, # K + a_scale_ptr, + b_scale_ptr, + b_zp_ptr, + # Strides + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_ase, + stride_asm, + stride_ask, + stride_bse, + stride_bsk, + stride_bsn, + # offsets + offs_bn, + # Blockwise quantization data + group_n, + group_k, + # Quantization schemes + use_fp8_w8a8, + use_int8_w8a16, + per_act_token_quant, + # Kernel config + BLOCK_M, + BLOCK_N, + BLOCK_K, + ) + + +def invoke_moe_batched_triton_kernel( + A: torch.Tensor, # [E, max_tokens, K] + B: torch.Tensor, # [E, N, K] + C: torch.Tensor, # [E, max_tokens, N] + expert_num_tokens: torch.Tensor, # [E] + compute_type: tl.dtype, + # Quantization data + A_scale: torch.Tensor | None, + B_scale: torch.Tensor | None, + B_zp: torch.Tensor, + # Quantization schemes + use_fp8_w8a8: bool, + use_int8_w8a16: bool, + use_int4_w4a16: bool, + config: dict[str, int], + per_act_token_quant: bool, + block_shape: list[int] | None = None, +): + assert not use_int4_w4a16 + max_num_tokens = A.size(1) + K = A.size(2) + N = C.size(2) + + BLOCK_M = config["BLOCK_SIZE_M"] + BLOCK_N = config["BLOCK_SIZE_N"] + BLOCK_K = config["BLOCK_SIZE_K"] + + grid = ( + expert_num_tokens.size(0), + triton.cdiv(max_num_tokens, BLOCK_M) * triton.cdiv(B.size(1), BLOCK_N), + ) + + A_scale = normalize_batched_scales_shape(A_scale, expert_num_tokens.shape[0]) + + if B_scale is not None and B_scale.ndim == 1: + assert B_scale.numel() == expert_num_tokens.shape[0] + B_scale = B_scale.view(-1, 1, 1) + + assert A_scale is None or A_scale.ndim == 3, ( + f"{0 if A_scale is None else A_scale.shape}" + ) + assert B_scale is None or B_scale.ndim == 1 or B_scale.ndim == 3, ( + f"{0 if B_scale is None else B_scale.shape}" + ) + + if B_scale is not None: + if B_scale.ndim == 1: + stride_bse = 1 + stride_bsk = 0 + stride_bsn = 0 + else: + stride_bse = B_scale.stride(0) + stride_bsk = B_scale.stride(2) + stride_bsn = B_scale.stride(1) + + else: + stride_bse = 0 + stride_bsk = 0 + stride_bsn = 0 + + if A_scale is not None: + stride_ase = A_scale.stride(0) + stride_asm = A_scale.stride(1) + stride_ask = A_scale.stride(2) + else: + stride_ase = 0 + stride_asm = 0 + stride_ask = 0 + + batched_triton_kernel[grid]( + A, + B, + C, + expert_num_tokens, + compute_type, + # Dimensions + max_num_tokens, + K, + N, + # Quantization data + A_scale, + B_scale, + B_zp, + # Strides + A.stride(0), + A.stride(1), + A.stride(2), + B.stride(0), + B.stride(2), + B.stride(1), + C.stride(0), + C.stride(1), + C.stride(2), + stride_ase, + stride_asm, + stride_ask, + stride_bse, + stride_bsk, + stride_bsn, + # Blockwise quantization data + 0 if block_shape is None else block_shape[0], + 0 if block_shape is None else block_shape[1], + # Quantization schemes + use_fp8_w8a8, + use_int8_w8a16, + per_act_token_quant, + # Kernel config + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_K=BLOCK_K, + ) + + +class NaiveBatchedExperts(mk.FusedMoEExpertsModular): + """ + A reference MoE expert class that operates on expert batched format, + i.e. E x max_num_tokens x K. This is the format that the batched + dispatch/combine kernels use. + """ + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + max_num_tokens: int, + num_dispatchers: int, + ): + super().__init__( + moe_config=moe_config, + quant_config=quant_config, + max_num_tokens=max_num_tokens, + num_dispatchers=num_dispatchers, + ) + assert not self.quant_config.use_int8_w8a8, "NYI" + assert not self.quant_config.use_int8_w8a16, "NYI" + assert not self.quant_config.use_int4_w4a16, "NYI" + assert self.quant_config.ocp_mx_scheme is None, "NYI" + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.BatchedExperts + + @staticmethod + def _supports_current_device() -> bool: + raise NotImplementedError( + "NaiveBatchedExperts is not yet used by an Oracle. " + "This method should not be called." + ) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + raise NotImplementedError( + "NaiveBatchedExperts is not yet used by an Oracle. " + "This method should not be called." + ) + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + raise NotImplementedError( + "NaiveBatchedExperts is not yet used by an Oracle. " + "This method should not be called." + ) + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + raise NotImplementedError( + "NaiveBatchedExperts is not yet used by an Oracle. " + "This method should not be called." + ) + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + raise NotImplementedError( + "NaiveBatchedExperts is not yet used by an Oracle. " + "This method should not be called." + ) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + # Let PrepareAndFinalize::finalize() decide the impl. + return TopKWeightAndReduceDelegate() + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + assert self.num_dispatchers is not None + assert self.max_num_tokens is not None + num_dp = self.num_dispatchers + num_experts = local_num_experts + workspace13 = (num_experts, self.max_num_tokens * num_dp, K) + workspace2 = (self.max_num_tokens * num_dp, N) + output = workspace13 + return (workspace13, workspace2, output) + + def dequant(self, t: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + assert self.quant_config.is_quantized + f32 = torch.float32 + if self.quant_config.is_per_act_token or self.quant_config.is_per_tensor: + return t.to(f32) * scale + else: + return t.to(f32) * group_broadcast(scale, t.shape) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + assert hidden_states.dim() == 3 + assert expert_tokens_meta is not None + expert_num_tokens = expert_tokens_meta.expert_num_tokens + + num_local_experts = w1.size(0) + assert num_local_experts == w1.size(0), f"{num_local_experts} == {w1.size(0)}" + + N = w1.size(1) // 2 + + for expert in range(num_local_experts): + # Indexing expert_num_tokens doesn't work w/cudagraphs or inductor + if ( + torch.compiler.is_compiling() + or torch.cuda.is_current_stream_capturing() + ): + num = hidden_states.shape[1] + else: + num = int(expert_num_tokens[expert].item()) + + if num == 0: + continue + + tmp = _resize_cache(workspace2, (num, N)) + + if self.quant_config.is_quantized: + assert a1q_scale is not None and self.w1_scale is not None + input = self.dequant(hidden_states[expert, :, :], a1q_scale[expert]) + w1_dq = self.dequant(w1[expert], self.w1_scale[expert]) + input = input[:num] @ w1_dq.transpose(0, 1) + else: + input = hidden_states[expert, :num, :] @ w1[expert].transpose(0, 1) + + self.activation(activation, tmp, input.to(tmp.dtype)) + + if self.quant_config.is_quantized: + assert self.w2_scale is not None + w2_dq = self.dequant(w2[expert], self.w2_scale[expert]) + else: + w2_dq = w2[expert] + + output[expert, :num, :] = tmp @ w2_dq.transpose(0, 1).to(tmp.dtype) + + +def batched_moe_kernel_quantize_input( + A: torch.Tensor, + A_scale: torch.Tensor | None, + num_tokens: int, + E: int, + N: int, + expert_num_tokens: torch.Tensor, + qtype: torch.dtype | None, + per_act_token_quant: bool, + block_shape: list[int] | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + if torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing(): + # Note: this does a bunch of extra work because expert_num_tokens is + # ignored but it does support torch.compile + cudagraphs. + hidden_dim = A.size(-1) + assert A_scale is None or A_scale.ndim <= 2, ( + f"{A_scale.shape if A_scale is not None else None}" + ) + A_q, A_q_scale = moe_kernel_quantize_input( + A.view(-1, hidden_dim), A_scale, qtype, per_act_token_quant, block_shape + ) + A_q = A_q.view(E, -1, hidden_dim) + A_q_scale = normalize_batched_scales_shape(A_q_scale, E) + + return A_q, A_q_scale + elif qtype is None: + return A, normalize_batched_scales_shape(A_scale, E) + else: + A_q = torch.empty_like(A, dtype=qtype) + + if per_act_token_quant: + assert block_shape is None + scale_shape = (E, num_tokens, 1) + elif block_shape is not None: + _, block_k = block_shape + k_tiles = (A.shape[-1] + block_k - 1) // block_k + scale_shape = (E, num_tokens, k_tiles) + else: + scale_shape = (E, 1, 1) + + A_q_scale = torch.zeros(scale_shape, dtype=torch.float32, device=A.device) + + num_experts = expert_num_tokens.numel() + + A_scale = normalize_batched_scales_shape(A_scale, num_experts) + + for e in range(E): + num_tokens = int(expert_num_tokens[e].item()) + if num_tokens > 0: + if A_scale is not None: + scales = A_scale[e, : min(num_tokens, A_scale.shape[1])] + else: + scales = None + A_q[e, :num_tokens], tmp_scale = moe_kernel_quantize_input( + A[e, :num_tokens], + scales, + qtype, + per_act_token_quant, + block_shape, + ) + assert tmp_scale is not None + A_q_scale[e, : tmp_scale.shape[0]] = tmp_scale + + return A_q, A_q_scale + + +class BatchedTritonExperts(mk.FusedMoEExpertsModular): + """ + A Triton based MoE expert class that operates on expert batched format, + i.e. E x max_num_tokens x K. This is the format that the batched + dispatch/combine kernels use. + """ + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + max_num_tokens: int, + num_dispatchers: int, + ): + super().__init__( + moe_config=moe_config, + quant_config=quant_config, + max_num_tokens=max_num_tokens, + num_dispatchers=num_dispatchers, + ) + assert not self.quant_config.use_int8_w8a8, "NYI" + assert not self.quant_config.use_int8_w8a16, "NYI" + assert not self.quant_config.use_int4_w4a16, "NYI" + assert self.quant_config.ocp_mx_scheme is None, "NYI" + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.BatchedExperts + + @staticmethod + def _supports_current_device() -> bool: + return current_platform.is_cuda_alike() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return True + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + p = current_platform + if p.is_rocm(): + from vllm.platforms.rocm import on_gfx9 + + is_rocm_on_gfx9 = on_gfx9() + else: + is_rocm_on_gfx9 = False + + device_supports_fp8 = is_rocm_on_gfx9 or ( + p.is_cuda() and p.has_device_capability((8, 9)) + ) + + supported: list[tuple[QuantKey | None, QuantKey | None]] = [(None, None)] + if device_supports_fp8: + supported += [ + (kFp8Static128BlockSym, kFp8Dynamic128Sym), + (kFp8StaticChannelSym, kFp8DynamicTokenSym), + (kFp8StaticTensorSym, kFp8DynamicTokenSym), + (kFp8StaticTensorSym, kFp8StaticTensorSym), + (kFp8StaticTensorSym, kFp8DynamicTensorSym), + ] + return (weight_key, activation_key) in supported + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in [ + MoEActivation.SILU, + MoEActivation.GELU, + MoEActivation.GELU_TANH, + MoEActivation.SWIGLUOAI, + MoEActivation.SILU_NO_MUL, + MoEActivation.GELU_NO_MUL, + MoEActivation.GELU_TANH_NO_MUL, + MoEActivation.RELU2_NO_MUL, + ] + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + return True + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + # Let PrepareAndFinalize::finalize() decide the impl. + return TopKWeightAndReduceDelegate() + + def activation( + self, activation: MoEActivation, output: torch.Tensor, input: torch.Tensor + ) -> None: + gemm1_clamp_limit = self.quant_config.gemm1_clamp_limit + if activation == MoEActivation.SILU and gemm1_clamp_limit is not None: + swiglu_limit_func(output, input, float(gemm1_clamp_limit)) + return + + super().activation(activation, output, input) + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + assert self.num_dispatchers is not None + assert self.max_num_tokens is not None + num_dp = self.num_dispatchers + num_experts = local_num_experts + max_num_tokens = self.max_num_tokens + activation_out_dim = self.adjust_N_for_activation(N, activation) + workspace13 = (num_experts, max_num_tokens * num_dp, max(K, N)) + workspace2 = (num_experts, max_num_tokens * num_dp, activation_out_dim) + output = (num_experts, max_num_tokens * num_dp, K) + return (workspace13, workspace2, output) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + # Check constraints. + if self.quant_config.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 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, + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + ] + assert expert_tokens_meta is not None + + expert_num_tokens = expert_tokens_meta.expert_num_tokens + + E, max_num_tokens, N, K, top_k_num = self.moe_problem_size( + hidden_states, w1, w2, topk_ids + ) + + assert w1.size(0) == E + assert w2.size(0) == E + + config_dtype = self.quant_config.config_name(hidden_states.dtype) + + config = try_get_optimal_moe_config( + w1.size(), + w2.size(), + top_k_num, + config_dtype, + max_num_tokens, + block_shape=self.block_shape, + ) + + 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 + elif hidden_states.dtype == current_platform.fp8_dtype(): + compute_type = tl.bfloat16 + else: + raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}") + + # We can reuse the memory between these because by the time we need + # cache3, we're done with cache1 + intermediate_cache1 = _resize_cache(workspace13, (E, max_num_tokens, N)) + activation_out_dim = self.adjust_N_for_activation(N, activation) + intermediate_cache2 = _resize_cache( + workspace2, (E, max_num_tokens, activation_out_dim) + ) + + # TODO(bnell): should this be done for any quantized type? + if self.quant_config.use_fp8_w8a8: + intermediate_cache1.fill_(0) + + a1q_scale = normalize_batched_scales_shape(a1q_scale, E) + + # MM1 + invoke_moe_batched_triton_kernel( + A=hidden_states, + B=w1, + C=intermediate_cache1, + expert_num_tokens=expert_num_tokens, + compute_type=compute_type, + A_scale=a1q_scale, + B_scale=self.w1_scale, + B_zp=self.w1_zp, + use_fp8_w8a8=self.quant_config.use_fp8_w8a8, + use_int8_w8a16=self.quant_config.use_int8_w8a16, + use_int4_w4a16=self.quant_config.use_int4_w4a16, + config=config, + per_act_token_quant=self.per_act_token_quant, + block_shape=self.block_shape, + ) + + intermediate_cache2.fill_(0) + + # TODO (bnell): use triton utility from batched deep gemm. + self.activation( + activation, + intermediate_cache2.view(-1, activation_out_dim), + intermediate_cache1.view(-1, N), + ) + + qintermediate_cache2, a2q_scale = batched_moe_kernel_quantize_input( + intermediate_cache2, + a2_scale, + max_num_tokens, + E, + N, + expert_num_tokens, + self.quant_dtype, + self.per_act_token_quant, + self.block_shape, + ) + + invoke_moe_batched_triton_kernel( + A=qintermediate_cache2, + B=w2, + C=output, + expert_num_tokens=expert_num_tokens, + compute_type=compute_type, + A_scale=a2q_scale, + B_scale=self.w2_scale, + B_zp=self.w2_zp, + use_fp8_w8a8=self.quant_config.use_fp8_w8a8, + use_int8_w8a16=self.quant_config.use_int8_w8a16, + use_int4_w4a16=self.quant_config.use_int4_w4a16, + config=config, + per_act_token_quant=self.per_act_token_quant, + block_shape=self.block_shape, + ) diff --git a/ex_engine/moe/fused_moe.py b/ex_engine/moe/fused_moe.py new file mode 100644 index 00000000..49957c8f --- /dev/null +++ b/ex_engine/moe/fused_moe.py @@ -0,0 +1,1740 @@ +# 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 diff --git a/ex_engine/moe/fused_moe_method_base.py b/ex_engine/moe/fused_moe_method_base.py new file mode 100644 index 00000000..888d064d --- /dev/null +++ b/ex_engine/moe/fused_moe_method_base.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import abstractmethod +from typing import TYPE_CHECKING + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.modular_kernel import ( + FusedMoEExpertsModular, + FusedMoEPrepareAndFinalizeModular, +) +from vllm.model_executor.layers.quantization.base_config import ( + QuantizeMethodBase, +) + +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + from vllm.model_executor.layers.fused_moe.runner.shared_experts import SharedExperts + +logger = init_logger(__name__) + + +class FusedMoEMethodBase(QuantizeMethodBase): + def __init__(self, moe: FusedMoEConfig): + super().__init__() + self.moe: FusedMoEConfig = moe + self.moe_quant_config: FusedMoEQuantConfig | None = None + self.moe_kernel: mk.FusedMoEKernel | None = None + + @property + def supports_internal_mk(self) -> bool: + # NOTE(rob): temporary attribute to indicate support for + # completed migration to the new internal MK interface. + return self.moe_kernel is not None + + @property + def mk_can_overlap_shared_experts(self) -> bool: + # NOTE(rob): temporary attribute to indicate support for + # completed migration to the new internal MK interface. + return ( + self.moe_kernel is not None and self.moe_kernel.can_overlap_shared_experts + ) + + @abstractmethod + def create_weights( + self, + layer: "RoutedExperts", + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + raise NotImplementedError + + def uses_weight_scale_2_pattern(self) -> bool: + """ + Returns True if this quantization method uses 'weight_scale_2' pattern + for per-tensor weight scales (e.g., FP4 variants), False otherwise. + + This method should be overridden by subclasses that use the + 'weight_scale_2' pattern instead of the standard 'weight_scale' pattern. + """ + return False + + def maybe_roundup_sizes( + self, + hidden_size: int, + intermediate_size_per_partition: int, + act_dtype: torch.dtype, + moe_parallel_config: FusedMoEParallelConfig, + ) -> tuple[int, int]: + """ + Given layer hidden size and intermediate size per partition and MoE + configurations, round up hidden_size and intermediate_size_per_partition + if necessary. + + Args: + hidden_size: Layer hidden-size + intermediate_size_per_partition: Intermediate size per partition for + the layer. + act_dtype: Data type of the layer activations. + moe_parallel_config: Fused MoE parallelization strategy configuration. + + Return: + A tuple of (rounded_hidden_size, rounded_intermediate_size_per_partition), + where: + - rounded_hidden_size is the possibly rounded up hidden size. + - rounded_intermediate_size_per_partition is the possibly rounded + up intermediate size per partition. + """ + from .all2all_utils import maybe_roundup_layer_hidden_size + + return maybe_roundup_layer_hidden_size( + hidden_size, act_dtype, moe_parallel_config + ), intermediate_size_per_partition + + def maybe_make_prepare_finalize( + self, + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + ) -> FusedMoEPrepareAndFinalizeModular | None: + from .all2all_utils import maybe_make_prepare_finalize + + pf = maybe_make_prepare_finalize( + self.moe, self.moe_quant_config, routing_tables + ) + assert pf is None or isinstance(pf, FusedMoEPrepareAndFinalizeModular) + return pf + + def select_gemm_impl( + self, + prepare_finalize: FusedMoEPrepareAndFinalizeModular, + layer: "RoutedExperts", + ) -> FusedMoEExpertsModular: + # based on the all2all implementation, select the appropriate + # gemm implementation + raise ValueError( + f"{self.__class__.__name__} uses the new modular kernel initialization " + "logic. This function should not be called." + ) + + @abstractmethod + def get_fused_moe_quant_config( + self, layer: "RoutedExperts" + ) -> FusedMoEQuantConfig | None: + raise NotImplementedError + + @property + def topk_indices_dtype(self) -> torch.dtype | None: + if self.moe_kernel is not None: + return self.moe_kernel.prepare_finalize.topk_indices_dtype() + return None + + @property + def skip_forward_padding(self) -> bool: + """Whether to skip the padding in the forward before applying the moe method.""" + return False + + @property + def has_unpadded_output(self) -> bool: + """ + Indicates that the hidden_states output might be the unpadded + hidden_states shape rather than the full padded shape. + """ + return False + + @property + def supports_eplb(self) -> bool: + return False + + @property + def method_name(self) -> str: + return self.__class__.__name__ + + @property + def is_monolithic(self) -> bool: + if self.moe_kernel is None: + if hasattr(self, "experts_cls"): + return self.experts_cls.is_monolithic() + else: + return False + return self.moe_kernel.is_monolithic + + def apply( + self, + layer: "RoutedExperts", + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts: "SharedExperts | None", + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + """ + Apply the MoE operation using modular kernels. + + Args: + layer: RoutedExperts instance containing weight parameters + x: Input tensor + topk_weights: Expert weights from router + topk_ids: Selected expert IDs from router + shared_experts_input: Input for shared experts (if any) + + Returns: + Output tensor from routed experts + """ + raise NotImplementedError + + def apply_monolithic( + self, + layer: "RoutedExperts", + x: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + """ + Apply the MoE operation using monolithic kernels. + + Args: + layer: RoutedExperts instance containing weight parameters + x: Input tensor + router_logits: Router logits (routing done internally) + + Returns: + Output tensor from routed experts + """ + raise NotImplementedError diff --git a/ex_engine/moe/fused_moe_modular_method.py b/ex_engine/moe/fused_moe_modular_method.py new file mode 100644 index 00000000..fb8e1793 --- /dev/null +++ b/ex_engine/moe/fused_moe_modular_method.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import TYPE_CHECKING + +import torch + +from vllm.logger import init_logger +from vllm.model_executor.custom_op import CustomOp +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.modular_kernel import ( + FusedMoEKernel, + FusedMoEPrepareAndFinalizeModular, +) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) + +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.routed_experts import ( + RoutedExperts, + ) + +logger = init_logger(__name__) + + +# --8<-- [start:modular_fused_moe] +@CustomOp.register("modular_fused_moe") +class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp): + # --8<-- [end:modular_fused_moe] + + def __init__( + self, old_quant_method: FusedMoEMethodBase, moe_kernel: FusedMoEKernel + ): + super().__init__(moe_kernel.moe_config) + self.moe_quant_config = old_quant_method.moe_quant_config + self.moe_kernel = moe_kernel + self.old_quant_method = old_quant_method + logger.debug("Swapping out %s", self.old_quant_method.__class__.__name__) + + @property + def wraps_legacy_quant_method(self) -> bool: + return not self.old_quant_method.supports_internal_mk + + @staticmethod + def make( + routed_experts: "RoutedExperts", + old_quant_method: FusedMoEMethodBase, + prepare_finalize: FusedMoEPrepareAndFinalizeModular, + ) -> "FusedMoEModularMethod": + return FusedMoEModularMethod( + old_quant_method, + FusedMoEKernel( + prepare_finalize, + old_quant_method.select_gemm_impl(prepare_finalize, routed_experts), + ), + ) + + @property + def skip_forward_padding(self) -> bool: + return self.old_quant_method.skip_forward_padding + + @property + def has_unpadded_output(self) -> bool: + return self.old_quant_method.has_unpadded_output + + @property + def supports_eplb(self) -> bool: + return self.old_quant_method.supports_eplb + + @property + def method_name(self) -> str: + return self.old_quant_method.method_name + + def create_weights( + self, + layer: "RoutedExperts", + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + raise NotImplementedError + + def get_fused_moe_quant_config( + self, layer: "RoutedExperts" + ) -> FusedMoEQuantConfig | None: + return self.moe_quant_config + + def apply( + self, + layer: "RoutedExperts", + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts: SharedExperts | None, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + assert self.moe_kernel is not None + return self.moe_kernel.apply( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + expert_map=layer.expert_map, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + ) diff --git a/ex_engine/moe/layer.py b/ex_engine/moe/layer.py new file mode 100644 index 00000000..15806ca4 --- /dev/null +++ b/ex_engine/moe/layer.py @@ -0,0 +1,406 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable +from typing import Any + +import torch + +from vllm._aiter_ops import rocm_aiter_ops +from vllm.config import ParallelConfig, get_current_vllm_config +from vllm.distributed import ( + get_dp_group, + get_pcp_group, + get_tensor_model_parallel_world_size, +) +from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, +) +from vllm.model_executor.layers.fused_moe.expert_map_manager import ( + ExpertMapManager, +) +from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( + FusedMoERouter, +) +from vllm.model_executor.layers.fused_moe.router.router_factory import ( + create_fused_moe_router, +) +from vllm.model_executor.layers.fused_moe.runner.moe_runner import ( + MoERunner, +) +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, +) + +logger = init_logger(__name__) + + +def make_parallel_config( + tp_size: int | None, + dp_size: int | None, + pcp_size: int | None, + is_sequence_parallel: bool, + parallel_config: ParallelConfig, +) -> FusedMoEParallelConfig: + tp_size_ = ( + tp_size if tp_size is not None else get_tensor_model_parallel_world_size() + ) + dp_size_ = dp_size if dp_size is not None else get_dp_group().world_size + pcp_size_ = pcp_size if pcp_size is not None else get_pcp_group().world_size + sp_size = tp_size_ if is_sequence_parallel else 1 + + moe_parallel_config = FusedMoEParallelConfig.make( + tp_size_=tp_size_, + pcp_size_=pcp_size_, + dp_size_=dp_size_, + sp_size_=sp_size, + vllm_parallel_config=parallel_config, + ) + + assert moe_parallel_config.is_sequence_parallel == is_sequence_parallel + + logger.debug("FusedMoEParallelConfig = %s", str(moe_parallel_config)) + + return moe_parallel_config + + +def determine_expert_counts( + num_experts: int, + num_redundant_experts: int, + n_shared_experts: int | None, + is_act_and_mul: bool, +) -> tuple[int, int, int]: + global_num_experts = num_experts + num_redundant_experts + logical_num_experts = num_experts + # ROCm aiter shared experts fusion + # AITER only supports gated activations (silu/gelu), so disable it + # for non-gated MoE (is_act_and_mul=False) + # rocm_aiter_fmoe_enabled = rocm_aiter_ops.is_fused_moe_enabled() and is_act_and_mul + aiter_fmoe_shared_expert_enabled = ( + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() and is_act_and_mul + ) + + num_fused_shared_experts = ( + n_shared_experts + if n_shared_experts is not None and aiter_fmoe_shared_expert_enabled + else 0 + ) + if not aiter_fmoe_shared_expert_enabled and num_fused_shared_experts != 0: + raise ValueError( + "n_shared_experts is only supported on ROCm aiter when " + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled" + ) + + return global_num_experts, logical_num_experts, num_fused_shared_experts + + +# TODO: rename this +def FusedMoE( + num_experts: int, # Global number of experts + top_k: int, + hidden_size: int, + intermediate_size: int, + params_dtype: torch.dtype | None = None, + renormalize: bool = True, + use_grouped_topk: bool = False, + num_expert_group: int | None = None, + topk_group: int | None = None, + quant_config: QuantizationConfig | None = None, + tp_size: int | None = None, + dp_size: int | None = None, + pcp_size: int | None = None, + prefix: str = "", + custom_routing_function: Callable | None = None, + router: FusedMoERouter | None = None, + scoring_func: str = "softmax", + routed_scaling_factor: float = 1.0, + swiglu_limit: float | None = None, + e_score_correction_bias: torch.Tensor | None = None, + apply_router_weight_on_input: bool = False, + activation: str = "silu", + enable_eplb: bool = False, + num_redundant_experts: int = 0, + has_bias: bool = False, + is_sequence_parallel: bool = False, + expert_mapping: list[tuple[str, str, int, str]] | None = None, + n_shared_experts: int | None = None, + router_logits_dtype: torch.dtype | None = None, + gate: torch.nn.Module | None = None, + shared_experts: torch.nn.Module | None = None, + shared_expert_gate: torch.nn.Module | None = None, + routed_input_transform: torch.nn.Module | None = None, + routed_output_transform: torch.nn.Module | None = None, + apply_routed_scale_to_output: bool = False, + zero_expert_type: str | None = None, + hash_indices_table: torch.Tensor | None = None, + runner_cls: type[MoERunner] | None = None, + runner_args: dict[str, Any] | None = None, + routed_experts_cls: type[RoutedExperts] | None = None, + routed_experts_args: dict[str, Any] | None = None, +) -> MoERunner: + """Factory function for creating MoE execution pipeline. + + Creates and configures a complete MoE execution pipeline including: + - Router (for token-to-expert assignment) + - RoutedExperts (containing expert weight parameters) + - MoERunner (orchestrates the complete forward pass) + + The experts contain both MergedColumnParallel weights (gate_up_proj/w13) + and RowParallelLinear weights (down_proj/w2). + + Note: Mixtral uses w1, w2, and w3 for gate, up, and down_proj. We + copy that naming convention here and handle any remapping in the + load_weights function in each model implementation. + + Args: + num_experts: Number of experts in the model (global count) + top_k: Number of experts selected for each token + hidden_size: Input hidden state size of the transformer + intermediate_size: Intermediate size of the experts + params_dtype: Data type for the parameters + renormalize: Whether to renormalize the logits in the router + use_grouped_topk: Whether to use grouped top-k routing + num_expert_group: Number of expert groups for grouped top-k + topk_group: Top-k value per group for grouped top-k + quant_config: Quantization configuration + tp_size: Tensor parallelism size (None = use global default) + dp_size: Data parallelism size (None = use global default) + pcp_size: Pipeline context parallelism size (None = use global default) + prefix: Layer name prefix for weight loading + custom_routing_function: Custom routing function override + router: Pre-configured router instance (None = create default) + scoring_func: Scoring function for routing ("softmax" or others) + routed_scaling_factor: Scaling factor applied to topk_weights or output + swiglu_limit: SwiGLU activation limit + e_score_correction_bias: Expert score correction bias tensor + apply_router_weight_on_input: Whether to apply router weights on input + activation: Activation function name ("silu", "gelu", etc.) + enable_eplb: Whether to enable expert parallelism load balancer + num_redundant_experts: Number of redundant experts for EPLB + has_bias: Whether expert layers have bias terms + is_sequence_parallel: Whether sequence parallelism is enabled + expert_mapping: Expert parameter mapping for weight loading + n_shared_experts: Number of shared experts (ROCm aiter only) + router_logits_dtype: Data type for router logits buffers + gate: Pre-configured gate module + shared_experts: Pre-configured shared experts module + shared_expert_gate: Pre-configured shared expert gate module + routed_input_transform: Input transformation module + routed_output_transform: Output transformation module + apply_routed_scale_to_output: Whether to apply routed_scaling_factor to + output instead of topk_weights + zero_expert_type: Type of zero expert handling + hash_indices_table: Hash table for expert indices + runner_cls: Custom MoERunner class (None = use default MoERunner) + runner_args: Additional arguments for runner constructor + routed_experts_cls: Custom RoutedExperts class (None = use default) + routed_experts_args: Additional arguments for routed_experts constructor + + Returns: + MoERunner: Configured MoE execution pipeline ready for forward passes + """ + vllm_config = get_current_vllm_config() + + layer_name = prefix + + moe_activation = MoEActivation.from_str(activation) + is_act_and_mul = moe_activation.is_gated + + moe_parallel_config = make_parallel_config( + tp_size=tp_size, + dp_size=dp_size, + pcp_size=pcp_size, + is_sequence_parallel=is_sequence_parallel, + parallel_config=vllm_config.parallel_config, + ) + + global_num_experts, logical_num_experts, num_fused_shared_experts = ( + determine_expert_counts( + num_experts, + num_redundant_experts, + n_shared_experts, + is_act_and_mul, + ) + ) + + # Initialize EPLB manager (or None?) + eplb_state: EplbLayerState | None = None + if enable_eplb: + use_ep = moe_parallel_config.use_ep + ep_size = moe_parallel_config.ep_size + if use_ep and global_num_experts % ep_size != 0: + raise ValueError( + f"EPLB currently only supports even distribution of " + f"experts across ranks. Got {global_num_experts} experts " + f"and {ep_size} EP ranks." + ) + eplb_state = EplbLayerState() + else: + assert num_redundant_experts == 0, ( + "Redundant experts are only supported with EPLB." + ) + + max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens + + # Create ExpertMapManager to handle expert mapping and placement for EP. + # See ExpertMapManager for a detailed description of what it does and when + # it is required. + expert_map_manager = ExpertMapManager( + max_num_batched_tokens=max_num_batched_tokens, + top_k=top_k, + global_num_experts=global_num_experts, + num_redundant_experts=num_redundant_experts, + num_expert_group=num_expert_group, + moe_parallel_config=moe_parallel_config, + placement_strategy=vllm_config.parallel_config.expert_placement_strategy, + enable_eplb=eplb_state is not None, + num_fused_shared_experts=num_fused_shared_experts, + rocm_aiter_enabled=rocm_aiter_ops.is_fused_moe_enabled() and is_act_and_mul, + ) + + # TODO(bnell): we should not have to create a router if the kernel is + # monolithic. + if router is None: + router = create_fused_moe_router( + top_k=top_k, + global_num_experts=global_num_experts, + eplb_state=eplb_state, + renormalize=renormalize, + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + custom_routing_function=custom_routing_function, + scoring_func=scoring_func, + # When apply_routed_scale_to_output is True, we set the scaling factor + # to 1.0 so it ends up being a nop. Applying the scale will be handled + # by the runner in this case. + # The member variable must be set in the same way as the router since + # some quantization methods can access it. + routed_scaling_factor=routed_scaling_factor + if not apply_routed_scale_to_output + else 1.0, + e_score_correction_bias=e_score_correction_bias, + num_fused_shared_experts=num_fused_shared_experts, + zero_expert_type=zero_expert_type, + num_logical_experts=logical_num_experts, + hash_indices_table=hash_indices_table, + ) + + if params_dtype is None: + params_dtype = torch.get_default_dtype() + + # FIXME (varun): We should have a better way of inferring the activation + # datatype. This works for now as the tensor datatype entering the MoE + # operation is typically unquantized (i.e. float16/bfloat16). + if vllm_config.model_config is not None: + moe_in_dtype = vllm_config.model_config.dtype + else: + # TODO (bnell): This is a hack to get test_mixtral_moe to work + # since model_config is not set in the pytest test. + moe_in_dtype = params_dtype + + moe_config = FusedMoEConfig( + num_experts=global_num_experts, + experts_per_token=top_k, + hidden_dim=hidden_size, + intermediate_size=intermediate_size, + num_local_experts=expert_map_manager.local_num_experts, + num_logical_experts=logical_num_experts, + moe_parallel_config=moe_parallel_config, + in_dtype=moe_in_dtype, + moe_backend=vllm_config.kernel_config.moe_backend, + router_logits_dtype=router_logits_dtype, + max_num_tokens=max_num_batched_tokens, + has_bias=has_bias, + is_lora_enabled=vllm_config.lora_config is not None, + activation=moe_activation, + device=vllm_config.device_config.device, + routing_method=router.routing_method_type, # Not ideal + swiglu_limit=swiglu_limit, + max_capture_size=vllm_config.compilation_config.max_cudagraph_capture_size, + ) + + logger.debug("FusedMoEConfig = %s", moe_config) + + # Create RoutedExperts instance BEFORE create_weights() + # This will hold all expert weight parameters + if routed_experts_cls is None: + routed_experts_cls = RoutedExperts + + assert params_dtype is not None + routed_experts = routed_experts_cls( + layer_name, + params_dtype, + moe_config, + quant_config, + expert_map_manager=expert_map_manager, + expert_mapping=expert_mapping, + # Extra params that are needed by quant_methods, pass along for now + # Prefer getting these from other sources, e.g. moe_config or + # router object + renormalize=renormalize, + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + custom_routing_function=custom_routing_function, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor + if not apply_routed_scale_to_output + else 1.0, + swiglu_limit=swiglu_limit, + # TODO get from router? needs to be truncated? + e_score_correction_bias=e_score_correction_bias, + apply_router_weight_on_input=apply_router_weight_on_input, + **routed_experts_args if routed_experts_args is not None else {}, + ) + + if runner_cls is None: + runner_cls = MoERunner + + runner = runner_cls( + layer_name=layer_name, + moe_config=moe_config, + router=router, + routed_experts=routed_experts, + enable_dbo=vllm_config.parallel_config.enable_dbo, + gate=gate, + shared_expert_gate=shared_expert_gate, + shared_experts=shared_experts, + routed_input_transform=routed_input_transform, + routed_output_transform=routed_output_transform, + # When apply_routed_scale_to_output is True, we allow + # the scaling factor to be passed to the runner, otherwise + # we pass 1.0 so it ends up being a nop. + routed_scaling_factor=routed_scaling_factor + if apply_routed_scale_to_output + else 1.0, + **runner_args if runner_args is not None else {}, + ) + + return runner + + +def fused_moe_make_expert_params_mapping( + model: torch.nn.Module, + ckpt_gate_proj_name: str, + ckpt_down_proj_name: str, + ckpt_up_proj_name: str, + num_experts: int, + num_redundant_experts: int = 0, + routed_experts_prefix: str = "routed_experts", +) -> list[tuple[str, str, int, str]]: + """Delegate to EPLB manager.""" + return RoutedExperts.make_expert_params_mapping( + model, + ckpt_gate_proj_name, + ckpt_down_proj_name, + ckpt_up_proj_name, + num_experts, + num_redundant_experts, + routed_experts_prefix, + ) diff --git a/ex_engine/moe/modular_kernel.py b/ex_engine/moe/modular_kernel.py new file mode 100644 index 00000000..d3176668 --- /dev/null +++ b/ex_engine/moe/modular_kernel.py @@ -0,0 +1,1630 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from abc import ABC, abstractmethod +from collections.abc import Callable +from dataclasses import dataclass +from enum import Enum +from math import prod +from typing import final + +import torch + +import vllm.envs as envs +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 ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, + SharedExpertsOrder, +) +from vllm.model_executor.layers.fused_moe.utils import ( + _resize_cache, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, +) +from vllm.platforms import current_platform +from vllm.v1.worker.ubatching import ( + dbo_enabled, + dbo_maybe_run_recv_hook, + dbo_register_recv_hook, + dbo_yield, +) +from vllm.v1.worker.workspace import current_workspace_manager + +logger = init_logger(__name__) + +# +# This file defines a set of base classes used to make MoE kernels more modular. +# The goal is to be able to utilize different communication mechanisms with +# any fused MoE kernel without needing to have combinatoric implementations. +# +# The fused moe kernels are broken down into the following components: +# +# [Router] → [Quantize-Dispatch] → [Permute-Experts-Unpermute] → [Combine] +# +# Each component will be independent of (but may inform) the others except for +# [Quantize-Dispatch] and `[Combine] (see below). The components can then be +# mixed and matched with so that DP+EP can be supported easily for multiple +# MoE kernel implementations. +# +# The following main classes are defined: +# * FusedMoEPrepareAndFinalizeModular - an abstract base class for preparation of MoE +# inputs (e.g. quantization, distribution) and finalization of Moe outputs. +# The prepare method must take care of any needed quantization and the +# finalize method, informed by the FusedMoEExpertsModular method, +# may apply weights and/or do the final reduction of the output. +# * FusedMoEExpertsModular - an abstract base class for the main fused +# MoE operation, i.e matmul + act_mul + optionally quant + matmul. +# Some FusedMoEExpertsModular implementations may choose to do +# the weight application and/or reduction. The class communicates this +# to [Finalize] via a TopKWeightAndReduce object. +# * FusedMoEModularKernel - an interface class that combines a +# FusedMoEPrepareAndFinalizeModular and a FusedMoEExpertsModular to +# provide the standard fused MoE kernel interface. +# * TopKWeightAndReduce - A TopKWeightAndReduce implementation chosen +# by the FusedMoEExpertsModular implementation that is passed +# on to [Finalize]. +# +# [Quantize-Prepare] and [Finalize] functionality are bundled into a single +# class `FusedMoEPrepareAndFinalizeModular` since they could use collective +# communication mechanisms that need to be consistent. +# + + +class FusedMoEActivationFormat(Enum): + """ + The standard activation format (num_tokens, hidden dim). + """ + + Standard = ("standard",) + """ + The batched experts format (num experts, max tokens per expert, hidden dim) + """ + BatchedExperts = ("batched_experts",) + + +@dataclass +class ExpertTokensMetadata: + """ + Metadata regarding expert-token routing. + """ + + expert_num_tokens: torch.Tensor + expert_num_tokens_cpu: torch.Tensor | None + + @staticmethod + def make_from_list( + expert_num_tokens_list: list[int], device: str + ) -> "ExpertTokensMetadata": + expert_num_tokens_cpu = torch.tensor( + expert_num_tokens_list, device="cpu", dtype=torch.int32 + ) + return ExpertTokensMetadata( + expert_num_tokens=expert_num_tokens_cpu.to(device, non_blocking=True), + expert_num_tokens_cpu=expert_num_tokens_cpu, + ) + + +class TopKWeightAndReduce(ABC): + """ + An abstract base class for weight application and reduction implementations. + """ + + @abstractmethod + def apply( + self, + output: torch.Tensor | None, + fused_expert_output: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + ) -> torch.Tensor: + """ + Apply topk_weights to the fused_experts_outputs and/or reduce. + If an output tensor is not passed, it will be created in the + function. + """ + raise NotImplementedError + + +# +# PrepareResultType is a tuple of: +# - quantized + dispatched a. +# - quantized + dispatched a1_scales. +# - Optional ExpertTokensMetadata containing gpu/cpu tensors +# as big as the number of local experts with the information about the +# number of tokens assigned to each local expert. +# - Optional dispatched expert topk IDs +# - Optional dispatched expert topk weight +# +# See `prepare` method below. +# +PrepareResultType = tuple[ + torch.Tensor, + torch.Tensor | None, + ExpertTokensMetadata | None, + torch.Tensor | None, + torch.Tensor | None, +] + +# +# PrepareResultType is a tuple of: +# - quantized + dispatched a. +# - quantized + dispatched a1_scales. +# - dispatched router logits. +# +# See `prepare_monolithic` method below. +# +PrepareMonolithicResultType = tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor, +] + +ReceiverType = Callable[[], PrepareResultType] + +################################################################################ +# Prepare/Finalize +################################################################################ + + +class FusedMoEPrepareAndFinalize(ABC): + """ + An abstract base class for the [Quantize-Prepare] and [Finalize] steps + described above. + + There are two variants of this class: + * FusedMoEPrepareAndFinalizeModular - this operates on topk ids and weights + * FusedMoEPrepareAndFinalizeMonolithic - the operates on router_logits + """ + + def post_init_setup(self, fused_experts: "FusedMoEExperts"): + """ + Initialize FusedMoEPrepareAndFinalizeModular settings that depend on + FusedMoEExpertsModular experts object. + The FusedMoEPrepareAndFinalizeModular implementations that have such + dependencies may choose to override this function. + """ + return + + @property + @abstractmethod + def activation_format(self) -> FusedMoEActivationFormat: + """ + A property indicating the output format of the activations for the + 'prepare' method. + """ + raise NotImplementedError + + @abstractmethod + def topk_indices_dtype(self) -> torch.dtype | None: + """ + The PrepareFinalize All2All implementations generally constrain the + dtype of the topk_ids they support. This function returns the + required topk indices dtype so it can be respected. + Return None if there are no such restrictions. + """ + raise NotImplementedError + + @abstractmethod + def max_num_tokens_per_rank(self) -> int | None: + """ + Some PrepareFinalize All2All implementations are batched. Meaning, + they can process only as set of tokens at a time. This + function returns the batch size i.e the maximum number of tokens + the implementation can process at a time. + Return None if there are no such restrictions. + """ + raise NotImplementedError + + @abstractmethod + def num_dispatchers(self) -> int: + raise NotImplementedError + + @abstractmethod + def output_is_reduced(self) -> bool: + """ + Indicates whether or not the output of finalize is reduced across all + ranks. + """ + raise NotImplementedError + + def supports_async(self) -> bool: + """ + Indicates whether or not this class implements prepare_async and + finalize_async. + """ + return False + + def on_commit(self) -> None: + """ + Runs after this prepare/finalize has been committed to the active + MoE kernel. + """ + return + + +# TODO: pass FusedMoEParallelConfig in as ctor parameter? +class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize): + """ + An abstract base class for the [Quantize-Prepare] and [Finalize] steps + described above for the Modular case. + """ + + @abstractmethod + def prepare( + self, + a1: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool, + ) -> PrepareResultType: + """ + Perform any quantization (and/or) dispatching needed for this kernel. + - a1: The (unquantized) input to the MoE layer. + - topk_ids: The topk ids. + - topk_weights: The topk weights. + - num_experts: The total number of experts in the global expert space. + - expert_map: A tensor mapping expert indices from the global expert + space to the local expert space of the expert parallel shard. + - apply_router_weight_on_input: When True, apply the weights to the + activations, before quantization + dispatching. + - quant_config: Quantization info provided by the fused experts. + - defer_input_quant: Runtime parameter indicating whether or not to + defer input quantization to the FusedMoEExpertsModular + in cases where the compute kernel expects unquantized inputs + + Returns a tuple of: + - quantized + dispatched a. + - Optional quantized + dispatched a1_scales. + - Optional ExpertTokensMetadata containing gpu/cpu tensors + as big as the number of local experts with the information about the + number of tokens assigned to each local expert. + - Optional dispatched expert topk IDs + - Optional dispatched expert topk weight + """ + raise NotImplementedError + + def prepare_async( + self, + a1: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool, + ) -> tuple[Callable, ReceiverType] | ReceiverType: + """ + Perform any quantization (and/or) dispatching needed for this kernel + but do not wait for results from other workers. + - a1: The (unquantized) input to the MoE layer. + - a1_scale: Optional scales for a1 + - a2_scale: Optional scales for the second MoE gemm. Required to make + sure the quantization is consistent for both gemms. + - topk_ids: The topk ids. + - topk_weights: The topk weights. + - num_experts: The total number of experts in the global expert space. + - expert_map: A tensor mapping expert indices from the global expert + space to the local expert space of the expert parallel shard. + - apply_router_weight_on_input: When True, apply the weights to the + activations, before quantization + dispatching. + - defer_input_quant: Runtime parameter indicating whether or not to + defer input quantization to the FusedMoEExpertsModular + in cases where the compute kernel expects unquantized inputs + + Returns a callback or a hook callback pair that when invoked waits for + results from other workers and has the same return signature as + `prepare`, if a hook is returned this is more lightweight check that + the recv is complete without doing extra work (used by DBO, will be + refactored in the very near future) + + e.g. + + ret = obj.prepare_async(...) + + if isinstance(ret, tuple): + hook, receiver = ret + hook() + + if hook is not None: + a, a_scales, expert_meta, topk_ids, topk_weights = receiver() + + is equivalent to: + + a, a_scales, expert_meta, topk_ids, topk_weights = obj.prepare(...) + """ + raise NotImplementedError + + @abstractmethod + def finalize( + self, + output: torch.Tensor, + fused_expert_output: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + weight_and_reduce_impl: TopKWeightAndReduce, + ) -> None: + """ + Perform any combine plus apply weights and perform a reduction on the + fused experts output. + - output: The output tensor, written in place. Must be (M, K) shape. + - fused_expert_output: The unweighted, unreduced output of the fused + experts, it will have (M, topk, K) shape. + - topk_weights: The weights to be applied to the fused_experts_output. + - topk_ids: The topk_ids. + - apply_router_weight_on_input: When False, apply the weights to + fused_expert_output. + - weight_and_reduce_impl: An optional TopKWeightAndReduce + implementation. + """ + raise NotImplementedError + + def finalize_async( + self, + output: torch.Tensor, + fused_expert_output: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + weight_and_reduce_impl: TopKWeightAndReduce, + ) -> tuple[Callable, Callable] | Callable: + """ + Perform any combine plus apply weights and perform a reduction on the + fused experts output but do not wait for results from other workers. + - output: The output tensor, written in place. Must be (M, K) shape. + - fused_expert_output: The unweighted, unreduced output of the fused + experts, it will have (M, topk, K) shape. + - topk_weights: The weights to be applied to the fused_experts_output. + - topk_ids: The topk_ids. + - apply_router_weight_on_input: When False, apply the weights to + fused_expert_output. + - weight_and_reduce_impl: An optional TopKWeightAndReduce + implementation. + + Returns a callback or a hook callback pair that when invoked waits for + results from other workers and has the same return signature as + `finalize`, if a hook is returned this is more lightweight check that + the recv is complete without doing extra work (used by DBO, will be + refactored in the very near future) + + ret = obj.finalize_async(output, ...) + ... output not valid yet ... + if isinstance(ret, tuple): + hook, receiver = ret + hook() + receiver() + ... output valid here ... + + is equivalent to: + + obj.finalize(output, ...) + """ + raise NotImplementedError + + +class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize): + """ + An abstract base class for the [Quantize-Prepare] and [Finalize] steps + described above for the monolithic case. + """ + + @abstractmethod + def prepare( + self, + a1: torch.Tensor, + router_logits: torch.Tensor, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, + ) -> PrepareMonolithicResultType: + """ + Optional method for subclasses compatible with monolithic + FusedMoEExpertsModular kernels. + + Perform any quantization (and/or) dispatching needed for this kernel. + - a1: The (unquantized) input to the MoE layer. + - quant_config: Quantization info provided by the fused experts. + - defer_input_quant: Runtime parameter indicating whether or not to + defer input quantization to the FusedMoEExpertsModular + + Returns a tuple of: + - quantized + dispatched a. + - Optional quantized + dispatched a1_scales. + """ + raise NotImplementedError + + @abstractmethod + def finalize(self, fused_expert_output: torch.Tensor) -> torch.Tensor: + """ + Optional method for subclasses compatible with monolithic + FusedMoEExpertsModular kernels. + + Perform any combine plus apply weights and perform a reduction on the + fused experts output. + - fused_expert_output: The unweighted, unreduced output of the fused + experts, it will have (M, topk, K) shape. + """ + raise NotImplementedError + + +################################################################################ +# Experts +################################################################################ + + +# TODO: add supported activations method (return string) +class FusedMoEExperts(ABC): + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + max_num_tokens: int | None = None, + num_dispatchers: int | None = None, + ): + """ + moe_config: MoE layer configuration. + quant_config: Quantization parameters for this experts instance. + """ + if self.activation_format() == FusedMoEActivationFormat.Standard and ( + max_num_tokens is not None or num_dispatchers is not None + ): + raise ValueError( + "max_num_tokens and num_dispatchers should only be set for " + "BatchedExperts activation format." + ) + elif self.activation_format() == FusedMoEActivationFormat.BatchedExperts and ( + max_num_tokens is None or num_dispatchers is None + ): + raise ValueError( + "max_num_tokens and num_dispatchers must be set for " + "BatchedExperts activation format." + ) + + self.moe_config = moe_config + self.quant_config = quant_config + self.max_num_tokens = max_num_tokens + self.num_dispatchers = num_dispatchers + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # noqa: B027 + pass + + @staticmethod + def is_monolithic() -> bool: + raise NotImplementedError("Implemented by subclasses.") + + @property + def expects_unquantized_inputs(self) -> bool: + """ + Whether or not the PrepareFinalize should defer input quantization + in the prepare step. If True, then the Experts kernel will + execute the input quantization itself. + + Sample subclasses that override are AITER and FlashInfer CUTLASS. + """ + return False + + @staticmethod + @abstractmethod + def activation_format() -> FusedMoEActivationFormat: + """ + A property which is a tuple of the input and output activation formats + for the 'apply' method. + """ + raise NotImplementedError + + # + # Various helpers for registering support for various features. + # Used by the oracle to select a particular kernel for a deployment. + # + + @staticmethod + def is_supported_config( + cls: type["FusedMoEExperts"], + moe_config: FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + activation_format: FusedMoEActivationFormat, + ) -> tuple[bool, str | None]: + def _make_reason(reason: str) -> str: + return f"kernel does not support {reason}" + + if not cls._supports_current_device(): + return False, _make_reason(f"current device {current_platform.device_name}") + elif not (moe_config.is_act_and_mul or cls._supports_no_act_and_mul()): + return False, _make_reason("no act_and_mul MLP layer") + elif not cls._supports_activation(moe_config.activation): + return False, _make_reason(f"{moe_config.activation} activation") + elif not cls._supports_quant_scheme(weight_key, activation_key): + return False, _make_reason( + f"quantization scheme {weight_key}x{activation_key}" + ) + elif not cls._supports_parallel_config(moe_config.moe_parallel_config): + return False, _make_reason( + f"parallel config {moe_config.moe_parallel_config}" + ) + elif not cls._supports_routing_method( + moe_config.routing_method, weight_key, activation_key + ): + return False, _make_reason(f"routing method {moe_config.routing_method}") + elif not cls._supports_router_logits_dtype( + moe_config.router_logits_dtype, + moe_config.routing_method, + ): + return False, _make_reason( + f"router logits dtype {moe_config.router_logits_dtype}" + ) + elif not cls._supports_shape(moe_config.hidden_dim): + return False, _make_reason( + f"{moe_config.hidden_dim} hidden dim is not supported" + ) + elif activation_format != cls.activation_format(): + return False, _make_reason(f"{activation_format.value} activation format") + elif envs.VLLM_BATCH_INVARIANT and not cls._supports_batch_invariance(): + return False, _make_reason("batch invariance") + elif moe_config.is_lora_enabled and not cls.supports_lora(): + return False, _make_reason("LoRA") + return True, None + + @staticmethod + @abstractmethod + def _supports_current_device() -> bool: + """ + Whether the kernel supports the current device type + (compute cability and current platform). + """ + raise NotImplementedError + + @staticmethod + @abstractmethod + def _supports_no_act_and_mul() -> bool: + """ + Whether the kernel supports act_and_mul=False, i.e. + non-gated MoE models like Nemotron-Nano. + """ + raise NotImplementedError + + @staticmethod + @abstractmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + raise NotImplementedError + + @staticmethod + @abstractmethod + def _supports_activation(activation: MoEActivation) -> bool: + """ + Whether the kernel supports a particular act function. + """ + raise NotImplementedError + + @staticmethod + @abstractmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + """ + Whether the kernel supports deployment in particular parallel config. + + Can be overridden if a kernel does not support EP, SP or some other + configuration. + """ + raise NotImplementedError + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """ + Whether the kernel supports a routing method (e.g. GroupedTopK). + + Can be overridden by monolithic kernels that execute the router + in addition to the experts if certain routers are not supported. + """ + return True + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + """ + Whether a kernel supports a particular dtype for router logits input. + + Can be overridden by monolithic kernels that execute the router + in addition to the experts if certain dtypes are not supported. + """ + return True + + @staticmethod + def _supports_shape(hidden_dim: int) -> bool: + """ + Whether a kernel supports a particular shape. Can be overridden if a kernel + has specific shape requirements. + """ + return True + + @staticmethod + def _supports_batch_invariance() -> bool: + """ + Whether the kernel supports batch invariance, i.e. the output does not + depend on the order of the tokens in the input batch. This is useful + for determining if the kernel can used with VLLM_BATCH_INVARIANT=1. + """ + return False + + # + # Various helpers for accessing quantization parameters from the + # quant_config. + # + + @property + def quant_dtype(self) -> torch.dtype | str | None: + return self.quant_config.quant_dtype + + @property + def weight_quant_dtype(self) -> torch.dtype | str | None: + return self.quant_config.weight_quant_dtype + + @property + def block_shape(self) -> list[int] | None: + return self.quant_config.block_shape + + @property + def per_act_token_quant(self) -> bool: + return self.quant_config.per_act_token_quant + + @property + def per_out_ch_quant(self) -> bool: + return self.quant_config.per_out_ch_quant + + @property + def a1_scale(self) -> torch.Tensor | None: + return self.quant_config.a1_scale + + @property + def a2_scale(self) -> torch.Tensor | None: + return self.quant_config.a2_scale + + @property + def a1_gscale(self) -> torch.Tensor | None: + return self.quant_config.a1_gscale + + @property + def a2_gscale(self) -> torch.Tensor | None: + return self.quant_config.a2_gscale + + @property + def w1_scale(self) -> torch.Tensor | None: + return self.quant_config.w1_scale + + @property + def w2_scale(self) -> torch.Tensor | None: + return self.quant_config.w2_scale + + @property + def w1_zp(self) -> torch.Tensor | None: + return self.quant_config.w1_zp + + @property + def w2_zp(self) -> torch.Tensor | None: + return self.quant_config.w2_zp + + @property + def w1_bias(self) -> torch.Tensor | None: + return self.quant_config.w1_bias + + @property + def w2_bias(self) -> torch.Tensor | None: + return self.quant_config.w2_bias + + @property + def g1_alphas(self) -> torch.Tensor | None: + return self.quant_config.g1_alphas + + @property + def g2_alphas(self) -> torch.Tensor | None: + return self.quant_config.g2_alphas + + @staticmethod + def supports_lora() -> bool: + """Return True if this expert impl natively handles LoRA. + + LoRA-aware experts should mix in LoRAExpertsMixin, which flips this + to True and provides the per-forward LoRA state plumbing. + """ + return False + + def supports_packed_ue8m0_act_scales(self) -> bool: + """ + A flag indicating whether or not this class can process packed ue8m0 + activation scales. + """ + return False + + +class FusedMoEExpertsModular(FusedMoEExperts): + """ + An abstract base class for the [Permute-Experts-Unpermute] step described + above. + """ + + @staticmethod + def is_monolithic() -> bool: + return False + + def moe_problem_size( + self, + a1: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + ) -> tuple[int, int, int, int, int]: + """ + Extract the MoE problem size from the given tensor arguments: + - a: The hidden states, input to the MoE layer. + - w1: The first set of expert weights. + - w2: The second set of expert weights. + - topk_ids: The topk ids. + + Note: extracting the problem shape from the weight and activation + tensors is not obvious. It needs to be done this way specifically + due to subtle issues with particular kernels, e.g. the int4 kernels + divide the trailing dimension by two, so it's not "correct" to + extract N or K from the trailing dimension of w1 or w2. Similarly, + some kernels transpose the weights, so this needs to be kept in mind. + + Note: This implementation covers most cases. However, if experts + require a specialized implementation, like MarlinExperts, they are free + to override this function. + """ + assert len(w1.shape) == 3 and len(w2.shape) == 3 + E, N, _ = w1.shape + K = a1.size(-1) + + if a1.dim() == 2: + # Make sure we are using the correct a1 (pre-permute). + assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}" + M = a1.size(0) + else: + assert a1.dim() == 3 + assert a1.size(0) == E, f"{a1.size(0)} == {E}" + M = a1.size(1) # This is max_num_tokens + + assert topk_ids.dim() == 2 + topk = topk_ids.size(1) + + return E, M, N, K, topk + + def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype: + """ + Workspace type: The dtype to use for the workspace tensors. + """ + return act_dtype + + @abstractmethod + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + """ + Compute the shapes for the temporary and final outputs of the two gemms + and activation in the fused expert function. Since the gemms are + independent, the workspace for the first gemm can be shared with the + workspace for the last gemm. + + Inputs: + - M: number of tokens. + - N: Row (or column) dimension of expert weights. + - K: hidden dimension + - topk: The number of top-k experts to select. + - global_num_experts: global number of experts. + - local_num_experts: local number of experts due to DP/EP. + - expert_tokens_meta: number of tokens per expert metadata for batched + format. + + Returns a tuple of: + - workspace13 shape tuple: must be large enough to hold the + result of either expert gemm. + - workspace2 shape tuple: must be large enough to hold the + result of the activation function. + - output shape tuple: must be exact size of the final gemm output. + - Note: workspace shapes can be 0 if the workspace is not needed. + But in order for activation chunking to work, the first dimension + of each tuple must be the number of tokens when the shape is + not 0. + """ + raise NotImplementedError + + @staticmethod + def adjust_N_for_activation(N: int, activation: MoEActivation) -> int: + """ + Calculate the output dimension for the activation function. + + For *_no_mul activations (e.g. relu2_no_mul), + there's no gate/up split, so output size equals input size (N). + + For regular gated activations (e.g., silu, gelu, swigluoai), + output size is N // 2 due to gate × activation(up) multiplication. + + Args: + N: The intermediate size (width of w1/w3 weights). + activation: The activation function enum. + + Returns: + The output dimension after activation. + """ + return N if not activation.is_gated else N // 2 + + def activation( + self, activation: MoEActivation, output: torch.Tensor, input: torch.Tensor + ) -> None: + apply_moe_activation(activation, output, input) + + @abstractmethod + def finalize_weight_and_reduce_impl(self) -> TopKWeightAndReduce: + raise NotImplementedError + + @abstractmethod + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ) -> None: + """ + This function computes the intermediate result of a Mixture of Experts + (MoE) layer using two sets of weights, w1 and w2. + + Parameters: + - output: (torch.Tensor): The unweighted, unreduced output tensor. + - hidden_states: (torch.Tensor): The (quantized) input tensor to the MoE + layer. + - w1 (torch.Tensor): The first set of expert weights. + - w2 (torch.Tensor): The second set of expert weights. + - topk_weights: A map of row to expert weights. Some implementations + choose to do weight application. + - topk_ids (torch.Tensor): A map of row to expert id. + - activation (str): The activation function to apply after the first + MoE layer. + - global_num_experts (int): The total number of experts in the global + expert space. + - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices + from the global expert space to the local expert space of the expert + parallel shard. + - a1q_scale (Optional[torch.Tensor]): Optional quantized scale to be + used for a1. Result of quantization from prepare/finalize and not + from the FusedMoEQuantConfig. + - workspace13 (torch.Tensor): A scratch tensor used for gemm outputs + must be large enough to hold output of either MoE gemm. + - workspace2 (torch.Tensor): A scratch tensor used for the activation + function. + - expert_tokens_meta (Optional[ExpertTokensMetadata]) - An optional + ExpertTokensMetadata object containing gpu/cpu tensors + as big as the number of local experts with the information about the + number of tokens assigned to each local expert. + - apply_router_weight_on_input: True if router weights are already + applied on the input. This is relevant if the implementation + chooses to do weight application. + """ + raise NotImplementedError + + +class FusedMoEExpertsMonolithic(FusedMoEExperts): + """ + An abstract base class for the [Permute-Experts-Unpermute] step described + above, but with the monolithic interface (accepts router logits + rather than topk ids and weights). + """ + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """ + Whether the kernel supports a routing method (e.g. GroupedTopK). + + Monolithic kernels should explicitly opt-in to support. + """ + raise NotImplementedError + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + """ + Whether the kernel supports a dtype for router logits. + + Modular kernels should opt-in to support. + """ + raise NotImplementedError + + @staticmethod + def is_monolithic() -> bool: + return True + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + """ + Same as apply(), except uses router_logits as opposed + to the topk_ids and topk_weights. This is useful for kernels + with fused router and fused_experts (e.g. FLASHINFER_TRTLLM). + """ + raise NotImplementedError + + +################################################################################ +# Kernel +################################################################################ + + +@final +class FusedMoEKernelModularImpl: + def __init__( + self, + prepare_finalize: FusedMoEPrepareAndFinalizeModular, + fused_experts: FusedMoEExpertsModular, + ): + self.prepare_finalize = prepare_finalize + self.fused_experts = fused_experts + moe_parallel_config = fused_experts.moe_config.moe_parallel_config + self.moe_parallel_config = moe_parallel_config + self.is_dp_ep = ( + moe_parallel_config is not None + and moe_parallel_config.dp_size > 1 + and moe_parallel_config.use_ep + ) + + def _allocate_buffers( + self, + out_dtype: torch.dtype, + device: torch.device, + M_chunk: int, + M_full: int, + N: int, + K: int, + top_k: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Allocate temporary and output buffers for the fused experts op. + Inputs: + - out_dtype: output type of workspace and output tensors. + - device: the device of the workspace and output tensors. + See `workspace_shapes` for a description of the remainder of arguments. + Returns a tuple of (workspace13, workspace2, output) tensors. + """ + assert M_full > 0 and M_chunk > 0 + + workspace_dtype = self.fused_experts.workspace_dtype(out_dtype) + + # Get intermediate workspace shapes based off the chunked M size. + workspace13_shape, workspace2_shape, _ = self.fused_experts.workspace_shapes( + M_chunk, + N, + K, + top_k, + global_num_experts, + local_num_experts, + expert_tokens_meta, + activation, + ) + + # Get final output shape based on the full M size. + _, _, fused_out_shape = self.fused_experts.workspace_shapes( + M_full, + N, + K, + top_k, + global_num_experts, + local_num_experts, + expert_tokens_meta, + activation, + ) + + # We can reuse the memory between cache1 and cache3 because by the + # time we need cache3, we're done with cache1. + # Reuse workspace13 for the output since there is only one chunk. + max_shape_size = max(prod(workspace13_shape), prod(fused_out_shape)) + common_workspace, workspace2 = current_workspace_manager().get_simultaneous( + ((max_shape_size,), workspace_dtype), + (workspace2_shape, workspace_dtype), + ) + workspace13 = _resize_cache(common_workspace, workspace13_shape) + fused_out = _resize_cache(common_workspace, fused_out_shape) + + return workspace13, workspace2, fused_out + + def _maybe_apply_shared_experts( + self, + shared_experts: SharedExperts | None, + shared_experts_input: torch.Tensor | None, + ): + if shared_experts is not None: + assert self.prepare_finalize.supports_async() + assert shared_experts_input is not None + shared_experts( + shared_experts_input, + SharedExpertsOrder.MK_INTERNAL_OVERLAPPED, + ) + + def _prepare( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + global_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + ExpertTokensMetadata | None, + torch.Tensor, + torch.Tensor, + ]: + """ + The _prepare method is a wrapper around self.prepare_finalize.prepare + that handles DBO and async. + """ + if not self.prepare_finalize.supports_async(): + # We shouldn't be running an a2a kernel that doesn't + # support async prepare/finalize + # TODO(lucas): enable in follow-up + assert not dbo_enabled() + + ( + a1q, + a1q_scale, + expert_tokens_meta, + _expert_topk_ids, + _expert_topk_weights, + ) = self.prepare_finalize.prepare( + hidden_states, + topk_weights, + topk_ids, + global_num_experts, + expert_map, + apply_router_weight_on_input, + self.fused_experts.quant_config, + defer_input_quant=self.fused_experts.expects_unquantized_inputs, + ) + else: + # Overlap shared expert compute with all2all dispatch. + dbo_maybe_run_recv_hook() + prepare_ret = self.prepare_finalize.prepare_async( + hidden_states, + topk_weights, + topk_ids, + global_num_experts, + expert_map, + apply_router_weight_on_input, + self.fused_experts.quant_config, + defer_input_quant=self.fused_experts.expects_unquantized_inputs, + ) + + # TODO(lucas): refactor this in the alternative schedules followup + # currently unpack if we have hook + receiver pair or just + # receiver (see finalize_async docstring) + hook, receiver = ( + prepare_ret if isinstance(prepare_ret, tuple) else (None, prepare_ret) + ) + + if hook is not None: + if dbo_enabled(): + # If DBO is being used, register the hook with the ubatch + # context and call it in dbo_maybe_run_recv_hook instead of + # passing it to the receiver. + dbo_register_recv_hook(hook) + dbo_yield() + else: + hook() + + ( + a1q, + a1q_scale, + expert_tokens_meta, + _expert_topk_ids, + _expert_topk_weights, + ) = receiver() + + # Maybe prepare gathered topk_ids and topk_weights from other EP ranks. + topk_ids = topk_ids if _expert_topk_ids is None else _expert_topk_ids + topk_weights = ( + topk_weights if _expert_topk_weights is None else _expert_topk_weights + ) + + return a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights + + def _fused_experts( + self, + in_dtype: torch.dtype, + a1q: torch.Tensor, + a1q_scale: torch.Tensor | None, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + local_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + expert_tokens_meta: ExpertTokensMetadata | None, + output_alias: torch.Tensor | None = None, + ) -> torch.Tensor: + _, M_full, N, K, top_k = self.fused_experts.moe_problem_size( + a1q, w1, w2, topk_ids + ) + + # This happens when none of the tokens from the all2all reach this + # EP rank. Also, note that this is only relevant for CUDAGraph + # incompatible all2all kernels like the DeepEP high-throughput + # kernels. CUDAGraph compatible all2all kernels like the DeepEP + # low-latency kernels are always batched and can never run into + # the tensor.numel() == 0 case. + if M_full == 0: + return torch.empty_like(a1q, dtype=in_dtype) + + workspace13, workspace2, fused_out = self._allocate_buffers( + in_dtype, + a1q.device, + M_full, + M_full, + N, + K, + top_k, + global_num_experts, + local_num_experts, + expert_tokens_meta, + activation, + ) + + # If caller's output buffer already matches fused_out shape/dtype, alias + # to skip the redundant copy in TopKWeightAndReduceNoOP.apply downstream. + # This eliminates ~94% of __amd_rocclr_copyBuffer events (Copy 2 of the + # double-copy MoE write-back path). + if current_platform.is_rocm(): + from vllm._aiter_ops import rocm_aiter_ops + + if ( + rocm_aiter_ops.is_fused_moe_enabled() + and output_alias is not None + and output_alias.shape == fused_out.shape + and output_alias.dtype == fused_out.dtype + and output_alias.device == fused_out.device + and output_alias.is_contiguous() + ): + fused_out = output_alias + + self.fused_experts.apply( + output=fused_out, + hidden_states=a1q, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + a1q_scale=a1q_scale, + a2_scale=self.fused_experts.a2_scale, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=expert_tokens_meta, + apply_router_weight_on_input=apply_router_weight_on_input, + ) + + return fused_out + + def _finalize( + self, + output: torch.Tensor, + fused_out: torch.Tensor, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + shared_experts: SharedExperts | None, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + """ + The _finalize method is a wrapper around self.prepare_finalize.finalize + that handles DBO, async and shared expert overlap. + + Args: + shared_experts: SharedExperts | None. The shared experts if any. + shared_experts_input: Optional separate input for shared experts. + When latent MoE is used, hidden_states is the latent-projected + tensor (smaller dimension) used by routed experts, while + shared_experts_input is the original hidden_states (full + dimension) needed by the shared expert MLP. + """ + if not self.prepare_finalize.supports_async(): + assert not dbo_enabled() + + self.prepare_finalize.finalize( + output, + fused_out, + topk_weights, + topk_ids, + apply_router_weight_on_input, + self.fused_experts.finalize_weight_and_reduce_impl(), + ) + else: + finalize_ret = self.prepare_finalize.finalize_async( + output, + fused_out, + topk_weights, + topk_ids, + apply_router_weight_on_input, + self.fused_experts.finalize_weight_and_reduce_impl(), + ) + self._maybe_apply_shared_experts(shared_experts, shared_experts_input) + + # TODO(lucas): refactor this in the alternative schedules followup + # currently unpack if we have hook + receiver pair or just + # receiver (see finalize_async docstring) + hook, receiver = ( + finalize_ret + if isinstance(finalize_ret, tuple) + else (None, finalize_ret) + ) + + if hook is not None: + if dbo_enabled(): + # If DBO is being used, register the hook with the ubatch + # context and call it in dbo_maybe_run_recv_hook instead of + # passing it to the receiver. + dbo_register_recv_hook(hook) + dbo_yield() + else: + hook() + + receiver() + + return output + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + activation: MoEActivation = MoEActivation.SILU, + global_num_experts: int = -1, + expert_map: torch.Tensor | None = None, + apply_router_weight_on_input: bool = False, + shared_experts: SharedExperts | None = None, + shared_experts_input: torch.Tensor | None = None, + ) -> torch.Tensor: + """ + This function computes a Mixture of Experts (MoE) layer using two sets + of weights, w1 and w2, and top-k gating mechanism. + + Parameters: + - hidden_states: (torch.Tensor): The input tensor to the MoE layer. + - w1 (torch.Tensor): The first set of expert weights. + - w2 (torch.Tensor): The second set of expert weights. + - topk_weights (torch.Tensor): The topk weights applied at the end of the layer. + - topk_ids (torch.Tensor): A map of row to expert id. + - activation (MoEActivation): The activation function to apply after the first + MoE layer. + - global_num_experts (int): The total number of experts in the global + expert space. + - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices + from the global expert space to the local expert space of the expert + parallel shard. + - apply_router_weight_on_input (bool): When true, the topk weights are + applied directly on the inputs. This is only applicable when topk is + 1. + - shared_experts: SharedExperts | None. The shared experts if any. + - shared_experts_input (Optional[torch.Tensor]): Optional separate + input for shared experts. For latent MoE, this is the original + hidden_states before latent projection. + + Returns: + - torch.Tensor: The output tensor after applying the MoE layer. + """ + output = torch.empty_like(hidden_states) + + local_num_experts = w1.shape[0] + if global_num_experts == -1: + global_num_experts = local_num_experts + + a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights = self._prepare( + hidden_states, + topk_weights, + topk_ids, + global_num_experts, + expert_map, + apply_router_weight_on_input, + ) + + fused_out = self._fused_experts( + in_dtype=hidden_states.dtype, + a1q=a1q, + a1q_scale=a1q_scale, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=global_num_experts, + local_num_experts=local_num_experts, + expert_map=expert_map, + apply_router_weight_on_input=apply_router_weight_on_input, + expert_tokens_meta=expert_tokens_meta, + output_alias=output, + ) + + return self._finalize( + output, + fused_out, + hidden_states, + topk_weights, + topk_ids, + apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + ) + + +@final +class FusedMoEKernelMonolithicImpl: + def __init__( + self, + prepare_finalize: FusedMoEPrepareAndFinalizeMonolithic, + fused_experts: FusedMoEExpertsMonolithic, + ): + self.prepare_finalize = prepare_finalize + self.fused_experts = fused_experts + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + """ + Same as forward(), except uses router_logits as opposed + to the topk_ids and topk_weights. This is used for kernels + that have fused router + experts (e.g. FLASHINFER_TRTLLM). + """ + + a1q, a1q_scale, router_logits = self.prepare_finalize.prepare( + hidden_states, + router_logits=router_logits, + quant_config=self.fused_experts.quant_config, + defer_input_quant=self.fused_experts.expects_unquantized_inputs, + ) + + fused_out = self.fused_experts.apply( + hidden_states=a1q, + w1=w1, + w2=w2, + router_logits=router_logits, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + apply_router_weight_on_input=apply_router_weight_on_input, + a1q_scale=a1q_scale, + # grouped topk + fused topk bias parameters + num_expert_group=num_expert_group, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=routed_scaling_factor, + topk_group=topk_group, + ) + + output = self.prepare_finalize.finalize(fused_out) + + return output + + +@final +class FusedMoEKernel: + def __init__( + self, + prepare_finalize: FusedMoEPrepareAndFinalize, + fused_experts: FusedMoEExperts, + ): + super().__init__() + + # Initialize the implementation (monolithic or modular). + self.impl: FusedMoEKernelModularImpl | FusedMoEKernelMonolithicImpl + if isinstance( + prepare_finalize, FusedMoEPrepareAndFinalizeModular + ) and isinstance(fused_experts, FusedMoEExpertsModular): + self.impl = FusedMoEKernelModularImpl( + prepare_finalize, + fused_experts, + ) + + elif isinstance( + prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic + ) and isinstance(fused_experts, FusedMoEExpertsMonolithic): + self.impl = FusedMoEKernelMonolithicImpl( + prepare_finalize, + fused_experts, + ) + + else: + raise ValueError( + "prepare_finalize and fused_experts must both be either monolithic " + f"or non-monolithic but got {prepare_finalize.__class__.__name__} " + f"and {fused_experts.__class__.__name__}" + ) + + self._post_init_setup() + + @property + def can_overlap_shared_experts(self) -> bool: + if isinstance(self.impl, FusedMoEKernelModularImpl): + return self.impl.prepare_finalize.supports_async() + else: + return False + + @property + def is_monolithic(self) -> bool: + return isinstance(self.impl, FusedMoEKernelMonolithicImpl) + + @property + def prepare_finalize(self) -> FusedMoEPrepareAndFinalize: + return self.impl.prepare_finalize + + @property + def fused_experts(self) -> FusedMoEExperts: + return self.impl.fused_experts + + @property + def moe_config(self) -> FusedMoEConfig: + return self.fused_experts.moe_config + + def supports_lora(self) -> bool: + return self.fused_experts.supports_lora() + + def _post_init_setup(self): + """ + Resolve any leftover setup dependencies between self.prepare_finalize + and self.fused_experts here. + """ + self.prepare_finalize.post_init_setup(self.impl.fused_experts) + assert ( + self.prepare_finalize.activation_format + == self.fused_experts.activation_format() + ) + + def output_is_reduced(self) -> bool: + """ + Indicates whether or not the output of fused MoE kernel + is reduced across all ranks. + """ + return self.prepare_finalize.output_is_reduced() + + def apply_monolithic( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + assert isinstance(self.impl, FusedMoEKernelMonolithicImpl) + return self.impl.apply( + hidden_states=hidden_states, + w1=w1, + w2=w2, + router_logits=router_logits, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + apply_router_weight_on_input=apply_router_weight_on_input, + num_expert_group=num_expert_group, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=routed_scaling_factor, + topk_group=topk_group, + ) + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + shared_experts: SharedExperts | None = None, + shared_experts_input: torch.Tensor | None = None, + ) -> torch.Tensor: + assert isinstance(self.impl, FusedMoEKernelModularImpl) + return self.impl.apply( + hidden_states=hidden_states, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + apply_router_weight_on_input=apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + ) diff --git a/ex_engine/moe/moe_align_block_size.py b/ex_engine/moe/moe_align_block_size.py new file mode 100644 index 00000000..7fc8bfcf --- /dev/null +++ b/ex_engine/moe/moe_align_block_size.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm import _custom_ops as ops +from vllm.triton_utils import triton +from vllm.utils.math_utils import round_up + + +def moe_align_block_size( + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, + expert_map: torch.Tensor | None = None, + pad_sorted_ids: bool = False, + ignore_invalid_experts: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Aligns the token distribution across experts to be compatible with block + size for matrix multiplication. + + Note: In the case of expert_parallel, moe_align_block_size initially + considers all experts as valid and aligns all tokens appropriately. + Before the function returns it marks the experts_ids that are not in + the current GPU rank as -1 so the MoE matmuls could skip those blocks. + This requires the num_experts input arg to be the num global experts. + + Parameters: + - topk_ids: A tensor of shape [total_tokens, top_k] representing the + top-k expert indices for each token. + - block_size: The block size used in block matrix multiplication. + - num_experts: The total number of experts. + - expert_map: A tensor of shape [num_experts] that maps the expert index + from the global space to the local index space of the current + expert parallel shard. If the expert is not in the current expert + parallel shard, the mapping is set to -1. + - pad_sorted_ids: A flag indicating whether the sorted_token_ids length + should be padded to a multiple of block_size, + - ignore_invalid_experts: A flag indicating whether to ignore invalid + experts. When False, all expert_ids in topk_ids will participate in + counting and ranking, but invalid experts in expert_ids will be marked + as -1. When True, all invalid expert_ids in topk_ids will be ignored + and will not participate in counting or ranking, and there will be no + -1 in expert_ids. + + Returns: + - sorted_token_ids: A tensor containing the sorted token indices according + to their allocated expert. + - expert_ids: A tensor indicating the assigned expert index for each block. + - num_tokens_post_padded: The total number of tokens after padding, + ensuring divisibility by block_size. + + This function pads the number of tokens that each expert needs to process + so that it is divisible by block_size. + Padding ensures that during block matrix multiplication, the dimensions + align correctly. + + Example: + Given topk_ids = [[2, 3, 4], [1, 2, 4], [1, 3, 4], [1, 2, 3]], + block_size = 4, and num_experts = 4: + - We initially have 12 tokens (after repeating 'top_k' times) and 4 experts, + with each expert needing to process 3 tokens. + - As block_size is 4, we pad 1 token for each expert. + - First, flatten topk_ids to [2, 3, 4, 1, 2, 4, 1, 3, 4, 1, 2, 3]. + - Then append padding tokens [12, 12, 12, 12] for each block. + - After sorting by expert index, we obtain token_ids + [3, 6, 9, 12, 0, 4, 10, 12, 1, 7, 11, 12, 2, 5, 8, 12]. + Tokens 12 are non-existent (padding) and are ignored in + the subsequent matrix multiplication. + - The padding ensures that the total number of tokens is now divisible + by block_size for proper block matrix operations. + """ + max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1) + if pad_sorted_ids: + max_num_tokens_padded = round_up(max_num_tokens_padded, block_size) + if topk_ids.numel() < num_experts: + max_num_tokens_padded = min( + topk_ids.numel() * block_size, max_num_tokens_padded + ) + sorted_ids = torch.empty( + (max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device + ) + max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size) + expert_ids = torch.empty( + (max_num_m_blocks,), dtype=torch.int32, device=topk_ids.device + ) + num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device=topk_ids.device) + + ops.moe_align_block_size( + topk_ids, + num_experts, + block_size, + sorted_ids, + expert_ids, + num_tokens_post_pad, + expert_map if ignore_invalid_experts else None, + ) + + if expert_map is not None and not ignore_invalid_experts: + expert_ids = expert_map[expert_ids] + + return sorted_ids, expert_ids, num_tokens_post_pad + + +def batched_moe_align_block_size( + max_tokens_per_batch: int, block_size: int, expert_num_tokens: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Given num_batches, max_tokens_per_batch, block_size and the number of + valid-tokens in each batch, prepare sorted_token_ids, expert_ids and + num_tokens_post_pad. sorted_token_ids, expert_ids and num_tokens_post_pad + have the same semantics as in moe_align_block_size. + + This function is intended to be a drop in replacement for + moe_align_batch_size for the batched case. + + Parameters: + - max_tokens_per_batch (int): Number of tokens in each batch (both + valid and invalid). + - block_size (int): block_size to align the data to. + - expert_num_tokens (torch.Tensor): expert_num_tokens[i], indicates + the number of valid tokens in batch i. + + Returns: + - sorted_token_ids (torch.Tensor): Torch tensor of size + (num_batches * max_tokens_per_batch) indicating the token indices for + that block. + - expert_ids (torch.Tensor): Torch tensor of size + ceil((num_batches * max_tokens_per_batch) / block_size) indicating + what expert to use for each block. + - num_tokens_post_pad (torch.Tensor): Torch tensor of size 1 + indicating the number of valid blocks with actual data to + process. This is represented in terms of num tokens. + Example: + Let num_batches=5, max_tokens_per_batch=8, block_size=4, and + expert_num_tokens=[2, 3, 0, 6, 8]. This expert_num_tokens tensor + indicates that, + - The first 2 tokens in the 0th batch are valid and the rest 6 are + invalid (i.e. in the 2D hidden_states tensor of shape, + [num_batches * max_tokens_per_batch, K], indices 0, 1 are valid) + - The first 3 tokens in the 1st batch are valid. i.e. indices 8, 9, 10 + - 0 tokens in the 2nd batch are valid + - first 6 tokens in the 3rd batch are valid. i.e. indices, + 24, 25, 26, 27, 28, 29 + - so on ... + + In this case, + sorted_token_ids will be [0, 1, 40, 40, + 8, 9, 10, 40, + 24, 25, 26, 27, + 28, 29, 40, 40, + 32, 33, 34, 35, + 36, 37, 38, 39, + 40, 40, 40, 40, + (rest all 40, 40, 40, 40) + ...] + Here, 40 represents an invalid index. as there is no token index 40. + The gemm kernel using this sorted_token_ids is expected to skip the + gemm computation when it encounters this invalid index. + + expert_ids will be [0, 1, 3, 3, 4, 5, 5, -1, -1, (rest all -1) ...] + Here, -1 represents an invalid expert. The gemm kernel using this + expert_ids is expected to skip the gemm computation when it encounters + an expert of id -1. + + num_tokens_post_pad will be 24 as sorted_token_ids has valid entries + until 24. + """ + + B = expert_num_tokens.size(0) + device = expert_num_tokens.device + + # Round up so each batch can be split to blocks evenly. + max_num_tokens_padded = B * round_up(max_tokens_per_batch, block_size) + + sorted_ids = torch.empty((max_num_tokens_padded,), dtype=torch.int32, device=device) + assert max_num_tokens_padded % block_size == 0 + max_num_m_blocks = max_num_tokens_padded // block_size + expert_ids = torch.empty((max_num_m_blocks,), dtype=torch.int32, device=device) + num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device=device) + + ops.batched_moe_align_block_size( + max_tokens_per_batch, + block_size, + expert_num_tokens, + sorted_ids, + expert_ids, + num_tokens_post_pad, + ) + + return sorted_ids, expert_ids, num_tokens_post_pad diff --git a/ex_engine/moe/moe_fused_mul_sum.py b/ex_engine/moe/moe_fused_mul_sum.py new file mode 100644 index 00000000..768f41db --- /dev/null +++ b/ex_engine/moe/moe_fused_mul_sum.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch +from torch._subclasses.fake_tensor import FakeTensor + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + + +@triton.jit +def moe_fused_mul_sum_kernel( + inputs_ptr, + topk_weights_ptr, + outputs_ptr, + top_ids_ptr, + expert_map_ptr, + num_tokens, + stride_m, + has_expert_map: tl.constexpr, + top_k: tl.constexpr, + size: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, +): + pid_k = tl.program_id(0) + pid_m = tl.program_id(1) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) + + m_mask = offs_m < num_tokens + k_mask = offs_k < size + mask = m_mask[:, None] & k_mask[None, :] + + a_base = inputs_ptr + (offs_m * stride_m)[:, None] + offs_k[None, :] + b_base = topk_weights_ptr + offs_m * top_k + + acc = tl.zeros((BLOCK_M, BLOCK_K), dtype=tl.float32) + + for n in tl.static_range(top_k): + b_val = tl.load(b_base + n, mask=m_mask, other=0.0).to(tl.float32) + if has_expert_map: + id_val = tl.load(top_ids_ptr + offs_m * top_k + n, mask=m_mask, other=0) + expert_mask = tl.load(expert_map_ptr + id_val) >= 0 + a_vec = tl.load( + a_base + n * size, + mask=mask & expert_mask[:, None], + other=0.0, + ).to(tl.float32) + else: + a_vec = tl.load( + a_base + n * size, + mask=mask, + other=0.0, + ).to(tl.float32) + acc += a_vec * b_val[:, None] + + out_ptrs = outputs_ptr + (offs_m * size)[:, None] + offs_k[None, :] + tl.store( + out_ptrs, + acc.to(outputs_ptr.dtype.element_ty), + mask=mask, + ) + + +def _heuristic_config( + num_tokens: int, + top_k: int, + size: int, + element_size: int, +): + is_fp32 = element_size > 2 + is_sm90_plus = current_platform.has_device_capability(90) + is_sm80_before = not current_platform.has_device_capability(80) + + if current_platform.has_device_capability(90): + # SM90/SM100+: prefer small tiles + many CTAs. + if is_fp32: + BLOCK_M = 1 if num_tokens <= 4 else 2 + else: + if num_tokens <= 4: + BLOCK_M = 1 + elif num_tokens <= 128: + BLOCK_M = 2 + else: + BLOCK_M = 4 + elif is_fp32: + if num_tokens <= 4: + BLOCK_M = 1 + elif num_tokens <= 32: + BLOCK_M = 2 + elif num_tokens <= 128: + BLOCK_M = 4 + else: + BLOCK_M = 4 + else: + if num_tokens <= 4: + BLOCK_M = 1 + elif num_tokens <= 32: + BLOCK_M = 2 + elif num_tokens <= 128: + BLOCK_M = 4 + elif num_tokens <= 1024: + BLOCK_M = 16 + else: + BLOCK_M = 8 + + if is_fp32: + max_block_k = 256 + elif is_sm80_before or is_sm90_plus: + max_block_k = 512 + else: + max_block_k = 1024 + BLOCK_K = min(triton.next_power_of_2(size), max_block_k) + BLOCK_K = max(BLOCK_K, 256) + + total = BLOCK_M * BLOCK_K + if is_fp32: + num_warps = max(8, min(16, total // 64)) + else: + num_warps = max(4, min(16, total // 256)) + + if is_sm80_before: + num_warps = min(num_warps, 8) + num_stages = 2 + elif is_sm90_plus: + num_warps = min(num_warps, 8) + num_stages = 4 if total <= 2048 else 2 + else: + num_stages = 4 if total <= 2048 else 2 + + return BLOCK_M, BLOCK_K, num_warps, num_stages + + +def moe_fused_mul_sum( + inputs: torch.Tensor, + topk_weights: torch.Tensor, + outputs: torch.Tensor | None = None, + topk_ids: torch.Tensor | None = None, + expert_map: torch.Tensor | None = None, +) -> torch.Tensor: + """ + Fused kernel for MoE (Mixture of Experts) to perform weighted summation + of expert outputs. + + Args: + inputs: The output from experts. + Shape: (num_tokens, top_k, hidden_size). + topk_weights: The weights assigned to each expert for each token. + Shape: (num_tokens, top_k). + outputs: Optional pre-allocated output tensor. + Shape: (num_tokens, hidden_size). + topk_ids: Optional indices of the top-k experts. Used when + `expert_map` is provided. Shape: (num_tokens, top_k). + expert_map: Optional mapping for Expert Parallelism. A value < 0 + indicates an invalid token/expert pair that will be skipped. + + Returns: + The fused weighted sum of expert outputs. + Shape: (num_tokens, hidden_size). + """ + assert inputs.ndim == 3 + assert topk_weights.ndim == 2 + assert inputs.is_contiguous() + assert topk_weights.is_contiguous() + assert inputs.dtype in (torch.float32, torch.float16, torch.bfloat16) + assert topk_weights.dtype in (torch.float32, torch.float16, torch.bfloat16) + + num_tokens, top_k, size = inputs.shape + output_shape = (num_tokens, size) + if outputs is None: + outputs = torch.empty(output_shape, dtype=inputs.dtype, device=inputs.device) + + assert outputs.shape == output_shape + assert topk_weights.shape == (num_tokens, top_k) + + if not isinstance(inputs, FakeTensor): + BLOCK_M, BLOCK_K, num_warps, num_stages = _heuristic_config( + num_tokens, + top_k, + size, + inputs.element_size(), + ) + grid = (triton.cdiv(size, BLOCK_K), triton.cdiv(num_tokens, BLOCK_M)) + moe_fused_mul_sum_kernel[grid]( + inputs, + topk_weights, + outputs, + topk_ids, + expert_map, + num_tokens, + top_k * size, + expert_map is not None, + top_k, + size, + BLOCK_M, + BLOCK_K, + num_warps=num_warps, + num_stages=num_stages, + ) + + return outputs diff --git a/ex_engine/moe/moe_permute_unpermute.py b/ex_engine/moe/moe_permute_unpermute.py new file mode 100644 index 00000000..ad9fb509 --- /dev/null +++ b/ex_engine/moe/moe_permute_unpermute.py @@ -0,0 +1,283 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass, field + +import torch + + +@dataclass +class MoEPermuteScratch: + # Reused metadata buffers for repeated grouped-MoE permutes. + max_num_tokens: int + topk: int + num_experts: int + num_local_experts: int + device: torch.device + hidden_size: int | None = None + hidden_dtype: torch.dtype | None = None + token_expert_indices: torch.Tensor = field(init=False) + expert_first_token_offset: torch.Tensor = field(init=False) + permuted_idx: torch.Tensor = field(init=False) + inv_permuted_idx: torch.Tensor = field(init=False) + permuted_hidden_states: torch.Tensor | None = field(init=False, default=None) + sort_workspace: torch.Tensor = field(init=False) + permuted_experts_id: torch.Tensor = field(init=False) + sorted_row_idx: torch.Tensor = field(init=False) + topk_ids_int32: torch.Tensor = field(init=False) + topk_ids_for_sort: torch.Tensor = field(init=False) + max_expanded_rows: int = field(init=False) + + def __post_init__(self) -> None: + assert self.max_num_tokens > 0 + assert self.topk > 0 + assert self.num_experts > 0 + assert self.num_local_experts > 0 + if self.hidden_size is None: + assert self.hidden_dtype is None + else: + assert self.hidden_dtype is not None + + self.max_expanded_rows = self.max_num_tokens * self.topk + self.token_expert_indices = torch.arange( + self.max_expanded_rows, dtype=torch.int32, device=self.device + ) + self.expert_first_token_offset = torch.empty( + self.num_local_experts + 1, dtype=torch.int64, device=self.device + ) + self.permuted_idx = torch.empty( + self.max_expanded_rows, dtype=torch.int32, device=self.device + ) + self.inv_permuted_idx = torch.empty( + self.max_expanded_rows, dtype=torch.int32, device=self.device + ) + if self.hidden_size is not None: + hidden_numel = self.max_expanded_rows * self.hidden_size + self.permuted_hidden_states = torch.empty( + hidden_numel, dtype=self.hidden_dtype, device=self.device + ) + self.permuted_experts_id = torch.empty( + self.max_expanded_rows, dtype=torch.int32, device=self.device + ) + self.sorted_row_idx = torch.empty( + self.max_expanded_rows, dtype=torch.int32, device=self.device + ) + self.topk_ids_int32 = torch.empty( + self.max_expanded_rows, dtype=torch.int32, device=self.device + ) + self.topk_ids_for_sort = torch.empty( + self.max_expanded_rows, dtype=torch.int32, device=self.device + ) + sorter_size = torch.ops._moe_C.moe_permute_sort_workspace_size( + self.max_expanded_rows, self.num_experts + ) + self.sort_workspace = torch.empty( + sorter_size, dtype=torch.int8, device=self.device + ) + # torch.device("cuda") in config, after initialized, + # will be changed to cuda:{index}, so we need to refresh here. + self.device = self.token_expert_indices.device + + def validate(self, hidden_states: torch.Tensor, topk_ids: torch.Tensor) -> None: + n_token, n_hidden = hidden_states.shape + assert hidden_states.device == self.device + assert topk_ids.device == self.device + assert n_token <= self.max_num_tokens + assert topk_ids.size(1) == self.topk + assert topk_ids.size(0) == n_token + if self.hidden_size is not None: + assert n_hidden == self.hidden_size + assert hidden_states.dtype == self.hidden_dtype + assert self.permuted_hidden_states is not None + + def token_expert_indices_view(self, n_token: int) -> torch.Tensor: + return self.token_expert_indices[: n_token * self.topk].view(n_token, self.topk) + + def prepare_topk_ids(self, topk_ids: torch.Tensor) -> torch.Tensor: + if topk_ids.dtype == torch.int32: + return topk_ids + numel = topk_ids.numel() + topk_ids_int32 = self.topk_ids_int32[:numel].view_as(topk_ids) + topk_ids_int32.copy_(topk_ids) + return topk_ids_int32 + + +def moe_permute( + hidden_states: torch.Tensor, + a1q_scale: torch.Tensor | None, + topk_ids: torch.Tensor, + n_expert: int, + n_local_expert: int = -1, + expert_map: torch.Tensor | None = None, + permuted_hidden_states: torch.Tensor | None = None, + scratch: MoEPermuteScratch | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + This function expands and permutes activation to gather uncontinuous tokens + for each expert. + Parameters: + - hidden_states (torch.Tensor): The input tensor to the MoE layer. + - a1q_scale (Optional[torch.Tensor]): quant scale for hidden_states + - topk_ids (torch.Tensor): topk expert route id for each token. + - n_expert (int): The number of expert. + - n_local_expert (int): The number of expert in current EP rank. + - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices + from the global expert space to the local expert space of the expert + parallel shard. + - permuted_hidden_states (Optional[torch.Tensor]): Optional output tensor. + If None, the output tensor will be created in this function. + Returns: + - permuted_hidden_states (torch.Tensor): permuted activation. + - a1q_scale (Optional[torch.Tensor]): permuted quant scale for hidden_states + if original scale not per-tensor scaling + - expert_first_token_offset (torch.Tensor): offset of the first token + of each expert for standard grouped gemm. + - inv_permuted_idx (torch.Tensor): idx map for moe_unpermute. + - permuted_idx (torch.Tensor): idx map from hidden to permuted_hidden. + """ + n_token, n_hidden = hidden_states.size() + topk = topk_ids.size(1) + assert (n_hidden * hidden_states.element_size()) % 16 == 0, ( + "permue kernel need hidden dim align to 16B" + ) + permuted_row_size = n_token * topk + if n_local_expert == -1: + n_local_expert = n_expert + if permuted_hidden_states is None: + if scratch is None: + permuted_hidden_states = torch.empty( + (permuted_row_size, n_hidden), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + else: + scratch.validate(hidden_states, topk_ids) + hidden_numel = permuted_row_size * n_hidden + scratch_hidden_states = scratch.permuted_hidden_states + assert scratch_hidden_states is not None + permuted_hidden_states = scratch_hidden_states[:hidden_numel].view( + permuted_row_size, n_hidden + ) + assert permuted_hidden_states.size() == (permuted_row_size, n_hidden), ( + f"Expected permuted hidden states to be {(permuted_row_size, n_hidden)}" + f" but got {permuted_hidden_states.size()}" + ) + + if scratch is None: + token_expert_indices = torch.arange( + 0, n_token * topk, dtype=torch.int32, device=hidden_states.device + ).reshape((n_token, topk)) + + expert_first_token_offset = torch.empty( + n_local_expert + 1, dtype=torch.int64, device=hidden_states.device + ) + permuted_idx = torch.full( + (permuted_row_size,), + n_token * topk, + dtype=torch.int32, + device=hidden_states.device, + ) + inv_permuted_idx = torch.empty( + (n_token, topk), dtype=torch.int32, device=hidden_states.device + ) + topk_ids_int32 = topk_ids.to(torch.int32) + torch.ops._moe_C.moe_permute( + hidden_states, + topk_ids_int32, + token_expert_indices, + expert_map, + n_expert, + n_local_expert, + topk, + permuted_hidden_states, + expert_first_token_offset, + inv_permuted_idx, + permuted_idx, + ) + else: + scratch.validate(hidden_states, topk_ids) + assert n_expert == scratch.num_experts + assert n_local_expert == scratch.num_local_experts + token_expert_indices = scratch.token_expert_indices_view(n_token) + expert_first_token_offset = scratch.expert_first_token_offset + permuted_idx = scratch.permuted_idx[:permuted_row_size] + permuted_idx.fill_(permuted_row_size) + inv_permuted_idx = scratch.inv_permuted_idx[:permuted_row_size].view( + n_token, topk + ) + permuted_experts_id = scratch.permuted_experts_id[:permuted_row_size].view( + n_token, topk + ) + sorted_row_idx = scratch.sorted_row_idx[:permuted_row_size].view(n_token, topk) + topk_ids_for_sort = scratch.topk_ids_for_sort[:permuted_row_size].view( + n_token, topk + ) + topk_ids_int32 = scratch.prepare_topk_ids(topk_ids) + torch.ops._moe_C.moe_permute_with_scratch( + hidden_states, + topk_ids_int32, + token_expert_indices, + expert_map, + n_expert, + n_local_expert, + topk, + permuted_hidden_states, + expert_first_token_offset, + inv_permuted_idx, + permuted_idx, + scratch.sort_workspace, + permuted_experts_id, + sorted_row_idx, + topk_ids_for_sort, + ) + + if a1q_scale is not None and a1q_scale.dim() > 1: + a1q_scale = a1q_scale[permuted_idx.clamp(max=n_token * topk - 1) // topk] + return ( + permuted_hidden_states, + a1q_scale, + expert_first_token_offset, + inv_permuted_idx.flatten(), + permuted_idx, + ) + + +def moe_unpermute( + out: torch.Tensor, + permuted_hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + inv_permuted_idx: torch.Tensor, + expert_first_token_offset: torch.Tensor | None = None, +) -> None: + """ + This function expands and permutes activation to gathering uncontinuous + tokens for each expert. + Parameters: + - out (torch.Tensor): output tensor + - permuted_hidden_states (torch.Tensor): permuted activation. + - topk_weights (torch.Tensor): topk expert route weight for each token. + - inv_permuted_idx (torch.Tensor): row idx map for moe_unpermute. + - expert_first_token_offset (Optional[torch.Tensor]): offset of the first + token of each expert for grouped gemm. + Returns: + - hidden_states (torch.Tensor): The reduced and unpermuted activation + tensor. + """ + topk = topk_weights.size(1) + n_hidden = permuted_hidden_states.size(-1) + assert (n_hidden * permuted_hidden_states.element_size()) % 16 == 0, ( + "unpermue kernel need hidden dim align to 16B" + ) + + torch.ops._moe_C.moe_unpermute( + permuted_hidden_states, + topk_weights, + inv_permuted_idx, + expert_first_token_offset, + topk, + out, + ) + + +def moe_permute_unpermute_supported(): + return torch.ops._moe_C.moe_permute_unpermute_supported() diff --git a/ex_engine/moe/prepare_finalize/__init__.py b/ex_engine/moe/prepare_finalize/__init__.py new file mode 100644 index 00000000..b3529c99 --- /dev/null +++ b/ex_engine/moe/prepare_finalize/__init__.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.model_executor.layers.fused_moe.prepare_finalize.batched import ( + BatchedPrepareAndFinalize, +) +from vllm.model_executor.layers.fused_moe.prepare_finalize.naive_dp_ep import ( + MoEPrepareAndFinalizeNaiveDPEPModular, + MoEPrepareAndFinalizeNaiveDPEPMonolithic, + make_moe_prepare_and_finalize_naive_dp_ep, +) +from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import ( + MoEPrepareAndFinalizeNoDPEPModular, + MoEPrepareAndFinalizeNoDPEPMonolithic, + make_moe_prepare_and_finalize_no_dp_ep, +) + +__all__ = [ + "BatchedPrepareAndFinalize", + "MoEPrepareAndFinalizeNaiveDPEPMonolithic", + "MoEPrepareAndFinalizeNaiveDPEPModular", + "make_moe_prepare_and_finalize_naive_dp_ep", + "MoEPrepareAndFinalizeNoDPEPMonolithic", + "MoEPrepareAndFinalizeNoDPEPModular", + "make_moe_prepare_and_finalize_no_dp_ep", + # deepep_ht, deepep_ll, and flashinfer_a2a are not + # imported here as they have optional dependencies (deep_ep, flashinfer). + # Import them directly from their modules as needed. +] diff --git a/ex_engine/moe/prepare_finalize/batched.py b/ex_engine/moe/prepare_finalize/batched.py new file mode 100644 index 00000000..94302771 --- /dev/null +++ b/ex_engine/moe/prepare_finalize/batched.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceDelegate, + TopKWeightAndReduceNaiveBatched, +) +from vllm.model_executor.layers.fused_moe.utils import ( + moe_kernel_quantize_input, + normalize_scales_shape, +) + + +class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): + """ + A reference prepare/finalize class that reorganizes the tokens into + expert batched format, i.e. E x max_num_tokens x K. This is the format + that the batched dispatch/combine kernels use. + """ + + def __init__( + self, + max_num_tokens: int, + num_local_experts: int, + num_dispatchers: int, + rank: int, + ): + super().__init__() + self.max_num_tokens = max_num_tokens + self.num_local_experts = num_local_experts + self.rank = rank + self.num_dispatchers_ = num_dispatchers + + @property + def activation_format(self) -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.BatchedExperts + + def max_num_tokens_per_rank(self) -> int | None: + return self.max_num_tokens + + def topk_indices_dtype(self) -> torch.dtype | None: + return None + + def num_dispatchers(self) -> int: + return self.num_dispatchers_ + + def output_is_reduced(self) -> bool: + return False + + def prepare( + self, + a1: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, + ) -> mk.PrepareResultType: + if defer_input_quant: + raise NotImplementedError( + f"{self.__class__.__name__} does not support defer_input_quant=True. " + "Please select an MoE kernel that accepts quantized inputs." + ) + assert a1.dim() == 2 + assert topk_ids.dim() == 2 + assert topk_ids.size(0) == a1.size(0) + + if apply_router_weight_on_input: + topk = topk_ids.size(1) + # TODO: this only works for topK=1, will need to update for topK>1 + assert topk == 1, ( + "apply_router_weight_on_input is only implemented for topk=1" + ) + a1.mul_(topk_weights.to(a1.dtype)) + + num_tokens, hidden_dim = a1.size() + topk = topk_ids.size(1) + + tokens_per_expert = torch.zeros(num_experts, dtype=torch.int, device=a1.device) + + num_local_experts = self.num_local_experts + + if quant_config.quant_dtype is None: + b_type = a1.dtype + else: + b_type = quant_config.quant_dtype + + b_a1 = torch.zeros( + (num_local_experts, self.max_num_tokens, hidden_dim), + dtype=b_type, + device=a1.device, + ) + + if quant_config.is_quantized: + scale_shape = quant_config.batched_scale_shape( + num_local_experts, self.max_num_tokens, hidden_dim + ) + + b_a1_scale = torch.empty(scale_shape, dtype=torch.float32, device=a1.device) + else: + assert quant_config.a1_scale is None + b_a1_scale = None + + first_expert = num_local_experts * self.rank + last_expert = first_expert + num_local_experts + + a1_scale = normalize_scales_shape(quant_config.a1_scale) + + for expert_id in range(first_expert, last_expert): + topks = torch.any(topk_ids == expert_id, dim=1).flatten() + rows = torch.count_nonzero(topks.flatten()) + if rows == 0: + continue + idx = expert_id - first_expert + tokens_per_expert[idx] = rows + rhs = a1[: topks.numel()][topks] + if quant_config.quant_dtype is not None: + if a1_scale is not None: + if quant_config.is_per_act_token: + rhs_a1_scale = a1_scale[: topks.numel()][topks] + else: + rhs_a1_scale = a1_scale + else: + rhs_a1_scale = None + b_a1[idx, :rows, :], b_s = moe_kernel_quantize_input( + rhs, + rhs_a1_scale, + quant_config.quant_dtype, + quant_config.per_act_token_quant, + quant_config.block_shape, + ) + assert b_s is not None + if quant_config.is_per_act_token: + b_a1_scale[idx, :rows] = b_s[:rows] + else: + b_a1_scale[idx, : b_s.shape[0]] = b_s + else: + b_a1[idx, :rows, :] = rhs + + assert b_a1_scale is None or b_a1_scale.ndim == 3 + + expert_tokens_meta = mk.ExpertTokensMetadata( + expert_num_tokens=tokens_per_expert, expert_num_tokens_cpu=None + ) + + return b_a1, b_a1_scale, expert_tokens_meta, None, None + + def finalize( + self, + output: torch.Tensor, + fused_expert_output: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + weight_and_reduce_impl: mk.TopKWeightAndReduce, + ) -> None: + if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate): + weight_and_reduce_impl = TopKWeightAndReduceNaiveBatched(self.rank) + weight_and_reduce_impl.apply( + output=output, + fused_expert_output=fused_expert_output, + topk_weights=topk_weights, + topk_ids=topk_ids, + apply_router_weight_on_input=apply_router_weight_on_input, + ) diff --git a/ex_engine/moe/prepare_finalize/no_dp_ep.py b/ex_engine/moe/prepare_finalize/no_dp_ep.py new file mode 100644 index 00000000..69587770 --- /dev/null +++ b/ex_engine/moe/prepare_finalize/no_dp_ep.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceContiguous, + TopKWeightAndReduceDelegate, +) +from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input + + +def _quantize_input( + a1: torch.Tensor, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None]: + # Defer input quant to moe kernel for backends (e.g. AITER, FI) + # which use a single kernel call for quant + experts. + if defer_input_quant: + return a1, None + + input_sf = ( + quant_config.a1_gscale if quant_config.use_nvfp4_w4a4 else quant_config.a1_scale + ) + a1q, a1q_scale = moe_kernel_quantize_input( + a1, + input_sf, + quant_dtype=quant_config.quant_dtype, + per_act_token_quant=quant_config.per_act_token_quant, + block_shape=quant_config.block_shape, + is_scale_swizzled=quant_config.is_scale_swizzled, + mx_alignment=quant_config.mx_alignment, + ) + + return a1q, a1q_scale + + +class MoEPrepareAndFinalizeNoDPEPModular(mk.FusedMoEPrepareAndFinalizeModular): + @property + def activation_format(self) -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def max_num_tokens_per_rank(self) -> int | None: + return None + + def topk_indices_dtype(self) -> torch.dtype | None: + return None + + def num_dispatchers(self) -> int: + return 1 + + def output_is_reduced(self) -> bool: + return False + + def prepare( + self, + a1: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, + ) -> mk.PrepareResultType: + if apply_router_weight_on_input: + topk = topk_ids.size(1) + # TODO: this only works for topK=1, will need to update for topK>1 + assert topk == 1, ( + "apply_router_weight_on_input is only implemented for topk=1" + ) + a1 = a1 * topk_weights.to(a1.dtype) + + a1q, a1q_scale = _quantize_input(a1, quant_config, defer_input_quant) + + return a1q, a1q_scale, None, None, None + + def finalize( + self, + output: torch.Tensor, + fused_expert_output: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + weight_and_reduce_impl: mk.TopKWeightAndReduce, + ) -> None: + if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate): + weight_and_reduce_impl = TopKWeightAndReduceContiguous() + weight_and_reduce_impl.apply( + output=output, + fused_expert_output=fused_expert_output, + topk_weights=topk_weights, + topk_ids=topk_ids, + apply_router_weight_on_input=apply_router_weight_on_input, + ) + + +class MoEPrepareAndFinalizeNoDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMonolithic): + @property + def activation_format(self) -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def max_num_tokens_per_rank(self) -> int | None: + return None + + def topk_indices_dtype(self) -> torch.dtype | None: + return None + + def num_dispatchers(self) -> int: + return 1 + + def output_is_reduced(self) -> bool: + return False + + def prepare( + self, + a1: torch.Tensor, + router_logits: torch.Tensor, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, + ) -> mk.PrepareMonolithicResultType: + a1q, a1q_scale = _quantize_input(a1, quant_config, defer_input_quant) + return a1q, a1q_scale, router_logits + + def finalize( + self, + fused_expert_output: torch.Tensor, + ) -> torch.Tensor: + return fused_expert_output + + +def make_moe_prepare_and_finalize_no_dp_ep( + use_monolithic: bool, +) -> MoEPrepareAndFinalizeNoDPEPModular | MoEPrepareAndFinalizeNoDPEPMonolithic: + return ( + MoEPrepareAndFinalizeNoDPEPMonolithic() + if use_monolithic + else MoEPrepareAndFinalizeNoDPEPModular() + ) diff --git a/ex_engine/moe/topk_weight_and_reduce.py b/ex_engine/moe/topk_weight_and_reduce.py new file mode 100644 index 00000000..837c1498 --- /dev/null +++ b/ex_engine/moe/topk_weight_and_reduce.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import torch + +import vllm._custom_ops as ops +import vllm.model_executor.layers.fused_moe.modular_kernel as mk + + +class TopKWeightAndReduceDelegate(mk.TopKWeightAndReduce): + """ + Useful in the case when some FusedMoEExpertsModular + implementation does not perform weight application and reduction + but cannot address the needs of all the compatible PrepareAndFinalize + implementations. + For example, BatchedTritonExperts is compatible with both batched + PrepareAndFinalize implementations like DeepEPLLPrepareAndFinalize and + BatchedPrepareAndFinalize. Some PrepareAndFinalize implementations do + the weight-application + reduction as part of the combine kernel, while + BatchedPrepareAndFinalize needs an explicit implementation. To facilitate + this case, the BatchedTritonExperts could use TopKWeightAndReduceDelegate + so the PrepareAndFinalize implementations could choose how to + weight + reduce. + """ + + def __eq__(self, other): + return isinstance(other, TopKWeightAndReduceDelegate) + + def apply( + self, + output: torch.Tensor | None, + fused_expert_output: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + ) -> torch.Tensor: + raise RuntimeError( + "The caller is expected to choose an appropriate " + "TopKWeightAndReduce implementation." + ) + + +class TopKWeightAndReduceNoOP(mk.TopKWeightAndReduce): + """ + The fused_experts outputs have already been weight applied and reduced. + This implementation is a no-op. + """ + + def __eq__(self, other): + return isinstance(other, TopKWeightAndReduceNoOP) + + def apply( + self, + output: torch.Tensor | None, + fused_expert_output: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + ) -> torch.Tensor: + # Weight application and reduction operations are already done. + if output is None: + return fused_expert_output + + # Skip self-copy when caller aliased fused_out to output upstream. + if output is fused_expert_output: + return output + + # MoEPrepareAndFinalizeNoDPEPModular needs the output to be in the `output` + # tensor. + assert output.size() == fused_expert_output.size(), ( + "output shape is expected to match the fused_expert_output shape. " + f"But got output={output.size()}, " + f"used_expert_output={fused_expert_output.size()}" + ) + output.copy_(fused_expert_output, non_blocking=True) + return output + + +class TopKWeightAndReduceContiguous(mk.TopKWeightAndReduce): + """ + TopKWeightAndReduce implementation for a fused_experts output + of shape (m, topk, K) + """ + + def __eq__(self, other): + return isinstance(other, TopKWeightAndReduceContiguous) + + def apply( + self, + output: torch.Tensor | None, + fused_expert_output: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + ) -> torch.Tensor: + m, num_topk = topk_ids.size() + k = fused_expert_output.size(-1) + if fused_expert_output.ndim == 2: + fused_expert_output = fused_expert_output.view(m, num_topk, k) + + assert fused_expert_output.size() == (m, num_topk, k), ( + f"Expected fused_expert_output size {(m, num_topk, k)}. But got " + f"{fused_expert_output.size()}" + ) + + if not apply_router_weight_on_input: + fused_expert_output.mul_(topk_weights.view(m, -1, 1)) + + if output is None: + output = torch.empty( + (m, k), + device=fused_expert_output.device, + dtype=fused_expert_output.dtype, + ) + assert output.size() == (m, k), ( + f"Expected output size {(m, k)}. But got {output.size()}" + ) + + ops.moe_sum(fused_expert_output, output) + return output + + +class TopKWeightAndReduceNaiveBatched(mk.TopKWeightAndReduce): + """ + TopKWeightAndReduce implementation for a fused_experts output + of shape (num_experts, batch_size, K) + """ + + def __init__(self, rank: int): + self.rank = rank + + def __eq__(self, other): + return isinstance(other, TopKWeightAndReduceNaiveBatched) and ( + other.rank == self.rank + ) + + def apply( + self, + output: torch.Tensor | None, + fused_expert_output: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + ) -> torch.Tensor: + assert fused_expert_output.ndim == 3 + num_tokens = topk_ids.size(0) + num_local_experts = fused_expert_output.size(0) + K = fused_expert_output.size(-1) + + if output is None: + output = torch.zeros( + (num_tokens, K), + device=fused_expert_output.device, + dtype=fused_expert_output.dtype, + ) + else: + output.fill_(0) + + assert output.size() == (num_tokens, K), ( + f"Expected output size {(num_tokens, K)}, but got {output.size()}" + ) + + first_expert = num_local_experts * self.rank + last_expert = first_expert + num_local_experts + + for expert_id in range(first_expert, last_expert): + matching_tokens = topk_ids == expert_id + topks = torch.any(matching_tokens, dim=1).flatten() + rows = torch.count_nonzero(topks) + rhs = fused_expert_output[expert_id - first_expert, :rows, :] + if not apply_router_weight_on_input: + rhs.mul_(topk_weights[matching_tokens].view(rhs.size(0), 1)) + output[topks] = output[topks] + rhs + + return output diff --git a/ex_engine/moe/utils.py b/ex_engine/moe/utils.py new file mode 100644 index 00000000..cb2cd5e9 --- /dev/null +++ b/ex_engine/moe/utils.py @@ -0,0 +1,441 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from math import prod + +import torch +import torch.nn.functional as F + +from vllm import _custom_ops as ops +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + per_token_group_quant_fp8, +) +from vllm.model_executor.layers.quantization.utils.int8_utils import ( + per_token_group_quant_int8, + per_token_quant_int8, +) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + quant_dequant_mxfp4, +) +from vllm.model_executor.layers.quantization.utils.mxfp6_utils import ( + quant_dequant_mxfp6, +) +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, +) +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( + ref_nvfp4_quant_dequant, +) +from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( + per_tensor_dequantize, +) +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import cdiv + + +@triton.jit +def _count_expert_num_tokens( + topk_ids_ptr, + expert_num_tokens_ptr, + num_experts, + topk_numel, + expert_map, + HAS_EXPERT_MAP: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + curr_expert = tl.program_id(0) + + offsets = tl.arange(0, BLOCK_SIZE) + topk_ids_ptrs = topk_ids_ptr + offsets + + acc = tl.zeros((BLOCK_SIZE,), dtype=tl.int32) + for x in range(tl.cdiv(topk_numel, BLOCK_SIZE)): + mask = offsets < (topk_numel - x * BLOCK_SIZE) + expert_ids = tl.load(topk_ids_ptrs, mask=mask, other=-1) + if HAS_EXPERT_MAP: + expert_map_ptrs = expert_map + expert_ids + expert_map_mask = expert_ids >= 0 + expert_ids = tl.load(expert_map_ptrs, mask=expert_map_mask, other=-1) + + has_curr_expert = tl.where(expert_ids == curr_expert, 1, 0) + acc = acc + has_curr_expert + topk_ids_ptrs += BLOCK_SIZE + + if curr_expert < num_experts: + tl.store(expert_num_tokens_ptr + curr_expert, tl.sum(acc)) + + +def count_expert_num_tokens( + topk_ids: torch.Tensor, num_local_experts: int, expert_map: torch.Tensor | None +) -> torch.Tensor: + """ + Count the number to tokens assigned to each expert. + + Parameters: + - topk_ids (torch.Tensor): Tensor mapping each token to its + list of experts. + - num_local_experts (int): Number of experts in this rank. + - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices + from the global expert space to the local expert space of the expert + parallel shard. + + Returns: + A tensor of size num_local_experts, where tensor[i] holds the number + of tokens assigned to the ith expert. + """ + assert topk_ids.dtype.is_signed, "The kernel uses -1 to represent invalid topk_ids" + expert_num_tokens = torch.empty( + (num_local_experts), device=topk_ids.device, dtype=torch.int32 + ) + + grid = num_local_experts + BLOCK_SIZE = min(topk_ids.numel(), 1024) + BLOCK_SIZE = triton.next_power_of_2(BLOCK_SIZE) + + _count_expert_num_tokens[(grid,)]( + topk_ids, + expert_num_tokens, + num_local_experts, + topk_ids.numel(), + expert_map, + HAS_EXPERT_MAP=expert_map is not None, + BLOCK_SIZE=BLOCK_SIZE, + ) + + return expert_num_tokens + + +def _resize_cache(x: torch.Tensor, v: tuple[int, ...]) -> torch.Tensor: + """ + Shrink the given tensor and apply the given view to it. This is + used to resize the intermediate fused_moe caches. + """ + assert prod(v) <= x.numel(), ( + f"{v} ({prod(v)}) <= {x.shape} ({x.numel()})" + ) # CUDAGRAPH unfriendly? + return x.flatten()[: prod(v)].view(*v) + + +def _nvfp4_quantize( + A: torch.Tensor, + A_scale: torch.Tensor | None, + is_sf_swizzled_layout: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + return ops.scaled_fp4_quant(A, A_scale, is_sf_swizzled_layout=is_sf_swizzled_layout) + + +def _fp8_quantize( + A: torch.Tensor, + A_scale: torch.Tensor | None, + per_act_token: bool, + block_shape: list[int] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Perform fp8 quantization on the inputs. If a block_shape + is provided, the output will be blocked. + """ + if block_shape is None: + # TODO(luka): use QuantFP8 custom op + # https://github.com/vllm-project/vllm/issues/20711 + A, A_scale = ops.scaled_fp8_quant( + A, A_scale, use_per_token_if_dynamic=per_act_token + ) + else: + assert not per_act_token + assert len(block_shape) == 2 + _, block_k = block_shape[0], block_shape[1] + A, A_scale = per_token_group_quant_fp8(A, block_k) + assert cdiv(A.size(-1), block_k) == A_scale.size(-1) + + return A, A_scale + + +def _int8_quantize( + A: torch.Tensor, + A_scale: torch.Tensor | None, + per_act_token: bool, + block_shape: list[int] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Perform int8 quantization on the inputs. If a block_shape + is provided, the output will be blocked. + """ + + # If weights are per-channel (per_channel_quant=True), then + # activations apply per-token quantization. Otherwise, assume + # activation tensor-wise fp8/int8 quantization, dynamic or static + if block_shape is None: + if per_act_token: + A, A_scale = per_token_quant_int8(A) + elif A_scale is not None: + # Static per-tensor: use the optimized CUDA kernel + A, A_scale, _ = ops.scaled_int8_quant(A, scale=A_scale) + elif A_scale is None: + # Dynamic per-tensor: compute scale then quantize via kernel + A_scale = torch.clamp(A.abs().max() / 127.0, min=1e-10) + A, A_scale, _ = ops.scaled_int8_quant(A, scale=A_scale) + else: + assert not per_act_token + assert len(block_shape) == 2 + _, block_k = block_shape[0], block_shape[1] + A, A_scale = per_token_group_quant_int8(A, block_k) + assert cdiv(A.size(-1), block_k) == A_scale.size(-1) + + return A, A_scale + + +def _mxfp4_quantize( + A: torch.Tensor, + A_scale: torch.Tensor | None, + per_act_token_quant: bool, + block_shape: list[int] | None = None, +) -> tuple[torch.Tensor, None]: + assert block_shape is None + # TODO: native mxfp4 is currently not integrated in vllm, + # so simulating even on devices supporting this data type natively. + # Once integrated, `current_platform.supports_mx()` should be used to + # control quantize+dequantize, or simply quantize here down to mxfp4. + A = quant_dequant_mxfp4(A) + + return A, None + + +def _mxfp8_e4m3_quantize( + A: torch.Tensor, + A_scale: torch.Tensor | None, + per_act_token_quant: bool, + block_shape: list[int] | None = None, + is_sf_swizzled_layout: bool = False, + mx_alignment: int = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + assert A_scale is None + assert not per_act_token_quant + assert block_shape is None or block_shape == [1, 32] + return mxfp8_e4m3_quantize(A, is_sf_swizzled_layout, mx_alignment) + + +def _mxfp6_e3m2_quantize( + A: torch.Tensor, + A_scale: torch.Tensor | None, + per_act_token_quant: bool, + block_shape: list[int] | None = None, +) -> tuple[torch.Tensor, None]: + assert block_shape is None + + # TODO: native mxfp6 is currently not integrated in vllm, + # so simulating even on devices supporting this data type natively. + # Eventually, there should be a check based on + # `current_platform.supports_mx()` here. + A = quant_dequant_mxfp6(A, quant_dtype="fp6_e3m2") + + return A, None + + +def _mxfp6_e2m3_quantize( + A: torch.Tensor, + A_scale: torch.Tensor | None, + per_act_token_quant: bool, + block_shape: list[int] | None = None, +) -> tuple[torch.Tensor, None]: + assert block_shape is None + + # TODO: native mxfp6 is currently not integrated in vllm, + # so simulating even on devices supporting this data type natively. + # Eventually, there should be a check based on + # `current_platform.supports_mx()` here. + A = quant_dequant_mxfp6(A, quant_dtype="fp6_e2m3") + + return A, None + + +def moe_kernel_quantize_input( + A: torch.Tensor, + A_scale: torch.Tensor | None, + quant_dtype: None | torch.dtype | str, + per_act_token_quant: bool, + block_shape: list[int] | None = None, + is_scale_swizzled: bool = True, + ocp_mx_scheme: str | None = None, + quantization_emulation: bool = False, + mx_alignment: int = 0, +) -> tuple[torch.Tensor, torch.Tensor | None]: + # Handle OCP MX scheme that requires QDQ (quantize-dequantize) for emulation + if ocp_mx_scheme is not None: + if ocp_mx_scheme in {"w_mxfp4", "w_mxfp4_a_mxfp4"}: + pass # No QDQ needed for these schemes + elif ocp_mx_scheme.endswith("a_fp8"): + # Perform QDQ (quantize and dequantize) on activation for emulation + # purpose, because there is no native kernel for weight in ocp_mx_scheme + # and activation in FP8. The implementation is based on existing + # non-emulation ops. + qA, qA_scale = ops.scaled_fp8_quant( + A, A_scale, use_per_token_if_dynamic=False + ) + A = per_tensor_dequantize(qA, qA_scale).to(A.dtype) + # After QDQ, we don't need further quantization + return A, None + # else: For other schemes (e.g., *_a_mxfp6_e3m2, *_a_mxfp6_e2m3), + # weights are already dequantized, and we proceed with normal + # activation quantization below. + + if quant_dtype == current_platform.fp8_dtype(): + if quantization_emulation: + raise NotImplementedError( + f"moe_kernel_quantize_input does not support quant_dtype={quant_dtype}" + " MOE quantization emulation. Please open an issue." + ) + return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape) + elif quant_dtype == torch.int8: + if quantization_emulation: + raise NotImplementedError( + "moe_kernel_quantize_input does not support quant_dtype=torch.int8" + " MOE quantization emulation. Please open an issue." + ) + return _int8_quantize(A, A_scale, per_act_token_quant, block_shape) + elif quant_dtype == "nvfp4": + if not quantization_emulation: + return _nvfp4_quantize(A, A_scale, is_sf_swizzled_layout=is_scale_swizzled) + else: + A = ref_nvfp4_quant_dequant(A, A_scale, block_size=16) + return A, None + elif quant_dtype == "mxfp4": + if not quantization_emulation: + raise NotImplementedError( + "moe_kernel_quantize_input should not be used for native" + " quant_dtype='mxfp4' MOE. Please open an issue." + ) + return _mxfp4_quantize(A, A_scale, per_act_token_quant, block_shape) + elif quant_dtype == "mxfp8": + # TODO: `quant_dtype == "mxfp8"` is ambiguous, + # should be fp8_e4m3. OCP MX also defines `fp8_e5m2`. + if quantization_emulation: + raise NotImplementedError( + "moe_kernel_quantize_input does not support quant_dtype='mxfp8' MOE " + "quantization emulation. Please open an issue." + ) + return _mxfp8_e4m3_quantize( + A, + A_scale, + per_act_token_quant, + block_shape, + is_sf_swizzled_layout=is_scale_swizzled, + mx_alignment=mx_alignment, + ) + elif quant_dtype == "mxfp6_e3m2": + if not quantization_emulation: + raise NotImplementedError( + "moe_kernel_quantize_input should not be used for native " + " quant_dtype='mxfp6_e3m2'MOE. Please open an issue." + ) + + return _mxfp6_e3m2_quantize(A, A_scale, per_act_token_quant, block_shape) + elif quant_dtype == "mxfp6_e2m3": + if not quantization_emulation: + raise NotImplementedError( + "moe_kernel_quantize_input should not be used for native" + " quant_dtype='mxfp6_e2m3' MOE. Please open an issue." + ) + + return _mxfp6_e2m3_quantize(A, A_scale, per_act_token_quant, block_shape) + else: + return A, A_scale + + +def normalize_scales_shape(scales: torch.Tensor | None) -> torch.Tensor | None: + if scales is not None: + if scales.numel() == 1: + scales = scales.view(1, 1) + else: + scales = scales.view(-1, scales.size(-1)) + return scales + + +def normalize_batched_scales_shape( + scales: torch.Tensor | None, + num_experts: int, +) -> torch.Tensor | None: + if scales is not None and scales.ndim < 3: + if scales.numel() == 1: + scales = scales.view(1) + scales = torch.repeat_interleave(scales, num_experts, dim=0).view( + num_experts, 1, 1 + ) + else: + scales = scales.view(num_experts, -1, scales.size(-1)) + + return scales + + +@triton.jit +def _pack_topk_ids_weights_kernel( + topk_ids_ptr, + topk_weights_ptr, + output_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, + USE_GDC: tl.constexpr, + launch_pdl: tl.constexpr, # triton metadata +): + pid = tl.program_id(axis=0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + if USE_GDC: + tl.extra.cuda.gdc_launch_dependents() + tl.extra.cuda.gdc_wait() + expert_id = tl.load(topk_ids_ptr + offsets, mask=mask, other=0).to(tl.int32) + expert_id_shifted = expert_id << 16 + + weight = tl.load(topk_weights_ptr + offsets, mask=mask, other=0.0) + weight_bf16 = weight.to(tl.bfloat16) + weight_int16 = weight_bf16.to(tl.int16, bitcast=True) + + weight_int32 = weight_int16.to(tl.int32) & 0xFFFF + + packed = expert_id_shifted | weight_int32 + tl.store(output_ptr + offsets, packed, mask=mask) + + +def trtllm_moe_pack_topk_ids_weights( + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + block_size: int = 1024, +) -> torch.Tensor: + assert topk_ids.shape == topk_weights.shape + assert topk_ids.is_contiguous() and topk_weights.is_contiguous() + + original_shape = topk_ids.shape + ids_flat = topk_ids.reshape(-1) + weights_flat = topk_weights.reshape(-1) + + n_elements = ids_flat.numel() + output = torch.empty(n_elements, dtype=torch.int32, device=topk_ids.device) + + use_gdc = current_platform.is_cuda() and current_platform.has_device_capability(90) + grid = (triton.cdiv(n_elements, block_size),) + _pack_topk_ids_weights_kernel[grid]( + ids_flat, + weights_flat, + output, + n_elements, + BLOCK_SIZE=block_size, + USE_GDC=use_gdc, + launch_pdl=use_gdc, + ) + return output.reshape(original_shape) + + +@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) +def swiglu_limit_func( + output: torch.Tensor, + input: torch.Tensor, # first half is gate, second half is up + swiglu_limit: float = 0.0, +) -> None: + d = input.shape[1] // 2 + gate = input[:, :d] + up = input[:, d:] + + if swiglu_limit > 0: + gate = torch.clamp(gate, max=swiglu_limit) + up = torch.clamp(up, min=-swiglu_limit, max=swiglu_limit) + + output.copy_(F.silu(gate) * up) diff --git a/ex_engine/xllm_kernels/ilu/activation.cpp b/ex_engine/xllm_kernels/ilu/activation.cpp index 1ad364a4..ae2a16ba 100644 --- a/ex_engine/xllm_kernels/ilu/activation.cpp +++ b/ex_engine/xllm_kernels/ilu/activation.cpp @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2025 The xLLM Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. diff --git a/ex_engine/xllm_kernels/ilu/attention.cpp b/ex_engine/xllm_kernels/ilu/attention.cpp index ad3cd295..aa257bf1 100644 --- a/ex_engine/xllm_kernels/ilu/attention.cpp +++ b/ex_engine/xllm_kernels/ilu/attention.cpp @@ -1,5 +1,5 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2025 The xLLM Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. diff --git a/ex_engine/xllm_kernels/ilu/fused_moe.cpp b/ex_engine/xllm_kernels/ilu/fused_moe.cpp index 21c15d8c..794f9bd9 100644 --- a/ex_engine/xllm_kernels/ilu/fused_moe.cpp +++ b/ex_engine/xllm_kernels/ilu/fused_moe.cpp @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2025 The xLLM Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. diff --git a/ex_engine/xllm_kernels/ilu/group_gemm.cpp b/ex_engine/xllm_kernels/ilu/group_gemm.cpp index 290299a0..38743e66 100644 --- a/ex_engine/xllm_kernels/ilu/group_gemm.cpp +++ b/ex_engine/xllm_kernels/ilu/group_gemm.cpp @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2026 The xLLM Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. diff --git a/ex_engine/xllm_kernels/ilu/ilu_ops_api.h b/ex_engine/xllm_kernels/ilu/ilu_ops_api.h index 3dedd7da..e4fd7853 100644 --- a/ex_engine/xllm_kernels/ilu/ilu_ops_api.h +++ b/ex_engine/xllm_kernels/ilu/ilu_ops_api.h @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2025 The xLLM Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. diff --git a/ex_engine/xllm_kernels/ilu/ixformer.h b/ex_engine/xllm_kernels/ilu/ixformer.h index 83bad88e..57ce66dc 100644 --- a/ex_engine/xllm_kernels/ilu/ixformer.h +++ b/ex_engine/xllm_kernels/ilu/ixformer.h @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2025 The xLLM Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. diff --git a/ex_engine/xllm_kernels/ilu/matmul.cpp b/ex_engine/xllm_kernels/ilu/matmul.cpp index f90c0d47..91b6868f 100644 --- a/ex_engine/xllm_kernels/ilu/matmul.cpp +++ b/ex_engine/xllm_kernels/ilu/matmul.cpp @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2025 The xLLM Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. diff --git a/ex_engine/xllm_kernels/ilu/norm.cpp b/ex_engine/xllm_kernels/ilu/norm.cpp index e451bc36..c5a98595 100644 --- a/ex_engine/xllm_kernels/ilu/norm.cpp +++ b/ex_engine/xllm_kernels/ilu/norm.cpp @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2025 The xLLM Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. diff --git a/ex_engine/xllm_kernels/ilu/rope.cpp b/ex_engine/xllm_kernels/ilu/rope.cpp index 45af7656..89370b79 100644 --- a/ex_engine/xllm_kernels/ilu/rope.cpp +++ b/ex_engine/xllm_kernels/ilu/rope.cpp @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2025 The xLLM Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. diff --git a/ex_engine/xllm_kernels/ilu/utils.h b/ex_engine/xllm_kernels/ilu/utils.h index 9fd15298..e8af0c3c 100644 --- a/ex_engine/xllm_kernels/ilu/utils.h +++ b/ex_engine/xllm_kernels/ilu/utils.h @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2025 The xLLM Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. diff --git a/ex_engine/xllm_layers/common/activation.cpp b/ex_engine/xllm_layers/common/activation.cpp new file mode 100644 index 00000000..83ba1451 --- /dev/null +++ b/ex_engine/xllm_layers/common/activation.cpp @@ -0,0 +1,38 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "activation.h" + +#include "kernels/ops_api.h" +namespace xllm { +namespace layer { + +ActivationImpl::ActivationImpl(const std::string& act_mode, bool is_gated) + : act_mode_(act_mode), is_gated_(is_gated) {} + +void ActivationImpl::forward(torch::Tensor& input, torch::Tensor& output) { + xllm::kernel::ActivationParams activation_params; + activation_params.input = input; + activation_params.output = output; + activation_params.act_mode = act_mode_; + activation_params.is_gated = is_gated_; + xllm::kernel::active(activation_params); + // Unified assignment: NPU returns new tensor, others modify in-place (no-op + // assignment) + output = activation_params.output; +} + +} // namespace layer +} // namespace xllm \ No newline at end of file diff --git a/ex_engine/xllm_layers/common/activation.h b/ex_engine/xllm_layers/common/activation.h new file mode 100644 index 00000000..981d97ed --- /dev/null +++ b/ex_engine/xllm_layers/common/activation.h @@ -0,0 +1,38 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include + +namespace xllm { +namespace layer { + +class ActivationImpl : public torch::nn::Module { + public: + ActivationImpl(const std::string& act_mode, bool is_gated); + + void forward(torch::Tensor& input, torch::Tensor& output); + + private: + std::string act_mode_; + bool is_gated_; +}; +TORCH_MODULE(Activation); + +} // namespace layer +} // namespace xllm \ No newline at end of file diff --git a/ex_engine/xllm_layers/common/dense_mlp.cpp b/ex_engine/xllm_layers/common/dense_mlp.cpp new file mode 100644 index 00000000..bb95dd0f --- /dev/null +++ b/ex_engine/xllm_layers/common/dense_mlp.cpp @@ -0,0 +1,141 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "dense_mlp.h" + +#include + +#include "kernels/ops_api.h" +#include "platform/device.h" + +namespace xllm { +namespace layer { + +DenseMLPImpl::DenseMLPImpl(int64_t hidden_size, + int64_t intermediate_size, + bool is_gated, + bool has_bias, + const std::string& hidden_act, + bool enable_result_reduction, + const QuantArgs& quant_args, + ProcessGroup* process_group, + const torch::TensorOptions& options, + const std::string& module_prefix) + : is_gated_(is_gated), + intermediate_size_(intermediate_size), + process_group_(process_group), + hidden_act_(hidden_act) { + // Check if using w8a8 smoothquant quantization + is_smoothquant_ = quant_args.quant_method() == kQuantMethodSmoothquant; + + if (is_smoothquant_) { + // Safety check: only w8a8 smoothquant is supported + if (quant_args.bits() != 8 || !quant_args.activation_dynamic()) { + LOG(FATAL) + << "DenseMLP w8a8 mode only supports w8a8 smoothquant quantization. " + << "Got bits=" << quant_args.bits() + << ", activation_dynamic=" << quant_args.activation_dynamic(); + } + } + + // Determine extra args based on quantization mode + LinearExtraArgs gate_up_proj_extra_args("none", false); + LinearExtraArgs down_proj_extra_args("none", false); + if (is_smoothquant_) { + // For per-token smoothquant, use specific args + down_proj_extra_args = LinearExtraArgs(hidden_act_, is_gated_); + } + + // 1. gate + up + int64_t out_feature = is_gated_ ? intermediate_size_ * 2 : intermediate_size_; + gate_up_proj_ = + register_module("gate_up_proj", + ColumnParallelLinear(hidden_size, + out_feature, + /*bias=*/has_bias, + /*gather_output=*/false, + quant_args, + process_group_, + options, + gate_up_proj_extra_args)); + + act_ = register_module("act", Activation(hidden_act_, is_gated_)); + + // 2. down + const auto down_proj_quant_args = + module_prefix.empty() + ? quant_args + : quant_args.for_module(module_prefix + ".down_proj"); + down_proj_ = register_module("down_proj", + RowParallelLinear(intermediate_size_, + hidden_size, + /*bias=*/has_bias, + /*input_is_parallelized=*/true, + enable_result_reduction, + down_proj_quant_args, + process_group_, + options, + down_proj_extra_args)); +} + +torch::Tensor DenseMLPImpl::forward(const torch::Tensor& hidden_states) { + // input shape: [num_tokens, hidden_size] + auto gate_up = gate_up_proj_->forward(hidden_states); + + if (is_smoothquant_) { + // For w8a8 quantization, the active operation is fused with the down_proj + return down_proj_->forward(gate_up); + } else { + torch::Tensor output; + if (Device::type_str() != "npu") { + int64_t batch_size = gate_up.sizes()[0]; + output = torch::empty( + {batch_size, intermediate_size_ / process_group_->world_size()}, + gate_up.options()); + } + + act_->forward(gate_up, output); + return down_proj_->forward(output); + } +} + +void DenseMLPImpl::load_state_dict(const StateDict& state_dict) { + gate_up_proj_->load_state_dict(state_dict, {"gate_proj.", "up_proj."}); + down_proj_->load_state_dict(state_dict.get_dict_with_prefix("down_proj.")); +} + +void DenseMLPImpl::load_state_dict(const StateDict& state_dict, + const std::vector& gate_up_name, + const std::string& down_name) { + if (is_gated_) { + CHECK_EQ(gate_up_name.size(), 2); + gate_up_proj_->load_state_dict(state_dict, gate_up_name); + } else { + CHECK_EQ(gate_up_name.size(), 1); + gate_up_proj_->load_state_dict( + state_dict.get_dict_with_prefix(gate_up_name[0])); + } + down_proj_->load_state_dict(state_dict.get_dict_with_prefix(down_name)); +} + +std::optional DenseMLPImpl::get_fp8_input_scale() const { + if (gate_up_proj_) { + return gate_up_proj_->get_input_scale(); + } + return std::nullopt; +} + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/xllm_layers/common/dense_mlp.h b/ex_engine/xllm_layers/common/dense_mlp.h new file mode 100644 index 00000000..8b4b2248 --- /dev/null +++ b/ex_engine/xllm_layers/common/dense_mlp.h @@ -0,0 +1,67 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include "activation.h" +#include "framework/model/model_args.h" +#include "framework/parallel_state/parallel_args.h" +#include "framework/quant_args.h" +#include "framework/state_dict/state_dict.h" +#include "linear.h" + +namespace xllm { +namespace layer { + +class DenseMLPImpl : public torch::nn::Module { + public: + DenseMLPImpl() = default; + DenseMLPImpl(int64_t hidden_size, + int64_t intermediate_size, + bool is_gated, + bool has_bias, + const std::string& hidden_act, + bool enable_result_reduction, + const QuantArgs& quant_args, + ProcessGroup* process_group, + const torch::TensorOptions& options, + const std::string& module_prefix = ""); + + torch::Tensor forward(const torch::Tensor& hidden_states); + + void load_state_dict(const StateDict& state_dict); + void load_state_dict(const StateDict& state_dict, + const std::vector& gate_up_name, + const std::string& down_name); + + // Get FP8 input scale from gate_up_proj for fused RMSNorm+FP8 quantization + std::optional get_fp8_input_scale() const; + + private: + bool is_gated_; + int64_t intermediate_size_; + ProcessGroup* process_group_; + ColumnParallelLinear gate_up_proj_{nullptr}; + RowParallelLinear down_proj_{nullptr}; + Activation act_{nullptr}; + bool is_smoothquant_; + std::string hidden_act_; +}; +TORCH_MODULE(DenseMLP); + +} // namespace layer +} // namespace xllm \ No newline at end of file diff --git a/ex_engine/xllm_layers/common/fused_moe.cpp b/ex_engine/xllm_layers/common/fused_moe.cpp new file mode 100644 index 00000000..b91dc08e --- /dev/null +++ b/ex_engine/xllm_layers/common/fused_moe.cpp @@ -0,0 +1,58 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "fused_moe.h" + +#include + +namespace xllm { +namespace layer { + +FusedMoEImpl::FusedMoEImpl(const ModelArgs& /*model_args*/, + const FusedMoEArgs& /*moe_args*/, + const QuantArgs& /*quant_args*/, + const ParallelArgs& /*parallel_args*/, + const torch::TensorOptions& /*options*/) { + NOT_IMPLEMENTED_WITH_MSG( + "FusedMoE is not supported for this backend. Please use CUDA, MLU or " + "ILU backend for MoE models."); +} + +torch::Tensor FusedMoEImpl::forward_experts( + const torch::Tensor& /*hidden_states*/, + const torch::Tensor& /*router_logits*/, + bool /*enable_all2all_communication*/) { + NOT_IMPLEMENTED_WITH_MSG( + "FusedMoE is not supported for this backend. Please use CUDA, MLU or " + "ILU backend for MoE models."); + return torch::Tensor(); +} + +torch::Tensor FusedMoEImpl::forward(const torch::Tensor& /*hidden_states*/, + const ModelInputParams& /*input_params*/) { + NOT_IMPLEMENTED_WITH_MSG( + "FusedMoE is not supported for this backend. Please use CUDA, MLU or " + "ILU backend for MoE models."); + return torch::Tensor(); +} + +void FusedMoEImpl::load_state_dict(const StateDict& /*state_dict*/) { + NOT_IMPLEMENTED_WITH_MSG( + "FusedMoE is not supported for this backend. Please use CUDA, MLU or " + "ILU backend for MoE models."); +} + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/xllm_layers/common/fused_moe.h b/ex_engine/xllm_layers/common/fused_moe.h new file mode 100644 index 00000000..6e148c15 --- /dev/null +++ b/ex_engine/xllm_layers/common/fused_moe.h @@ -0,0 +1,54 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include "dense_mlp.h" +#include "framework/model/model_args.h" +#include "framework/model/model_input_params.h" +#include "framework/parallel_state/parallel_args.h" +#include "framework/quant_args.h" +#include "framework/state_dict/state_dict.h" +#include "framework/state_dict/utils.h" +#include "fused_moe_base.h" +#include "linear.h" + +namespace xllm { +namespace layer { + +// FusedMoE common implementation - placeholder for unsupported backends +// Actual implementations are in backend-specific fused_moe.h files. +class FusedMoEImpl : public torch::nn::Module { + public: + FusedMoEImpl() = default; + FusedMoEImpl(const ModelArgs& model_args, + const FusedMoEArgs& moe_args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options); + + torch::Tensor forward_experts(const torch::Tensor& hidden_states, + const torch::Tensor& router_logits, + bool enable_all2all_communication); + torch::Tensor forward(const torch::Tensor& hidden_states, + const ModelInputParams& input_params); + void load_state_dict(const StateDict& state_dict); +}; +TORCH_MODULE(FusedMoE); + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/xllm_layers/common/rms_norm.cpp b/ex_engine/xllm_layers/common/rms_norm.cpp new file mode 100644 index 00000000..41947c14 --- /dev/null +++ b/ex_engine/xllm_layers/common/rms_norm.cpp @@ -0,0 +1,144 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "rms_norm.h" + +#include + +#include "kernels/ops_api.h" +#include "platform/device.h" + +namespace xllm { +namespace layer { + +const static std::string kLayerNormMode = "layernorm"; +const static std::string kRmsNormMode = "rmsnorm"; + +RMSNormImpl::RMSNormImpl(int64_t dim, + double eps, + const torch::TensorOptions& options) + : norm_dim_(dim), eps_(eps), mode_(kRmsNormMode) { + weight_ = register_parameter("weight", + torch::empty({dim}, options), + /*requires_grad=*/false); +} + +RMSNormImpl::RMSNormImpl(const ModelContext& context) + : RMSNormImpl(context.get_model_args().hidden_size(), + context.get_model_args().rms_norm_eps(), + context.get_tensor_options()) {} + +std::tuple> RMSNormImpl::forward( + torch::Tensor& input, + std::optional residual, + std::optional inplace_output) { + auto org_shape = input.sizes().vec(); + input = input.reshape({-1, norm_dim_}); + + torch::Tensor output; + if (Device::type_str() != "npu") { + if (inplace_output.has_value()) { + output = inplace_output.value(); + output = output.reshape({-1, norm_dim_}); + } else { + output = torch::empty_like(input); + } + } + + std::optional residual_out; + if (residual.has_value()) { + residual.value() = residual.value().reshape({-1, norm_dim_}); + if (Device::type_str() == "mlu" || Device::type_str() == "ilu") { + residual_out = residual.value(); + } + } + + xllm::kernel::FusedLayerNormParams fused_layernorm_params; + fused_layernorm_params.input = input; + fused_layernorm_params.residual = residual; + fused_layernorm_params.output = output; + fused_layernorm_params.residual_out = residual_out; + fused_layernorm_params.weight = weight_; + fused_layernorm_params.eps = eps_; + fused_layernorm_params.mode = mode_; + fused_layernorm_params.store_output_before_norm = residual_out.has_value(); + if (bias_.defined()) { + fused_layernorm_params.beta = bias_; + } + + xllm::kernel::fused_layernorm(fused_layernorm_params); + + output = fused_layernorm_params.output; + residual_out = fused_layernorm_params.residual_out; + + output = output.view(org_shape); + if (residual_out.has_value()) { + residual_out.value() = residual_out.value().view(org_shape); + } + return std::make_tuple(output, residual_out); +} + +std::tuple> +RMSNormImpl::forward_fp8(torch::Tensor& input, + const torch::Tensor& fp8_scale, + std::optional residual) { + // Only supported on CUDA for now + CHECK(Device::type_str() == "cuda") + << "forward_fp8 is only supported on CUDA"; + CHECK(mode_ == kRmsNormMode) + << "forward_fp8 only supports RMSNorm mode, not LayerNorm"; + + if (residual.has_value()) { + // Fused Add + RMSNorm + FP8 Quantization + xllm::kernel::FusedAddRmsNormStaticFp8QuantParams params; + params.input = input; + params.residual = residual.value(); + params.weight = weight_; + params.scale = fp8_scale; + params.epsilon = eps_; + + auto [output, updated_residual] = + xllm::kernel::fused_add_rms_norm_static_fp8_quant(params); + + return std::make_tuple(output, updated_residual); + } else { + // RMSNorm + FP8 Quantization (no residual) + xllm::kernel::RmsNormStaticFp8QuantParams params; + params.input = input; + params.weight = weight_; + params.scale = fp8_scale; + params.epsilon = eps_; + + auto output = xllm::kernel::rms_norm_static_fp8_quant(params); + + return std::make_tuple(output, std::nullopt); + } +} + +void RMSNormImpl::load_state_dict(const StateDict& state_dict) { + LOAD_WEIGHT(weight); + if (bias_.defined()) { + LOAD_WEIGHT(bias); + } +} + +void RMSNormImpl::set_layernorm_mode() { + mode_ = kLayerNormMode; + bias_ = register_parameter( + "bias", torch::empty({norm_dim_}, weight_.options()), false); +} + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/xllm_layers/common/rms_norm.h b/ex_engine/xllm_layers/common/rms_norm.h new file mode 100644 index 00000000..0c90c1c8 --- /dev/null +++ b/ex_engine/xllm_layers/common/rms_norm.h @@ -0,0 +1,64 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include "core/framework/model_context.h" +#include "framework/state_dict/state_dict.h" +#include "framework/state_dict/utils.h" + +namespace xllm { +namespace layer { + +class RMSNormImpl : public torch::nn::Module { + public: + RMSNormImpl(int64_t dim, double eps, const torch::TensorOptions& options); + RMSNormImpl(const ModelContext& context); + + // Standard forward: returns (normalized_output, updated_residual) + std::tuple> forward( + torch::Tensor& input, + std::optional residual = std::nullopt, + std::optional inplace_output = std::nullopt); + + // Fused forward with FP8 quantization output (for static quantization) + // Returns: (fp8_quantized_output, updated_residual) + // This combines RMSNorm + FP8 quantization to reduce memory bandwidth + std::tuple> forward_fp8( + torch::Tensor& input, + const torch::Tensor& fp8_scale, + std::optional residual = std::nullopt); + + void set_layernorm_mode(); + + void load_state_dict(const StateDict& state_dict); + + torch::Tensor weight() const { return weight_; } + torch::Tensor bias() const { return bias_; } + double eps() const { return eps_; } + + private: + DEFINE_WEIGHT(weight); + DEFINE_WEIGHT(bias); + int64_t norm_dim_; + double eps_; + std::string mode_; +}; +TORCH_MODULE(RMSNorm); + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/xllm_layers/common/rotary_embedding.cpp b/ex_engine/xllm_layers/common/rotary_embedding.cpp new file mode 100644 index 00000000..350dd14b --- /dev/null +++ b/ex_engine/xllm_layers/common/rotary_embedding.cpp @@ -0,0 +1,307 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "rotary_embedding.h" + +#include "kernels/ops_api.h" +#include "platform/device.h" + +namespace xllm { +namespace layer { + +RotaryEmbeddingImpl::RotaryEmbeddingImpl(const ModelContext& context) { + LOG(FATAL) << "Not implement currently."; +} + +RotaryEmbeddingImpl::RotaryEmbeddingImpl(int64_t rotary_dim, + int64_t max_position_embeddings, + int64_t rope_theta, + bool interleaved, + const torch::TensorOptions& options) + : interleaved_(interleaved) { + auto inv_freq = rotary::compute_inv_freq(rotary_dim, rope_theta, options); + const auto cos_sin = rotary::compute_cos_sin_cache( + rotary_dim, max_position_embeddings, interleaved, inv_freq, options); + cos_sin_cache_ = register_buffer("cos_sin_cache", cos_sin); + + auto cos_sin_vec = cos_sin_cache_.chunk(2, /*dim=*/-1); + cos_ = cos_sin_vec[0].view({-1, rotary_dim}); + sin_ = cos_sin_vec[1].view({-1, rotary_dim}); + + // Pre-compute [cos_half, sin_half] format used by the CUDA/ILU/MUSA kernels. + const auto dev = Device::type_str(); + if (dev == "cuda" || dev == "ilu" || dev == "musa") { + auto chunks = cos_sin_cache_.chunk(4, -1); + precomputed_cos_sin_cache_ = + torch::cat({chunks[0], chunks[2]}, -1).contiguous(); + } +} + +void RotaryEmbeddingImpl::forward(torch::Tensor& q, + torch::Tensor& k, + const torch::Tensor& positions, + const torch::Tensor& cu_query_lens, + int64_t max_query_len, + bool is_prompt) { + bool discrete; + std::optional position_ids; + if (is_prompt) { + discrete = false; + if (Device::type_str() == "cuda" || Device::type_str() == "npu" || + Device::type_str() == "ilu" || Device::type_str() == "musa") { + position_ids = positions; + } + } else { + discrete = true; + position_ids = positions; + } + + xllm::kernel::RotaryParams rotary_params; + rotary_params.q = q; + rotary_params.k = k; + rotary_params.sin = sin_; + rotary_params.cos = cos_; + rotary_params.cos_sin = cos_sin_cache_; + rotary_params.precomputed_cos_sin = precomputed_cos_sin_cache_; + rotary_params.position_ids = position_ids; + rotary_params.cu_query_lens = cu_query_lens; + rotary_params.interleaved = interleaved_; + rotary_params.discrete = discrete; + rotary_params.max_query_len = max_query_len; + xllm::kernel::apply_rotary(rotary_params); + + q = rotary_params.q; + k = rotary_params.k; +} + +// Single tensor forward for MLA architecture +void RotaryEmbeddingImpl::forward(torch::Tensor& input, + const torch::Tensor& positions, + const torch::Tensor& cu_query_lens, + int64_t max_query_len, + bool is_prompt) { + bool discrete; + std::optional position_ids; + if (is_prompt) { + discrete = false; + if (Device::type_str() == "cuda" || Device::type_str() == "npu" || + Device::type_str() == "ilu") { + position_ids = positions; + } + } else { + discrete = true; + position_ids = positions; + } + + xllm::kernel::RotaryParams rotary_params; + rotary_params.q = input; + rotary_params.sin = sin_; + rotary_params.cos = cos_; + rotary_params.cos_sin = cos_sin_cache_; + rotary_params.position_ids = position_ids; + rotary_params.cu_query_lens = cu_query_lens; + rotary_params.interleaved = interleaved_; + rotary_params.discrete = discrete; + rotary_params.max_query_len = max_query_len; + xllm::kernel::apply_rotary(rotary_params); + + input = rotary_params.q; +} + +MRotaryEmbeddingImpl::MRotaryEmbeddingImpl( + int64_t rotary_dim, + int64_t max_position_embeddings, + int64_t rope_theta, + bool interleaved, + const std::vector& rope_scaling_mrope_section, + const torch::TensorOptions& options) + : RotaryEmbeddingImpl(rotary_dim, + max_position_embeddings, + rope_theta, + interleaved, + options), + mrope_section_(rope_scaling_mrope_section) { + mrope_cu_seq_lens_ = torch::zeros(2, torch::kInt32).to(options.device()); +} + +void MRotaryEmbeddingImpl::forward(torch::Tensor& q, + torch::Tensor& k, + const torch::Tensor& positions, + const AttentionMetadata& attn_metadata) { + bool only_prefill = + (attn_metadata.is_prefill || attn_metadata.is_chunked_prefill); + if (!only_prefill || mrope_section_.empty()) { + torch::Tensor position_ids = positions; + if (positions.dim() == 2) { + position_ids = positions[0]; + } + return RotaryEmbeddingImpl::forward(q, + k, + position_ids, + attn_metadata.q_cu_seq_lens, + attn_metadata.max_query_len, + attn_metadata.is_prefill); + } + + int64_t num_tokens = positions.size(-1); + mrope_cu_seq_lens_[1] = num_tokens; + CHECK(attn_metadata.mrope_cos.defined() && attn_metadata.mrope_sin.defined()); + xllm::kernel::RotaryParams rotary_params; + rotary_params.q = q; + rotary_params.k = k; + rotary_params.sin = attn_metadata.mrope_sin; + rotary_params.cos = attn_metadata.mrope_cos; + rotary_params.cos_sin = cos_sin_cache_; + rotary_params.precomputed_cos_sin = precomputed_cos_sin_cache_; + rotary_params.position_ids = std::nullopt; + rotary_params.cu_query_lens = mrope_cu_seq_lens_; + rotary_params.interleaved = interleaved_; + rotary_params.discrete = false; + rotary_params.max_query_len = num_tokens; + xllm::kernel::apply_rotary(rotary_params); + + q = rotary_params.q; + k = rotary_params.k; +} + +DeepseekScalingRotaryEmbeddingImpl::DeepseekScalingRotaryEmbeddingImpl( + int64_t head_size, + int64_t rotary_dim, + int64_t max_position_embeddings, + int64_t rope_scaling_original_max_position_embeddings, + int64_t rope_theta, + bool interleaved, + float scaling_factor, + float extrapolation_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float mscale, + float mscale_all_dim, + const torch::TensorOptions& options) + : head_size_(head_size), + rotary_dim_(rotary_dim), + interleaved_(interleaved) { + auto inv_freq = rotary::apply_deepseek_yarn_rope_scaling( + scaling_factor, + extrapolation_factor, + beta_fast, + beta_slow, + rotary_dim, + rope_theta, + rope_scaling_original_max_position_embeddings); + const auto cos_sin = rotary::compute_cos_sin_cache(rotary_dim, + max_position_embeddings, + interleaved, + scaling_factor, + attn_factor, + mscale, + mscale_all_dim, + inv_freq, + options); + cos_sin_cache_ = register_buffer("cos_sin_cache", cos_sin); + + auto cos_sin_vec = cos_sin_cache_.chunk(2, /*dim=*/-1); + cos_ = cos_sin_vec[0].view({-1, rotary_dim}); + sin_ = cos_sin_vec[1].view({-1, rotary_dim}); + + // Pre-compute [cos_half, sin_half] format used by the CUDA/ILU/MUSA kernels. + const auto dev = Device::type_str(); + if (dev == "cuda" || dev == "ilu" || dev == "musa") { + auto chunks = cos_sin_cache_.chunk(4, -1); + precomputed_cos_sin_cache_ = + torch::cat({chunks[0], chunks[2]}, -1).contiguous(); + } +} + +void DeepseekScalingRotaryEmbeddingImpl::forward( + torch::Tensor& input, + const torch::Tensor& positions, + const torch::Tensor& cu_query_lens, + int64_t max_query_len, + bool is_prompt) { + const int32_t dim = -1; + bool discrete; + std::optional position_ids; + if (is_prompt) { + discrete = false; + position_ids = std::nullopt; + } else { + discrete = true; + position_ids = positions; + max_query_len = 1; + } + auto input_rot = input.slice(dim, 0, rotary_dim_); + torch::Tensor input_pass; + if (rotary_dim_ < head_size_) { + input_pass = input.slice(dim, rotary_dim_, head_size_); + } + + xllm::kernel::RotaryParams rotary_params; + rotary_params.q = input_rot; + rotary_params.sin = sin_; + rotary_params.cos = cos_; + rotary_params.cos_sin = cos_sin_cache_; + rotary_params.precomputed_cos_sin = precomputed_cos_sin_cache_; + rotary_params.position_ids = position_ids; + rotary_params.cu_query_lens = cu_query_lens; + rotary_params.interleaved = interleaved_; + rotary_params.discrete = discrete; + rotary_params.max_query_len = max_query_len; + xllm::kernel::apply_rotary(rotary_params); + input_rot = rotary_params.q; + + if (rotary_dim_ < head_size_) { + input = torch::cat({input_rot, input_pass}, dim); + } else { + input = input_rot; + } +} + +// Factory function: creates the appropriate RoPE type based on model args +std::shared_ptr create_mla_rotary_embedding( + const ModelArgs& args, + int64_t rotary_dim, + int64_t max_position_embeddings, + bool interleaved, + const torch::TensorOptions& options) { + if (args.rope_scaling_rope_type() == "deepseek_yarn") { + return std::make_shared( + rotary_dim, // head_size (same as rotary_dim for MLA) + rotary_dim, + max_position_embeddings, + args.rope_scaling_original_max_position_embeddings(), + args.rope_theta(), + interleaved, + args.rope_scaling_factor(), + args.rope_extrapolation_factor(), + args.rope_scaling_attn_factor(), + args.rope_scaling_beta_fast(), + args.rope_scaling_beta_slow(), + args.rope_scaling_mscale(), + args.rope_scaling_mscale_all_dim(), + options); + } else { + // default rope type + return std::make_shared(rotary_dim, + max_position_embeddings, + args.rope_theta(), + interleaved, + options); + } +} + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/xllm_layers/common/rotary_embedding.h b/ex_engine/xllm_layers/common/rotary_embedding.h new file mode 100644 index 00000000..fa72124d --- /dev/null +++ b/ex_engine/xllm_layers/common/rotary_embedding.h @@ -0,0 +1,158 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include + +#include + +#include "attention_metadata.h" +#include "core/framework/model_context.h" +#include "framework/model/model_args.h" +#include "rotary_embedding_util.h" + +namespace xllm { +namespace layer { + +class RotaryEmbeddingBase : public torch::nn::Module { + public: + ~RotaryEmbeddingBase() override = default; + + virtual void forward(torch::Tensor& input, + const torch::Tensor& positions, + const torch::Tensor& cu_query_lens, + int64_t max_query_len, + bool is_prompt) = 0; + virtual const torch::Tensor& get_sin_cache() const = 0; + virtual const torch::Tensor& get_cos_cache() const = 0; + virtual const bool get_interleaved() const = 0; +}; + +class RotaryEmbeddingImpl : public RotaryEmbeddingBase { + public: + RotaryEmbeddingImpl(int64_t rotary_dim, + int64_t max_position_embeddings, + int64_t rope_theta, + bool interleaved, + const torch::TensorOptions& options); + RotaryEmbeddingImpl(const ModelContext& context); + + void forward(torch::Tensor& q, + torch::Tensor& k, + const torch::Tensor& positions, + const torch::Tensor& cu_query_lens, + int64_t max_query_len, + bool is_prompt); + // Single tensor forward for MLA architecture + void forward(torch::Tensor& input, + const torch::Tensor& positions, + const torch::Tensor& cu_query_lens, + int64_t max_query_len, + bool is_prompt) override; + + const torch::Tensor& precomputed_cos_sin_cache() { + return precomputed_cos_sin_cache_; + } + + torch::Tensor get_cos_sin_cache() { return cos_sin_cache_; } + const torch::Tensor& get_sin_cache() const override { return sin_; } + const torch::Tensor& get_cos_cache() const override { return cos_; } + const bool get_interleaved() const override { return interleaved_; } + + protected: + bool interleaved_; + torch::Tensor cos_sin_cache_; + // Pre-formatted [cos_half, sin_half] cache for CUDA/MUSA/ILU kernels. + // Avoids chunk/cat operations on every forward call. + torch::Tensor precomputed_cos_sin_cache_; + + private: + torch::Tensor sin_; + torch::Tensor cos_; +}; +TORCH_MODULE(RotaryEmbedding); + +class MRotaryEmbeddingImpl : public RotaryEmbeddingImpl { + public: + MRotaryEmbeddingImpl(int64_t rotary_dim, + int64_t max_position_embeddings, + int64_t rope_theta, + bool interleaved, + const std::vector& rope_scaling_mrope_section, + const torch::TensorOptions& options); + + void forward(torch::Tensor& q, + torch::Tensor& k, + const torch::Tensor& positions, + const AttentionMetadata& attn_metadata); + + private: + std::vector mrope_section_; + torch::Tensor mrope_cu_seq_lens_; +}; +TORCH_MODULE(MRotaryEmbedding); + +class DeepseekScalingRotaryEmbeddingImpl : public RotaryEmbeddingBase { + public: + DeepseekScalingRotaryEmbeddingImpl( + int64_t head_size, + int64_t rotary_dim, + int64_t max_position_embeddings, + int64_t rope_scaling_original_max_position_embeddings, + int64_t rope_theta, + bool interleaved, + float scaling_factor, + float extrapolation_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float mscale, + float mscale_all_dim, + const torch::TensorOptions& options); + + void forward(torch::Tensor& input, + const torch::Tensor& positions, + const torch::Tensor& cu_query_lens, + int64_t max_query_len, + bool is_prompt) override; + const torch::Tensor& get_sin_cache() const override { return sin_; } + const torch::Tensor& get_cos_cache() const override { return cos_; } + const bool get_interleaved() const override { return interleaved_; } + + private: + int64_t head_size_; + int64_t rotary_dim_; + bool interleaved_; + torch::Tensor sin_; + torch::Tensor cos_; + torch::Tensor cos_sin_cache_; + // Pre-formatted [cos_half, sin_half] cache for CUDA/MUSA/ILU kernels. + // Avoids chunk/cat operations on every forward call. + torch::Tensor precomputed_cos_sin_cache_; +}; +TORCH_MODULE(DeepseekScalingRotaryEmbedding); + +// Factory function: creates the appropriate RoPE type based on model args +std::shared_ptr create_mla_rotary_embedding( + const ModelArgs& args, + int64_t rotary_dim, + int64_t max_position_embeddings, + bool interleaved, + const torch::TensorOptions& options); + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/xllm_layers/ilu/attention.cpp b/ex_engine/xllm_layers/ilu/attention.cpp new file mode 100644 index 00000000..b66f28a4 --- /dev/null +++ b/ex_engine/xllm_layers/ilu/attention.cpp @@ -0,0 +1,189 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "attention.h" + +#include "kernels/ilu/ilu_ops_api.h" +#include "kernels/ops_api.h" + +namespace xllm { +namespace layer { +AttentionImpl::AttentionImpl(int64_t num_heads, + int64_t head_size, + float scale, + int64_t num_kv_heads, + int64_t sliding_window) + : num_heads_(num_heads), + head_size_(head_size), + scale_(scale), + num_kv_heads_(num_kv_heads), + v_head_dim_(head_size), + use_fused_mla_qkv_(false), + enable_lighting_indexer_(false), + enable_mla_(false), + sliding_window_(sliding_window) { + if (sliding_window_ > -1) { + sliding_window_ = sliding_window_ - 1; + } +} + +AttentionImpl::AttentionImpl(int64_t num_heads, + int64_t head_size, + int64_t num_kv_heads, + int64_t v_head_dim, + int64_t sliding_window, + float scale, + bool use_fused_mla_qkv, + bool enable_lighting_indexer, + bool enable_mla) + : num_heads_(num_heads), + head_size_(head_size), + scale_(scale), + num_kv_heads_(num_kv_heads), + v_head_dim_(v_head_dim), + use_fused_mla_qkv_(use_fused_mla_qkv), + enable_lighting_indexer_(enable_lighting_indexer), + enable_mla_(enable_mla), + sliding_window_(sliding_window) { + if (sliding_window_ > -1) { + sliding_window_ = sliding_window_ - 1; + } +} + +std::tuple> AttentionImpl::forward( + const AttentionMetadata& attn_metadata, + torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + KVCache& kv_cache) { + std::optional output_lse = std::nullopt; + torch::Tensor output; + if (enable_mla_) { + output = torch::empty({query.size(0), num_heads_ * v_head_dim_}, + query.options()); + } else { + output = torch::empty_like(query); + } + if (attn_metadata.is_dummy) { + return std::make_tuple(output, output_lse); + } + + bool only_prefill = + attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; + int64_t num_kv_heads = (enable_mla_ && !only_prefill) ? 1 : num_kv_heads_; + torch::Tensor k_cache = kv_cache.get_k_cache(); + std::optional v_cache; + std::optional v; + if (!enable_mla_) { + v = value.view({-1, num_kv_heads, head_size_}); + v_cache = kv_cache.get_v_cache(); + } + + bool skip_process_cache = enable_mla_ && (only_prefill || use_fused_mla_qkv_); + if (!skip_process_cache) { + xllm::kernel::ReshapePagedCacheParams reshape_paged_cache_params; + reshape_paged_cache_params.key = key.view({-1, num_kv_heads, head_size_}); + reshape_paged_cache_params.value = v; + reshape_paged_cache_params.k_cache = k_cache; + reshape_paged_cache_params.v_cache = v_cache; + reshape_paged_cache_params.slot_mapping = attn_metadata.slot_mapping; + xllm::kernel::reshape_paged_cache(reshape_paged_cache_params); + } + + if (enable_lighting_indexer_ || !only_prefill) { + decoder_forward(query, output, k_cache, v_cache, attn_metadata); + } else { + prefill_forward(query, key, value, output, k_cache, v_cache, attn_metadata); + } + + int64_t head_size = enable_mla_ ? v_head_dim_ : head_size_; + output = output.view({-1, num_heads_ * head_size}); + return {output, output_lse}; +} + +void AttentionImpl::prefill_forward(torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata) { + int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_; + std::optional output_lse = std::nullopt; + query = query.view({-1, num_heads_, head_size_}); + output = output.view({-1, num_heads_, head_size_v}); + // torch::Tensor k_cache_ = k_cache; + // torch::Tensor v_cache_ = v_cache.value(); + xllm::kernel::ilu::batch_prefill(query, + k_cache, + v_cache, + output, + output_lse, + attn_metadata.q_cu_seq_lens, + attn_metadata.kv_cu_seq_lens, + /*alibi_slope=*/std::nullopt, + /*attn_bias=*/std::nullopt, + /*q_quant_scale=*/std::nullopt, + /*k_quant_scale=*/std::nullopt, + /*v_quant_scale=*/std::nullopt, + attn_metadata.block_table, + attn_metadata.max_query_len, + attn_metadata.max_seq_len, + scale_, + attn_metadata.is_causal, + sliding_window_, + /*window_size_right=*/-1, + attn_metadata.compute_dtype, + /*return_lse=*/false); +} + +void AttentionImpl::decoder_forward(torch::Tensor& query, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata) { + int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_; + query = query.view({-1, 1, num_heads_, head_size_}); + output = output.view({-1, 1, num_heads_, head_size_v}); + std::optional output_lse = std::nullopt; + + int64_t block_aligned_max_seq_len = + attn_metadata.block_table.size(-1) * k_cache.size(2); + + xllm::kernel::ilu::batch_decode(query, + k_cache, + output, + attn_metadata.block_table, + attn_metadata.kv_seq_lens, + v_cache, + output_lse, + /*q_quant_scale=*/std::nullopt, + /*k_quant_scale=*/std::nullopt, + /*v_quant_scale=*/std::nullopt, + /*out_quant_scale=*/std::nullopt, + /*alibi_slope=*/std::nullopt, + attn_metadata.attn_mask, + attn_metadata.compute_dtype, + block_aligned_max_seq_len, + sliding_window_, + /*window_size_right=*/-1, + scale_, + /*return_lse=*/false, + attn_metadata.is_causal, + /*kv_cache_quant_bit_size=*/-1); +} + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/xllm_layers/ilu/attention.h b/ex_engine/xllm_layers/ilu/attention.h new file mode 100644 index 00000000..a971835f --- /dev/null +++ b/ex_engine/xllm_layers/ilu/attention.h @@ -0,0 +1,82 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include + +#include "framework/kv_cache/kv_cache.h" +#include "framework/model/model_input_params.h" +#include "layers/common/attention_metadata.h" + +namespace xllm { +namespace layer { +class AttentionImpl : public torch::nn::Module { + public: + AttentionImpl() = default; + + AttentionImpl(int64_t num_heads, + int64_t head_size, + float scale, + int64_t num_kv_heads, + int64_t sliding_window); + AttentionImpl(int64_t num_heads, + int64_t head_size, + int64_t num_kv_heads, + int64_t v_head_dim, + int64_t sliding_window, + float scale, + bool use_fused_mla_qkv, + bool enable_lighting_indexer, + bool enable_mla); + + std::tuple> forward( + const AttentionMetadata& attn_metadata, + torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + KVCache& kv_cache); + + void prefill_forward(torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata); + + void decoder_forward(torch::Tensor& query, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata); + + private: + int64_t num_heads_; + int64_t head_size_; + float scale_; + int64_t num_kv_heads_; + int64_t v_head_dim_; + bool use_fused_mla_qkv_; + bool enable_lighting_indexer_; + bool enable_mla_; + int64_t sliding_window_; +}; +TORCH_MODULE(Attention); + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/xllm_layers/ilu/fused_moe.cpp b/ex_engine/xllm_layers/ilu/fused_moe.cpp new file mode 100644 index 00000000..4238012e --- /dev/null +++ b/ex_engine/xllm_layers/ilu/fused_moe.cpp @@ -0,0 +1,797 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "fused_moe.h" + +#include + +#include + +#include "common/global_flags.h" +#include "framework/parallel_state/parallel_state.h" +#include "kernels/ops_api.h" +#include "layers/common/dp_utils.h" +#include "util/utils.h" + +namespace { + +int32_t get_dtype_size(torch::ScalarType dtype) { + return static_cast(torch::elementSize(dtype)); +} + +} // namespace + +namespace xllm { +namespace layer { + +FusedMoEImpl::FusedMoEImpl(const ModelArgs& model_args, + const FusedMoEArgs& moe_args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options) + : num_total_experts_(static_cast(model_args.n_routed_experts())), + topk_(model_args.num_experts_per_tok()), + num_expert_group_(model_args.n_group()), + topk_group_(model_args.topk_group()), + route_scale_(model_args.routed_scaling_factor()), + hidden_size_(model_args.hidden_size()), + n_shared_experts_(model_args.n_shared_experts()), + is_gated_(moe_args.is_gated), + renormalize_(model_args.norm_topk_prob() ? 1 : 0), + hidden_act_(model_args.hidden_act()), + scoring_func_(model_args.scoring_func()), + quant_args_(quant_args), + parallel_args_(parallel_args), + options_(options), + device_(options.device()) { + const int64_t num_experts = num_total_experts_; + const int64_t intermediate_size = + static_cast(model_args.moe_intermediate_size()); + const std::string& topk_method = model_args.topk_method(); + int64_t ep_size = parallel_args.ep_size(); + int64_t ep_rank = 0; + tp_pg_ = parallel_args.tp_group_; + if (ep_size > 1) { + ep_rank = parallel_args.moe_ep_group_->rank(); + tp_pg_ = parallel_args.moe_tp_group_; + } + + // smoothquant check: If quant_method is not empty, only w8a8 smoothquant is + // supported + if (!quant_args.quant_method().empty()) { + if (quant_args.quant_method() != "smoothquant" || quant_args.bits() != 8 || + !quant_args.activation_dynamic()) { + LOG(FATAL) << "FusedMoE only supports w8a8 smoothquant quantization when " + "quant_method is set. " + << "Got quant_method=" << quant_args.quant_method() + << ", bits=" << quant_args.bits() + << ", activation_dynamic=" << quant_args.activation_dynamic(); + } + // If confirmed as smoothquant w8a8, set is_smoothquant_ to true + is_smoothquant_ = true; + } else { + is_smoothquant_ = false; + } + + // Deep EP initialization check + enable_deep_ep_ = FLAGS_expert_parallel_degree == 2 && ep_size > 1; + if (enable_deep_ep_) { + // for now, we only implement the deep ep for decode stage. + // so we will assume the max_token_num is limited to max_batch_size * (1+K) + // K is the number of speculative tokens. + int64_t dispatch_token_size; + if (quant_args.quant_method() == "smoothquant") { + // float32 is for the scale of the quantized input + dispatch_token_size = hidden_size_ * get_dtype_size(torch::kInt8) + + get_dtype_size(torch::kFloat32); + } else { + dispatch_token_size = + hidden_size_ * get_dtype_size(options_.dtype().toScalarType()); + } + torch::ScalarType combine_dtype = options_.dtype().toScalarType(); + int64_t combine_token_size = hidden_size_ * get_dtype_size(combine_dtype); + // Ensure calculation base is at least ep_size + int64_t effective_seqs = + std::max((int64_t)FLAGS_max_seqs_per_batch, (int64_t)ep_size); + // NOTE: FLAGS_max_seqs_per_batch represents the maximum total batch size, + // regardless of the dp size. To ensure robust scheduling and account + // for the worst-case scenario, we must guarantee that each rank is capable + // of handling the maximum possible number of tokens. Therefore, we define + // max_num_tokens_per_rank as the full maximum value, without dividing by + // either the rank count or the dp size. + int64_t max_num_tokens_per_rank = + (1 + FLAGS_num_speculative_tokens) * effective_seqs * topk_; + + // make sure that all layers share the same deep ep instance + // so that the memory footprint is minimized + deep_ep_ = DeepEPManager::get_instance(dispatch_token_size, + combine_token_size, + max_num_tokens_per_rank, + num_experts, + parallel_args, + options_); + + // obtain the buffer and parameters of deep ep + deep_ep_buffer_ = deep_ep_->get_buffer(); + deep_ep_params_ = deep_ep_->get_params(); + + // intermediate buffer that can be initialized once + // we place these tensor here in order to speed up forward pass + int64_t n_tokens_recv = deep_ep_params_.max_num_tokens_recv; + int64_t token_bytes = is_smoothquant_ + ? get_dtype_size(torch::kInt8) + : get_dtype_size(options_.dtype().toScalarType()); + token_bytes = token_bytes * hidden_size_; + int64_t head_size = n_tokens_recv * token_bytes; + dispatch_recv_token_tensor_head_ = + deep_ep_buffer_.combine_send_token_tensor.narrow(0, 0, head_size) + .view({n_tokens_recv, token_bytes}); + // input scale in smoothquant + if (is_smoothquant_) { + int64_t tail_size = n_tokens_recv * get_dtype_size(torch::kFloat32); + dispatch_recv_token_tensor_tail_ = + deep_ep_buffer_.combine_send_token_tensor + .narrow(0, head_size, tail_size) + .view({n_tokens_recv, -1}); + } + } + + // calculate the number of experts per rank + num_experts_per_rank_ = num_experts / ep_size; + start_expert_id_ = ep_rank * num_experts_per_rank_; + + if (topk_method == "noaux_tc") { + e_score_correction_bias_ = register_parameter( + "e_score_correction_bias", torch::empty({num_experts}, options), false); + } + + gate_ = register_module( + "gate_proj", + ReplicatedLinear(hidden_size_, num_experts, false, quant_args, options)); + if (n_shared_experts_ > 0) { + ProcessGroup* shared_expert_pg; + if (parallel_args_.ep_size() > 1) { + // we use tp=1 for shared experts computation in deep ep mode + CHECK(parallel_args_.ep_size() == parallel_args_.world_size()) + << "Models with shared experts only support ep_size equal to " + "world size for now."; + shared_expert_pg = parallel_args.moe_tp_group_; + } else { + shared_expert_pg = parallel_args.process_group_; + } + // The shared experts computation can proceed in parallel with the + // final communication step during the MoE computation, as long as it + // remains independent of any communication operations. For optimal + // performance, ensure that the shared experts layer on each rank always + // maintains its own unique weights. + shared_experts_ = + register_module("shared_experts", + DenseMLP(hidden_size_, + intermediate_size * n_shared_experts_, + is_gated_, + false, + hidden_act_, + /*enable_result_reduction=*/true, + quant_args, + shared_expert_pg, + options)); + } + + // create weight buffer + const int64_t world_size = tp_pg_->world_size(); + int64_t local_intermediate_size = intermediate_size / world_size; + if (is_smoothquant_) { + auto quant_option = options_.dtype(torch::kInt8); + auto fp_option = options_.dtype(torch::kFloat32); + w13_ = register_parameter( + "w13", + torch::empty( + {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_}, + quant_option), + false); + w13_scale_ = register_parameter( + "w13_scale", + torch::empty({num_experts_per_rank_, local_intermediate_size * 2}, + fp_option), + false); + // Note: We do not check enable_deep_ep_ here, since smooth quantization + // information may be needed even when deep EP mode is disabled. This allows + // retrieving quantization parameters for any subset of experts as required. + input_smooth_ = register_parameter( + "input_smooth", + torch::empty({num_total_experts_, hidden_size_}, fp_option), + false); + w2_ = register_parameter( + "w2", + torch::empty( + {num_experts_per_rank_, hidden_size_, local_intermediate_size}, + quant_option), + false); + w2_scale_ = register_parameter( + "w2_scale", + torch::empty({num_experts_per_rank_, hidden_size_}, fp_option), + false); + act_smooth_ = register_parameter( + "act_smooth", + torch::empty({num_experts_per_rank_, local_intermediate_size}, + fp_option), + false); + + } else { + w13_ = register_parameter( + "w13", + torch::empty( + {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_}, + options_), + false); + w2_ = register_parameter( + "w2", + torch::empty( + {num_experts_per_rank_, hidden_size_, local_intermediate_size}, + options_), + false); + } +} + +torch::Tensor FusedMoEImpl::create_group_gemm_output( + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& group_list, + torch::ScalarType dtype, + torch::Tensor& workspace) { + // unify shape logic: define the target shape once. + bool is_3d_weight = (b.dim() != 2); + int64_t num_tokens = a.size(0); + int64_t out_dim = is_3d_weight ? b.size(1) : b.size(0); + + std::vector output_shape; + int64_t required_elements = num_tokens * out_dim; + + if (is_3d_weight) { + output_shape = {num_tokens, out_dim}; + } else { + output_shape = {group_list.size(0), num_tokens, out_dim}; + required_elements *= group_list.size(0); + } + + auto options = a.options().dtype(dtype); + + // non-smoothquant: direct allocation + if (!is_smoothquant_) { + return torch::empty(output_shape, options); + } + + // smoothquant: managed workspace logic + if (!workspace.defined()) { + // Lazy initialization: allocate max buffer for the lifecycle + // Note: accessing class members w13_ and w2_ directly for context + int64_t max_width = std::max(w13_.size(1), w2_.size(1)); + workspace = torch::empty({num_tokens * max_width}, options); + } + + // view construction + CHECK(workspace.numel() >= required_elements) + << "FusedMoE Workspace too small! Alloc: " << workspace.numel() + << ", Req: " << required_elements; + + // utilize the pre-calculated output_shape + return workspace.slice(0, 0, required_elements).view(output_shape); +} + +torch::Tensor FusedMoEImpl::select_experts( + const torch::Tensor& hidden_states_2d, + const torch::Tensor& router_logits_2d, + SelectedExpertInfo& selected_expert_info, + bool enable_all2all_communication) { + // prepare the parameters for select_experts + std::optional e_score_correction_bias = std::nullopt; + if (e_score_correction_bias_.defined()) { + e_score_correction_bias = e_score_correction_bias_; + } + int64_t expert_size = w13_.size(0); + + // Step 1: apply softmax topk or sigmoid topk / routing logic + torch::Tensor reduce_weight; + torch::Tensor expert_id; + { + xllm::kernel::MoeFusedTopkParams moe_active_topk_params; + moe_active_topk_params.input = router_logits_2d; + moe_active_topk_params.topk = topk_; + moe_active_topk_params.num_expert_group = num_expert_group_; + moe_active_topk_params.topk_group = topk_group_; + moe_active_topk_params.normalize = renormalize_; + moe_active_topk_params.normed_by = "topk_logit"; + moe_active_topk_params.scoring_func = scoring_func_; + moe_active_topk_params.route_scale = route_scale_; + moe_active_topk_params.e_score_correction_bias = e_score_correction_bias; + std::tie(reduce_weight, expert_id) = + xllm::kernel::moe_active_topk(moe_active_topk_params); + } + + // Step 2: generate expert ids + torch::Tensor gather_idx; + torch::Tensor combine_idx; + torch::Tensor token_count; + std::optional cusum_token_count; + { + xllm::kernel::MoeGenIdxParams moe_gen_idx_params; + moe_gen_idx_params.expert_id = expert_id; + moe_gen_idx_params.expert_num = num_total_experts_; + std::vector output_vec = + xllm::kernel::moe_gen_idx(moe_gen_idx_params); + gather_idx = output_vec[0]; + combine_idx = output_vec[1]; + token_count = output_vec[2]; + // during all2all communication, we do not need cusum_token_count in the + // following computation + if (enable_all2all_communication) { + cusum_token_count = std::nullopt; + } else { + cusum_token_count = output_vec[3]; + } + } + + // Step 3: expand and quantize input if needed + torch::Tensor expand_hidden_states; + torch::Tensor hidden_states_scale; + torch::Tensor token_count_slice; + // all2all related variables + torch::Tensor dispatch_send_token_tensor; + // in all2all, the input is scattered, so there is no need to slice the token + // count, and we can use the dispatch buffer directly + if (enable_all2all_communication) { + token_count_slice = token_count; + int64_t num_token_expand = hidden_states_2d.size(0) * topk_; + int64_t dispatch_bytes = + num_token_expand * deep_ep_params_.dispatch_token_size; + dispatch_send_token_tensor = + deep_ep_buffer_.dispatch_send_token_tensor.slice(0, 0, dispatch_bytes) + .view({num_token_expand, deep_ep_params_.dispatch_token_size}); + } else { + token_count_slice = + token_count.slice(0, start_expert_id_, start_expert_id_ + expert_size); + } + + if (is_smoothquant_) { + xllm::kernel::ScaledQuantizeParams scaled_quantize_params; + scaled_quantize_params.x = hidden_states_2d; + // use dispatch_send_token_tensor buffer for input + // to reduce memory footprint + if (enable_all2all_communication) { + scaled_quantize_params.smooth = input_smooth_; + scaled_quantize_params.output = + dispatch_send_token_tensor.slice(1, 0, hidden_size_); + } else { + scaled_quantize_params.smooth = input_smooth_.slice( + 0, start_expert_id_, start_expert_id_ + expert_size); + scaled_quantize_params.gather_index_start_position = + cusum_token_count.value().index({start_expert_id_}).unsqueeze(0); + } + scaled_quantize_params.token_count = token_count_slice; + scaled_quantize_params.gather_index = gather_idx; + scaled_quantize_params.act_mode = "none"; + scaled_quantize_params.active_coef = 1.0; + scaled_quantize_params.is_gated = false; + scaled_quantize_params.quant_type = torch::kChar; + std::tie(expand_hidden_states, hidden_states_scale) = + xllm::kernel::scaled_quantize(scaled_quantize_params); + if (enable_all2all_communication) { + // since view_as_dtype has not supported stride yet, + // we need to copy the scale output to the dispatch buffer + torch::Tensor dispatch_scale_slice = + dispatch_send_token_tensor.slice(1, hidden_size_); + torch::Tensor hidden_states_scale_bytes = + view_as_dtype(hidden_states_scale, torch::kInt8) + .view_as(dispatch_scale_slice); + dispatch_scale_slice.copy_(hidden_states_scale_bytes); + } + } else { + xllm::kernel::MoeExpandInputParams moe_expand_input_params; + moe_expand_input_params.input = hidden_states_2d; + moe_expand_input_params.gather_index = gather_idx; + moe_expand_input_params.combine_idx = combine_idx; + moe_expand_input_params.topk = topk_; + expand_hidden_states = + xllm::kernel::moe_expand_input(moe_expand_input_params); + if (enable_all2all_communication) { + // use copy to place the output inside the dispatch buffer + torch::Tensor dispatch_tensor = + view_as_dtype(expand_hidden_states, torch::kChar); + dispatch_send_token_tensor.copy_(dispatch_tensor); + } + } + + // collect the selected tensor + selected_expert_info.reduce_weight = reduce_weight; + selected_expert_info.combine_idx = combine_idx; + selected_expert_info.token_count_slice = token_count_slice; + selected_expert_info.cusum_token_count = cusum_token_count; + if (is_smoothquant_) { + selected_expert_info.input_scale = hidden_states_scale; + } + + return expand_hidden_states; +} + +torch::Tensor FusedMoEImpl::forward_experts(const torch::Tensor& hidden_states, + const torch::Tensor& router_logits, + bool enable_all2all_communication) { + if (!stream_initialized_) { + // update device record + device_ = xllm::Device(hidden_states.device()); + + // acquire streams from the pool again + routed_stream_ = device_.get_stream_from_pool(); + shared_stream_ = device_.get_stream_from_pool(); + stream_initialized_ = true; + } + + std::optional e_score_correction_bias = std::nullopt; + if (e_score_correction_bias_.defined()) { + e_score_correction_bias = e_score_correction_bias_; + } + + // prepare the parameters for MoE computation + torch::Tensor shared_expert_output; + torch::IntArrayRef hidden_states_shape = hidden_states.sizes(); + torch::ScalarType hidden_states_dtype = hidden_states.dtype().toScalarType(); + torch::Tensor hidden_states_2d = + hidden_states.reshape({-1, hidden_states.size(-1)}); + torch::Tensor router_logits_2d = + router_logits.reshape({-1, router_logits.size(-1)}); + int64_t group_gemm_max_dim = enable_all2all_communication + ? deep_ep_params_.max_num_tokens_recv / topk_ + : hidden_states_2d.size(0); + int64_t expert_size = w13_.size(0); + + // Step 1-3: select experts + SelectedExpertInfo selected_expert_info; + torch::Tensor expand_hidden_states = + select_experts(hidden_states_2d, + router_logits_2d, + selected_expert_info, + enable_all2all_communication); + + // Communciation Step 1: Dipatch + // intermediate outputs that are used both in dispatch and combine + torch::Tensor gather_by_rank_index; + torch::Tensor token_sum; + if (enable_all2all_communication) { + int64_t dispatch_token_num = hidden_states_2d.size(0) * topk_; + + // 1. Dispatch Step: Generate layout and send data + deep_ep_->dispatch_step(dispatch_token_num, + selected_expert_info.token_count_slice); + + // 2. Process Result: Generate indices and unpack to computation buffer + // use the buffer during initialization for the output + expand_hidden_states = dispatch_recv_token_tensor_head_; + std::optional output_tail = std::nullopt; + if (is_smoothquant_) { + output_tail = dispatch_recv_token_tensor_tail_; + // update selected_expert_info with the tail (input scale) + selected_expert_info.input_scale = output_tail; + } + + DeepEPMetaResult deep_ep_meta = deep_ep_->process_dispatch_result( + num_experts_per_rank_, expand_hidden_states, output_tail); + + // Extract metadata for subsequent steps + gather_by_rank_index = deep_ep_meta.gather_rank_index; + selected_expert_info.token_count_slice = deep_ep_meta.token_count_slice; + token_sum = deep_ep_meta.token_sum; + } + + // common gemm workspace for reduce memory footprint + torch::Tensor gemm_workspace; + + // Step 4: group gemm 1 + torch::Tensor gemm1_out = + create_group_gemm_output(expand_hidden_states, + w13_, + selected_expert_info.token_count_slice, + hidden_states_dtype, + gemm_workspace); + // ensure the lifespan of these parameters via brace + { + xllm::kernel::GroupGemmParams group_gemm_params; + torch::ScalarType a_dtype = + is_smoothquant_ ? torch::kInt8 : hidden_states_dtype; + group_gemm_params.a = + view_as_dtype(expand_hidden_states, a_dtype).view({-1, hidden_size_}); + group_gemm_params.b = w13_; + group_gemm_params.token_count = + selected_expert_info.token_count_slice.to("cpu"); + if (is_smoothquant_) { + torch::Tensor a_scale = + selected_expert_info.input_scale.value().flatten(); + selected_expert_info.input_scale = + view_as_dtype(a_scale, torch::kFloat32); + group_gemm_params.a_scale = selected_expert_info.input_scale; + group_gemm_params.b_scale = w13_scale_; + } + group_gemm_params.max_dim = group_gemm_max_dim; + group_gemm_params.trans_a = false; + group_gemm_params.trans_b = true; + group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1; + group_gemm_params.output = gemm1_out; + group_gemm_params.combine_idx = std::nullopt; + gemm1_out = xllm::kernel::group_gemm(group_gemm_params); + } + + // Step 5: activation or scaled quantization(fused with activation) + torch::Tensor act_out; + torch::Tensor act_out_scale; + if (is_smoothquant_) { + int64_t slice_dim = gemm1_out.size(1); + if (is_gated_) slice_dim /= 2; + // slice operation is a view, does not take up extra memory, but points to + // the same memory + act_out = expand_hidden_states.slice(1, 0, slice_dim); + act_out_scale = + selected_expert_info.input_scale.value().slice(0, 0, gemm1_out.size(0)); + // call scaled quantization kernel (also fused with activation) + xllm::kernel::ScaledQuantizeParams scaled_quantize_params; + scaled_quantize_params.x = gemm1_out; + scaled_quantize_params.smooth = act_smooth_; + scaled_quantize_params.token_count = selected_expert_info.token_count_slice; + scaled_quantize_params.output = act_out; + scaled_quantize_params.output_scale = act_out_scale; + scaled_quantize_params.act_mode = hidden_act_; + scaled_quantize_params.active_coef = 1.0; + scaled_quantize_params.is_gated = is_gated_; + scaled_quantize_params.quant_type = torch::kChar; + std::tie(act_out, act_out_scale) = + xllm::kernel::scaled_quantize(scaled_quantize_params); + } else { + act_out = is_gated_ + ? gemm1_out.slice(1, 0, gemm1_out.size(1) / 2).contiguous() + : gemm1_out; + // call activation kernel + xllm::kernel::ActivationParams activation_params; + activation_params.input = gemm1_out; + activation_params.output = act_out; + activation_params.cusum_token_count = + selected_expert_info.cusum_token_count; + activation_params.act_mode = hidden_act_; + activation_params.is_gated = is_gated_; + activation_params.start_expert_id = start_expert_id_; + activation_params.expert_size = expert_size; + xllm::kernel::active(activation_params); + } + + // Step 6: group gemm 2 + torch::Tensor gemm2_out = + create_group_gemm_output(act_out, + w2_, + selected_expert_info.token_count_slice, + hidden_states_dtype, + gemm_workspace); + // ensure the lifespan of these parameters via brace + { + xllm::kernel::GroupGemmParams group_gemm_params; + group_gemm_params.a = act_out; + group_gemm_params.b = w2_; + group_gemm_params.token_count = + selected_expert_info.token_count_slice.to("cpu"); + if (is_smoothquant_) { + group_gemm_params.a_scale = act_out_scale; + group_gemm_params.b_scale = w2_scale_; + } + group_gemm_params.max_dim = group_gemm_max_dim; + group_gemm_params.trans_a = false; + group_gemm_params.trans_b = true; + group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1; + group_gemm_params.output = gemm2_out; + group_gemm_params.combine_idx = selected_expert_info.combine_idx; + gemm2_out = xllm::kernel::group_gemm(group_gemm_params); + } + + // Communciation Step 2: Combine + if (enable_all2all_communication) { + int64_t num_token_expand = hidden_states_2d.size(0) * topk_; + // Delegate pack, layout generation and combine to DeepEP + torch::Tensor combine_send_layout = + deep_ep_->combine_step_pack(gemm2_out, + gather_by_rank_index, + token_sum, + hidden_size_, + hidden_states_dtype); + + // create a wait event for the current stream to finish computation + auto current_stream = device_.current_stream(); + routed_stream_->wait_stream(*current_stream); + // pure communciation kernel: dispatch + { + torch::StreamGuard stream_guard = routed_stream_->set_stream_guard(); + gemm2_out = deep_ep_->combine_step_comm(combine_send_layout, + num_token_expand, + hidden_size_, + hidden_states_dtype); + } + + // pure computation kernel: shared experts + if (n_shared_experts_ > 0) { + shared_stream_->wait_stream(*current_stream); + torch::StreamGuard stream_guard = shared_stream_->set_stream_guard(); + shared_expert_output = shared_experts_(hidden_states); + } + + // join for parallelization + current_stream->wait_stream(*routed_stream_); + if (n_shared_experts_ > 0) { + current_stream->wait_stream(*shared_stream_); + } + } + + // After group gemm is finished, some tensors are no + // longer needed. We must explicitly release the memory. + expand_hidden_states = torch::Tensor(); + selected_expert_info.input_scale = std::nullopt; + act_out = torch::Tensor(); + + // Step 7: combine the intermediate results and get the final hidden states + torch::Tensor final_hidden_states; + // ensure the lifespan of these parameters via brace + { + xllm::kernel::MoeCombineResultParams moe_combine_result_params; + moe_combine_result_params.input = gemm2_out; + moe_combine_result_params.reduce_weight = + selected_expert_info.reduce_weight; + moe_combine_result_params.gather_ids = selected_expert_info.combine_idx; + moe_combine_result_params.cusum_token_count = + selected_expert_info.cusum_token_count; + moe_combine_result_params.start_expert_id = start_expert_id_; + moe_combine_result_params.expert_size = expert_size; + moe_combine_result_params.bias = std::nullopt; + // if all2all communication is enabled and shared output is provided, + // we will fused the add up to combine result + if (enable_all2all_communication && n_shared_experts_ > 0) { + moe_combine_result_params.residual = + shared_expert_output.reshape({-1, shared_expert_output.size(-1)}); + } + + final_hidden_states = + xllm::kernel::moe_combine_result(moe_combine_result_params); + } + + // reshape the final hidden states to the original shape + final_hidden_states = final_hidden_states.reshape(hidden_states_shape); + + if (enable_all2all_communication) { + return final_hidden_states; + } + + // Communciation Step 3: AllReduce for non-all2all communication + // shared experts can be parallelized with the final communication step + // during moe computation. + auto current_stream = device_.current_stream(); + routed_stream_->wait_stream(*current_stream); + { + torch::StreamGuard stream_guard = routed_stream_->set_stream_guard(); + if (tp_pg_->world_size() > 1) { + final_hidden_states = parallel_state::reduce(final_hidden_states, tp_pg_); + } + if (parallel_args_.ep_size() > 1) { + final_hidden_states = parallel_state::reduce( + final_hidden_states, parallel_args_.moe_ep_group_); + } + } + + if (n_shared_experts_ > 0) { + shared_stream_->wait_stream(*current_stream); + torch::StreamGuard stream_guard = shared_stream_->set_stream_guard(); + // for non all2all, we compute the shared experts parallelized with the + // final communication step + shared_expert_output = shared_experts_(hidden_states); + shared_expert_output = + shared_expert_output.reshape({-1, shared_expert_output.size(-1)}); + } + + // join for parallelization + current_stream->wait_stream(*routed_stream_); + if (n_shared_experts_ > 0) { + current_stream->wait_stream(*shared_stream_); + final_hidden_states += shared_expert_output; + } + + return final_hidden_states; +} + +torch::Tensor FusedMoEImpl::forward(const torch::Tensor& hidden_states, + const ModelInputParams& input_params) { + // we only support all2all communication for decode stage for now + bool enable_all2all_communication = + enable_deep_ep_ && std::all_of(input_params.dp_is_decode.begin(), + input_params.dp_is_decode.end(), + [](int32_t val) { return val == 1; }); + + bool is_dp_ep_parallel = + parallel_args_.dp_size() > 1 && parallel_args_.ep_size() > 1; + // during all2all communication, the output has been + // gathered and sliced by dispatch and combine steps, + // so we do not need to gather input and slice output again + bool need_gather_and_slice = + is_dp_ep_parallel && !enable_all2all_communication; + + auto input = hidden_states; + if (need_gather_and_slice) { + input = parallel_state::gather(input, + parallel_args_.dp_local_process_group_, + input_params.dp_global_token_nums); + } + // MoE Gate + auto router_logits = gate_(input); + + // MoE Experts + auto output = + forward_experts(input, router_logits, enable_all2all_communication); + + if (need_gather_and_slice) { + output = get_dp_local_slice(output, input_params, parallel_args_); + } + + return output; +} + +void FusedMoEImpl::load_e_score_correction_bias(const StateDict& state_dict) { + if (e_score_correction_bias_.defined() && + !e_score_correction_bias_is_loaded_) { + LOAD_WEIGHT(e_score_correction_bias); + } +} + +void FusedMoEImpl::load_experts(const StateDict& state_dict) { + const int64_t rank = tp_pg_->rank(); + const int64_t world_size = tp_pg_->world_size(); + const int64_t start_expert_id = start_expert_id_; + const int64_t num_experts_per_rank = num_experts_per_rank_; + const int64_t num_total_experts = num_total_experts_; + std::vector prefixes = {"gate_proj.", "up_proj."}; + if (is_smoothquant_) { + LOAD_MOE_FUSED_WEIGHT("qweight", w1, w3, w13); + LOAD_MOE_FUSED_WEIGHT("per_channel_scale", w1_scale, w3_scale, w13_scale); + // When supporting DeepEP All2All mode, + // we need to load the complete set of expert weights corresponding to + // "up_proj.smooth". Note that even if deep EP mode is not enabled, it + // remains possible to retrieve the smooth quantization information for a + // subset of experts. Therefore, we intentionally do not check whether + // deep_ep_ is enabled in this case. + LOAD_MOE_ALL_EXPERT_WEIGHT("up_proj.", "smooth", input_smooth, -1); + LOAD_MOE_WEIGHT("down_proj.", "qweight", w2, 1); + LOAD_MOE_WEIGHT("down_proj.", "per_channel_scale", w2_scale, -1); + LOAD_MOE_WEIGHT("down_proj.", "smooth", act_smooth, 0); + } else { + LOAD_MOE_FUSED_WEIGHT("weight", w1, w3, w13); + LOAD_MOE_WEIGHT("down_proj.", "weight", w2, 1); + } +} + +void FusedMoEImpl::load_state_dict(const StateDict& state_dict) { + if (state_dict.size() == 0) { + return; + } + + if (n_shared_experts_ > 0) { + shared_experts_->load_state_dict( + state_dict.get_dict_with_prefix("shared_experts.")); + } + gate_->load_state_dict(state_dict.get_dict_with_prefix("gate.")); + load_e_score_correction_bias(state_dict.get_dict_with_prefix("gate.")); + load_experts(state_dict.get_dict_with_prefix("experts.")); +} + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/xllm_layers/ilu/fused_moe.h b/ex_engine/xllm_layers/ilu/fused_moe.h new file mode 100644 index 00000000..3e477064 --- /dev/null +++ b/ex_engine/xllm_layers/ilu/fused_moe.h @@ -0,0 +1,131 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include "framework/model/model_args.h" +#include "framework/model/model_input_params.h" +#include "framework/parallel_state/parallel_args.h" +#include "framework/quant_args.h" +#include "framework/state_dict/state_dict.h" +#include "framework/state_dict/utils.h" +#include "layers/common/deep_ep.h" +#include "layers/common/dense_mlp.h" +#include "layers/common/fused_moe_base.h" +#include "layers/common/linear.h" +#include "platform/device.h" +#include "util/tensor_helper.h" + +namespace xllm { +namespace layer { + +class FusedMoEImpl : public torch::nn::Module { + public: + FusedMoEImpl() = default; + FusedMoEImpl(const ModelArgs& model_args, + const FusedMoEArgs& moe_args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options); + + torch::Tensor forward_experts(const torch::Tensor& hidden_states, + const torch::Tensor& router_logits, + bool enable_all2all_communication); + torch::Tensor forward(const torch::Tensor& hidden_states, + const ModelInputParams& input_params); + void load_state_dict(const StateDict& state_dict); + + private: + // struct to store the selected expert info + struct SelectedExpertInfo { + torch::Tensor reduce_weight; + torch::Tensor combine_idx; + torch::Tensor token_count_slice; + std::optional cusum_token_count; + std::optional input_scale; + }; + + // initial steps for MoE computation, select the experts for each token + torch::Tensor select_experts(const torch::Tensor& hidden_states_2d, + const torch::Tensor& router_logits_2d, + SelectedExpertInfo& selected_expert_info, + bool enable_all2all_communication); + + private: + int64_t num_total_experts_; + int64_t topk_; + int64_t num_expert_group_; + int64_t topk_group_; + double route_scale_; + int64_t hidden_size_; + int64_t n_shared_experts_; + bool is_gated_; + int64_t renormalize_; + std::string hidden_act_; + std::string scoring_func_; + bool is_smoothquant_; + + int64_t num_experts_per_rank_; + int64_t start_expert_id_; + + // Deep EP related parameters + bool enable_deep_ep_; + DeepEPBuffer deep_ep_buffer_; + DeepEPParams deep_ep_params_; + torch::Tensor dispatch_recv_token_tensor_head_; + torch::Tensor dispatch_recv_token_tensor_tail_; + + // steams for parallel shared experts + std::unique_ptr shared_stream_; + std::unique_ptr routed_stream_; + xllm::Device device_; + bool stream_initialized_ = false; + + ReplicatedLinear gate_{nullptr}; + DenseMLP shared_experts_{nullptr}; + DeepEP deep_ep_{nullptr}; + + QuantArgs quant_args_; + ParallelArgs parallel_args_; + torch::TensorOptions options_; + ProcessGroup* tp_pg_; + + DEFINE_WEIGHT(w13); + DEFINE_FUSED_WEIGHT(w1); + DEFINE_FUSED_WEIGHT(w3); + DEFINE_FUSED_WEIGHT(w2); + DEFINE_WEIGHT(e_score_correction_bias); + DEFINE_WEIGHT(w13_scale); + DEFINE_FUSED_WEIGHT(w1_scale); + DEFINE_FUSED_WEIGHT(w3_scale); + DEFINE_FUSED_WEIGHT(w2_scale); + DEFINE_FUSED_WEIGHT(input_smooth); + DEFINE_FUSED_WEIGHT(act_smooth); + + void load_e_score_correction_bias(const StateDict& state_dict); + void load_experts(const StateDict& state_dict); + // create the group gemm output tensor with the workspace + torch::Tensor create_group_gemm_output(const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& group_list, + torch::ScalarType dtype, + torch::Tensor& workspace); +}; +TORCH_MODULE(FusedMoE); + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.cpp b/ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.cpp index 0b97be6d..7d572476 100644 --- a/ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.cpp +++ b/ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.cpp @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2026 The xLLM Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at @@ -123,53 +123,19 @@ torch::Tensor Qwen3_5GatedDeltaNetImpl::merge_ba_from_split_activations( } std::pair -Qwen3_5GatedDeltaNetImpl::project_decode_inputs( - const torch::Tensor& hidden_states) { - const auto reshape_projection = [](const torch::Tensor& projection) { - return projection.view({projection.size(0), -1, projection.size(-1)}); - }; - auto qkv = reshape_projection(in_proj_qkv_->forward(hidden_states)); - auto z_proj = reshape_projection(in_proj_z_->forward(hidden_states)); - auto b_proj = reshape_projection(in_proj_b_->forward(hidden_states)); - auto a_proj = reshape_projection(in_proj_a_->forward(hidden_states)); - return {merge_qkvz_from_split_activations(qkv, z_proj), - merge_ba_from_split_activations(b_proj, a_proj)}; -} - -std::pair -Qwen3_5GatedDeltaNetImpl::project_flat_inputs( - const torch::Tensor& hidden_states) { - auto qkv = in_proj_qkv_->forward(hidden_states).unsqueeze(0); - auto z_proj = in_proj_z_->forward(hidden_states).unsqueeze(0); - auto b_proj = in_proj_b_->forward(hidden_states).unsqueeze(0); - auto a_proj = in_proj_a_->forward(hidden_states).unsqueeze(0); - auto qkvz = merge_qkvz_from_split_activations(qkv, z_proj); - auto ba = merge_ba_from_split_activations(b_proj, a_proj); - return {qkvz.view({hidden_states.size(0), qkvz.size(-1)}).contiguous(), - ba.view({hidden_states.size(0), ba.size(-1)}).contiguous()}; -} - -std::optional< - std::tuple> -Qwen3_5GatedDeltaNetImpl::project_split_inputs( +Qwen3_5GatedDeltaNetImpl::project_padded_inputs( const torch::Tensor& hidden_states, const AttentionMetadata& attn_metadata) { - auto qkv = reshape_projected_tokens_with_pad( - attn_metadata, in_proj_qkv_->forward(hidden_states)); - auto z_proj = reshape_projected_tokens_with_pad( - attn_metadata, in_proj_z_->forward(hidden_states)); - auto b_proj = reshape_projected_tokens_with_pad( - attn_metadata, in_proj_b_->forward(hidden_states)); - auto a_proj = reshape_projected_tokens_with_pad( - attn_metadata, in_proj_a_->forward(hidden_states)); - - const int64_t batch_size = qkv.size(0); - const int64_t seq_len = qkv.size(1); - auto z = - z_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_}); - auto b = b_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_}); - auto a = a_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_}); - return std::make_tuple(qkv, z, b, a); + auto qkv = reshape_qkvz_with_pad(attn_metadata, + in_proj_qkv_->forward(hidden_states)); + auto z_proj = + reshape_qkvz_with_pad(attn_metadata, in_proj_z_->forward(hidden_states)); + auto b_proj = + reshape_qkvz_with_pad(attn_metadata, in_proj_b_->forward(hidden_states)); + auto a_proj = + reshape_qkvz_with_pad(attn_metadata, in_proj_a_->forward(hidden_states)); + return {merge_qkvz_from_split_activations(qkv, z_proj), + merge_ba_from_split_activations(b_proj, a_proj)}; } void Qwen3_5GatedDeltaNetImpl::load_projection_state_dict( diff --git a/ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.h b/ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.h index 7c782e3f..bec6c1c6 100644 --- a/ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.h +++ b/ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.h @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2026 The xLLM Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -17,9 +17,7 @@ limitations under the License. #include -#include #include -#include #include #include "qwen3_next_gated_delta_net.h" @@ -36,15 +34,9 @@ class Qwen3_5GatedDeltaNetImpl : public Qwen3NextGatedDeltaNetImpl { const torch::TensorOptions& options); protected: - std::pair project_decode_inputs( - const torch::Tensor& hidden_states) override; - std::pair project_flat_inputs( - const torch::Tensor& hidden_states) override; - std::optional< - std::tuple> - project_split_inputs(const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata) override; - bool use_fla_ssm_state_layout() const override { return true; } + std::pair project_padded_inputs( + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata) override; void load_projection_state_dict(const StateDict& state_dict) override; void verify_projection_weights(const std::string& prefix) const override; diff --git a/ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.cpp b/ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.cpp index 7f8b4b5c..cec9a95e 100644 --- a/ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.cpp +++ b/ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.cpp @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2026 The xLLM Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at @@ -15,12 +15,9 @@ limitations under the License. #include #include -#include #include -#include "xllm/core/kernels/npu/npu_ops_api.h" #include "xllm/core/kernels/ops_api.h" -#include "xllm/core/platform/npu/acl_graph_task_update_context.h" namespace xllm { namespace layer { @@ -31,31 +28,6 @@ torch::Tensor l2norm(const torch::Tensor& x, int64_t dim, double eps = 1e-6) { return x / norm; } -torch::Tensor repeat_tensor_heads(const torch::Tensor& tensor, - int64_t target_heads, - int64_t head_dim) { - const int64_t current_heads = tensor.size(head_dim); - if (current_heads == target_heads) { - return tensor; - } - CHECK_GT(current_heads, 0) << "current heads must be positive"; - CHECK_EQ(target_heads % current_heads, 0) - << "target heads must be divisible by current heads, target_heads=" - << target_heads << ", current_heads=" << current_heads; - - const int64_t repeats = target_heads / current_heads; - std::vector view_shape = tensor.sizes().vec(); - view_shape.insert(view_shape.begin() + head_dim + 1, 1); - std::vector expand_shape = view_shape; - expand_shape[head_dim + 1] = repeats; - std::vector output_shape = tensor.sizes().vec(); - output_shape[head_dim] = target_heads; - return tensor.unsqueeze(head_dim + 1) - .expand(expand_shape) - .reshape(output_shape) - .contiguous(); -} - std::tuple torch_recurrent_gated_delta_rule( torch::Tensor query, torch::Tensor key, @@ -80,9 +52,6 @@ std::tuple torch_recurrent_gated_delta_rule( value = to_float32_and_transpose(value); beta = to_float32_and_transpose(beta); g = to_float32_and_transpose(g); - const int64_t value_num_heads = value.size(1); - query = repeat_tensor_heads(query, value_num_heads, 1); - key = repeat_tensor_heads(key, value_num_heads, 1); int64_t batch_size = key.size(0); int64_t num_heads = key.size(1); @@ -150,15 +119,12 @@ std::tuple torch_chunk_gated_delta_rule( value = to_float32(value); beta = to_float32(beta); g = to_float32(g); - const int64_t value_num_heads = value.size(1); - query = repeat_tensor_heads(query, value_num_heads, 1); - key = repeat_tensor_heads(key, value_num_heads, 1); - int64_t batch_size = query.size(0); - int64_t num_heads = query.size(1); - int64_t sequence_length = query.size(2); - int64_t k_head_dim = key.size(-1); - int64_t v_head_dim = value.size(-1); + auto batch_size = query.size(0); + auto num_heads = query.size(1); + auto sequence_length = query.size(2); + auto k_head_dim = key.size(-1); + auto v_head_dim = value.size(-1); int64_t pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size; query = torch::nn::functional::pad( @@ -276,164 +242,6 @@ std::tuple torch_chunk_gated_delta_rule( core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype); return std::make_tuple(core_attn_out, last_recurrent_state); } - -int64_t get_checkpoint_stride(const torch::Tensor& conv_cache, - const torch::Tensor& ssm_cache) { - if (!conv_cache.defined() || !ssm_cache.defined() || - conv_cache.numel() == 0 || ssm_cache.numel() == 0) { - return 1; - } - CHECK_GT(conv_cache.size(0), 0) << "conv cache must have positive batch dim"; - CHECK_EQ(ssm_cache.size(0) % conv_cache.size(0), 0) - << "ssm cache checkpoint layout mismatch, ssm_rows=" << ssm_cache.size(0) - << ", conv_rows=" << conv_cache.size(0); - return ssm_cache.size(0) / conv_cache.size(0); -} - -torch::Tensor build_linear_state_base_indices( - const torch::Tensor& logical_state_indices, - int64_t checkpoint_stride) { - if (checkpoint_stride == 1) { - return logical_state_indices; - } - return logical_state_indices * checkpoint_stride; -} - -torch::Tensor expand_sequence_tensor_to_batch(const torch::Tensor& tensor, - int64_t target_batch, - const char* tensor_name) { - CHECK(tensor.defined()) << tensor_name << " must be defined"; - CHECK_EQ(tensor.dim(), 1) << tensor_name << " must be a 1D tensor."; - const int64_t source_batch = tensor.size(0); - if (source_batch == target_batch) { - return tensor.contiguous(); - } - CHECK_GT(source_batch, 0) << tensor_name << " must not be empty."; - CHECK_EQ(target_batch % source_batch, 0) - << tensor_name << " cannot be expanded from " << source_batch << " to " - << target_batch; - const int64_t repeat_count = target_batch / source_batch; - return tensor.unsqueeze(1) - .expand({source_batch, repeat_count}) - .reshape({target_batch}) - .contiguous(); -} - -torch::Tensor run_causal_conv1d_graph_update( - const std::shared_ptr& graph_context, - const torch::Tensor& x, - const torch::Tensor& weight, - const torch::Tensor& conv_state, - const std::optional& bias, - const std::vector& query_start_loc, - const std::vector& cache_indices, - const std::vector& num_accepted_tokens, - xllm::npu::CausalConv1dGraphBranch branch) { - CHECK(graph_context != nullptr && graph_context->capturing) - << "causal_conv1d graph update can only be registered during capture"; - - c10_npu::NPUStream stream = c10_npu::getCurrentNPUStream(); - auto event = std::make_shared(ACL_EVENT_EXTERNAL); - event->block(stream); - event->reset(stream); - - torch::Tensor output; - c10_npu::graph_task_group_begin(stream); - const std::vector empty_host_args; - CHECK(!query_start_loc.empty()) - << "query_start_loc must be populated for causal_conv1d graph update"; - CHECK_EQ(query_start_loc.back(), x.size(0)) - << "query_start_loc must be padded to x.shape[0] during graph capture"; - CHECK_EQ(cache_indices.size() + 1, query_start_loc.size()) - << "cache_indices must be sequence-scoped"; - if (branch == xllm::npu::CausalConv1dGraphBranch::kSpecVerify) { - CHECK_EQ(num_accepted_tokens.size(), cache_indices.size()) - << "num_accepted_tokens must be sequence-scoped for spec verify"; - } - - output = torch::empty_like(x); - xllm::kernel::causal_conv1d_out(output, - x, - weight, - conv_state, - bias, - torch::IntArrayRef(query_start_loc), - torch::IntArrayRef(cache_indices), - torch::IntArrayRef(empty_host_args), - torch::IntArrayRef(num_accepted_tokens), - xllm::npu::kCausalConv1dActivationSilu, - xllm::npu::kCausalConv1dGraphPadSlotId, - xllm::npu::kCausalConv1dRunModeUpdate); - c10_npu::NPUTaskGroupHandle handle = c10_npu::graph_task_group_end(stream); - - xllm::npu::CausalConv1dGraphTask task; - task.output = output; - task.x = x; - task.weight = weight; - task.conv_state = conv_state; - task.bias = bias; - task.activation_mode = xllm::npu::kCausalConv1dActivationSilu; - task.pad_slot_id = xllm::npu::kCausalConv1dGraphPadSlotId; - task.run_mode = xllm::npu::kCausalConv1dRunModeUpdate; - task.branch = branch; - task.handle = handle; - task.event = std::move(event); - graph_context->causal_conv1d_tasks.emplace_back(std::move(task)); - return output; -} - -torch::Tensor run_spec_verify_gated_delta_rule( - torch::Tensor query, - torch::Tensor key, - torch::Tensor value, - torch::Tensor g, - torch::Tensor beta, - torch::Tensor& ssm_cache, - const torch::Tensor& checkpoint_indices, - const torch::Tensor& num_accepted_tokens, - const torch::Tensor& cu_seq_lens, - const std::vector& q_seq_lens_vec, - double scale) { - const auto device = value.device(); - const int64_t batch_size = value.size(0); - const int64_t seq_len = value.size(1); - const int64_t total_seq_len = batch_size * seq_len; - CHECK_EQ(cu_seq_lens.numel(), batch_size + 1) - << "GDN spec verify cu_seq_lens must be cumulative."; - CHECK_EQ(q_seq_lens_vec.size(), static_cast(batch_size)) - << "GDN spec verify q_seq_lens_vec must be per sequence."; - for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { - CHECK_EQ(q_seq_lens_vec[batch_idx], seq_len) - << "Qwen3.5 spec verify fused recurrent path expects dense " - "same-length validate tokens."; - } - - xllm::kernel::FusedRecurrentGatedDeltaRuleParams params; - params.q = query.reshape({1, total_seq_len, query.size(-2), query.size(-1)}) - .contiguous(); - params.k = - key.reshape({1, total_seq_len, key.size(-2), key.size(-1)}).contiguous(); - params.v = value.reshape({1, total_seq_len, value.size(-2), value.size(-1)}) - .contiguous(); - params.g = g.to(torch::kFloat32) - .reshape({1, total_seq_len, g.size(-1)}) - .contiguous(); - params.beta = beta.reshape({1, total_seq_len, beta.size(-1)}).contiguous(); - params.scale = static_cast(scale); - params.initial_state = ssm_cache; - params.inplace_final_state = true; - params.cu_seqlens = cu_seq_lens.to(torch::kLong).contiguous(); - params.ssm_state_indices = checkpoint_indices.contiguous(); - params.num_accepted_tokens = - num_accepted_tokens.to(device, torch::kInt32).contiguous(); - params.use_qk_l2norm_in_kernel = true; - - auto output_and_state = - xllm::kernel::fused_recurrent_gated_delta_rule(params); - return output_and_state.first.view( - {batch_size, seq_len, value.size(-2), value.size(-1)}); -} - } // namespace Qwen3GatedDeltaNetBaseImpl::Qwen3GatedDeltaNetBaseImpl( @@ -495,11 +303,7 @@ void Qwen3GatedDeltaNetBaseImpl::load_common_state_dict( if (auto w = state_dict.get_tensor("conv1d.weight"); w.defined()) { conv1d_->load_state_dict( - StateDict({{"weight", w.squeeze(1)}}, - static_cast(state_dict.prefix()) + "conv1d."), - shard_tensor_count, - shard_sizes); - conv1d_->weight().set_(conv1d_->weight().transpose(0, 1).contiguous()); + StateDict({{"weight", w.squeeze(1)}}), shard_tensor_count, shard_sizes); } o_proj_->load_state_dict(state_dict.get_dict_with_prefix("out_proj.")); if (auto w = state_dict.get_tensor("norm.weight"); w.defined()) { @@ -518,279 +322,87 @@ void Qwen3GatedDeltaNetBaseImpl::verify_common_loaded_weights( << prefix << "A_log"; } -std::pair -Qwen3GatedDeltaNetBaseImpl::project_padded_inputs( - const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata) { - if (attn_metadata.is_prefill || attn_metadata.is_chunked_prefill) { - auto [qkvz_flat, ba_flat] = project_flat_inputs(hidden_states); - return {reshape_projected_tokens_with_pad(attn_metadata, qkvz_flat), - reshape_projected_tokens_with_pad(attn_metadata, ba_flat)}; - } - return project_decode_inputs(hidden_states); -} - torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( const torch::Tensor& hidden_states, const AttentionMetadata& attn_metadata, KVCache& kv_cache, const ModelInputParams& input_params) { - // Early-return on dummy shards. Under dp>1, an empty shard is padded with a - // fake token by worker_impl but its GDN state tensors (kv_cache_tokens_nums, - // linear_state_ids etc.) are left undefined. This mirrors the is_dummy - // early-return in Attention::forward (npu_torch/attention.cpp). Uses - // zeros_like rather than empty_like so downstream post-norm / mlp do not - // read uninitialized data. Placed before FlashComm1 sequence gather so - // dummy shards do not enter the collective and waste bandwidth. - if (attn_metadata.is_dummy) { - return torch::zeros_like(hidden_states); - } - const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context(); - torch::Tensor h = hidden_states; - if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { - h = gather_sequence(hidden_states, *fc1_ctx); - } + auto [qkvz_padded, ba_padded] = + project_padded_inputs(hidden_states, attn_metadata); + int64_t batch_size = qkvz_padded.size(0); + int64_t seq_len = qkvz_padded.size(1); + + torch::Tensor qkvz_flat = + qkvz_padded.view({batch_size * seq_len, qkvz_padded.size(-1)}); + torch::Tensor ba_flat = + ba_padded.view({batch_size * seq_len, ba_padded.size(-1)}); + xllm::kernel::FusedQkvzbaSplitReshapeParams fused_params; + fused_params.mixed_qkvz = qkvz_flat; + fused_params.mixed_ba = ba_flat; + fused_params.num_heads_qk = static_cast(num_k_heads_ / tp_size_); + fused_params.num_heads_v = static_cast(num_v_heads_ / tp_size_); + fused_params.head_qk = static_cast(head_k_dim_); + fused_params.head_v = static_cast(head_v_dim_); - // Save the gathered hidden-state size for potential padding later. - const int64_t original_num_tokens = h.size(0); - const bool use_spec_verify = input_params.is_spec_verify; - const bool is_any_prefill = - attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; torch::Tensor mixed_qkv, z, b, a; - torch::Tensor processed_q, processed_k, processed_v; - int64_t batch_size = 0; - int64_t seq_len = 0; + std::tie(mixed_qkv, z, b, a) = + xllm::kernel::fused_qkvzba_split_reshape_cat(fused_params); - // Qwen3.5 stores qkv, z, b, and a as separate projection weights, so it can - // use their outputs directly in every forward mode. Qwen3Next stores qkvz - // and ba as packed weights and uses the fused-split fallback below. - auto split_inputs = project_split_inputs(h, attn_metadata); - if (split_inputs.has_value()) { - std::tie(mixed_qkv, z, b, a) = split_inputs.value(); - batch_size = mixed_qkv.size(0); - seq_len = mixed_qkv.size(1); - } else { - auto [qkvz_padded, ba_padded] = project_padded_inputs(h, attn_metadata); - batch_size = qkvz_padded.size(0); - seq_len = qkvz_padded.size(1); - - torch::Tensor qkvz_flat = - qkvz_padded.view({batch_size * seq_len, qkvz_padded.size(-1)}); - torch::Tensor ba_flat = - ba_padded.view({batch_size * seq_len, ba_padded.size(-1)}); - xllm::kernel::FusedQkvzbaSplitReshapeParams fused_params; - fused_params.mixed_qkvz = qkvz_flat; - fused_params.mixed_ba = ba_flat; - fused_params.num_heads_qk = static_cast(num_k_heads_ / tp_size_); - fused_params.num_heads_v = static_cast(num_v_heads_ / tp_size_); - fused_params.head_qk = static_cast(head_k_dim_); - fused_params.head_v = static_cast(head_v_dim_); - - std::tie(mixed_qkv, z, b, a) = - xllm::kernel::fused_qkvzba_split_reshape_cat(fused_params); - - mixed_qkv = mixed_qkv.view({batch_size, seq_len, mixed_qkv.size(-1)}); - z = z.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_}); - b = b.view({batch_size, seq_len, num_v_heads_ / tp_size_}); - a = a.view({batch_size, seq_len, num_v_heads_ / tp_size_}); - } - - const bool fla_ssm_state_layout = use_fla_ssm_state_layout(); - const int64_t local_q_heads = num_k_heads_ / tp_size_; - const int64_t local_v_heads = num_v_heads_ / tp_size_; - const int64_t local_conv_dim = - 2 * local_q_heads * head_k_dim_ + local_v_heads * head_v_dim_; - bool used_direct_prefill_qkv = false; + mixed_qkv = mixed_qkv.view({batch_size, seq_len, mixed_qkv.size(-1)}); + z = z.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_}); + b = b.view({batch_size, seq_len, num_v_heads_ / tp_size_}); + a = a.view({batch_size, seq_len, num_v_heads_ / tp_size_}); torch::Tensor conv_cache = kv_cache.get_conv_cache(); torch::Tensor ssm_cache = kv_cache.get_ssm_cache(); - torch::Device device = mixed_qkv.device(); - torch::Tensor conv_weight = conv1d_->weight(); - torch::Tensor logical_state_indices = - get_linear_state_indices(input_params, device); - const int64_t checkpoint_stride = - get_checkpoint_stride(conv_cache, ssm_cache); - torch::Tensor linear_state_base_indices = - build_linear_state_base_indices(logical_state_indices, checkpoint_stride); - auto graph_context = input_params.graph.acl_graph_task_update_context; - const bool register_conv1d_graph_update = - graph_context != nullptr && graph_context->capturing; + torch::Tensor g, beta, core_attn_out, last_recurrent_state; + auto device = mixed_qkv.device(); + auto conv_weight = conv1d_->weight(); + auto linear_state_indices = get_linear_state_indices(input_params, device); - if (!use_spec_verify && is_any_prefill) { - torch::IntArrayRef num_accepted_tokens_opt; - std::vector linear_state_indices_vec( - input_params.embedding.linear_state_ids.begin(), - input_params.embedding.linear_state_ids.end()); - torch::Tensor conv_input = reshape_qkvz_unpad(attn_metadata, mixed_qkv); + if (attn_metadata.is_prefill) { + mixed_qkv = mixed_qkv.transpose(1, 2); + torch::Tensor conv_state = + (seq_len < conv_kernel_size_ - 1) + ? torch::pad(mixed_qkv, {0, conv_kernel_size_ - 1 - seq_len}) + : (seq_len > conv_kernel_size_ - 1) + ? mixed_qkv.narrow( + -1, seq_len - conv_kernel_size_ + 1, conv_kernel_size_ - 1) + : mixed_qkv; + conv_state = conv_state.transpose(1, 2).contiguous(); + conv_cache.index_put_({linear_state_indices}, + conv_state.to(conv_cache.dtype())); + torch::Tensor bias; + auto conv_output = + torch::conv1d(mixed_qkv, + conv_weight.unsqueeze(1).to(device), + bias, + /*stride=*/std::vector{1}, + /*padding=*/std::vector{3}, + /*dilation=*/std::vector{1}, + /*groups=*/static_cast(mixed_qkv.size(1))); + mixed_qkv = torch::silu(conv_output.slice(2, 0, seq_len)); - const bool direct_qkv_model_supported = - fla_ssm_state_layout && num_k_heads_ % tp_size_ == 0 && - num_v_heads_ % tp_size_ == 0 && local_q_heads > 0 && - local_v_heads > 0 && head_k_dim_ == 128 && head_v_dim_ == 128; - const bool direct_qkv_metadata_available = - attn_metadata.q_seq_lens_vec.size() == - static_cast(batch_size) && - input_params.parallel.query_start_loc.size() == - static_cast(batch_size + 1) && - input_params.embedding.linear_state_ids.size() == - static_cast(batch_size) && - input_params.linear_state_validity_mask.size() == - static_cast(batch_size); - int64_t total_valid_tokens = 0; - bool direct_qkv_lengths_valid = direct_qkv_metadata_available; - if (direct_qkv_metadata_available) { - for (const int32_t valid_len : attn_metadata.q_seq_lens_vec) { - direct_qkv_lengths_valid = - direct_qkv_lengths_valid && valid_len >= 0 && valid_len <= seq_len; - total_valid_tokens += valid_len; - } - } - const bool direct_qkv_sequence_supported = - direct_qkv_model_supported && direct_qkv_lengths_valid && - conv_input.dim() == 2 && total_valid_tokens == conv_input.size(0); - const bool direct_qkv_shape_supported = - direct_qkv_sequence_supported && conv_input.size(1) == local_conv_dim && - conv_weight.dim() == 2 && conv_weight.size(0) == 4 && - conv_weight.size(1) == local_conv_dim && conv_cache.dim() == 3 && - conv_cache.size(1) >= 3 && conv_cache.size(2) == local_conv_dim; - const bool direct_qkv_dtype_supported = - direct_qkv_shape_supported && - conv_input.scalar_type() == torch::kBFloat16 && - conv_weight.scalar_type() == torch::kBFloat16 && - conv_cache.scalar_type() == torch::kBFloat16; - const bool use_direct_prefill_qkv = - direct_qkv_dtype_supported && conv_input.is_contiguous() && - conv_weight.is_contiguous() && conv_cache.is_contiguous(); - if (use_direct_prefill_qkv) { - std::tie(processed_q, processed_k, processed_v) = - xllm::kernel::npu::causal_conv1d_qkv( - conv_input, - conv_weight, - conv_cache, - torch::IntArrayRef(input_params.parallel.query_start_loc), - torch::IntArrayRef(linear_state_indices_vec), - torch::IntArrayRef(input_params.linear_state_validity_mask), - local_q_heads, - local_v_heads, - head_k_dim_, - head_v_dim_); - used_direct_prefill_qkv = true; - } else { - mixed_qkv = xllm::kernel::causal_conv1d( - conv_input, - conv_weight, - conv_cache, - std::optional(), // bias (no bias for qwen3) - torch::IntArrayRef(input_params.parallel.query_start_loc), - torch::IntArrayRef(linear_state_indices_vec), - torch::IntArrayRef(input_params.linear_state_validity_mask), - num_accepted_tokens_opt, - xllm::npu::kCausalConv1dActivationSilu, - xllm::npu::kCausalConv1dGraphPadSlotId, - xllm::npu::kCausalConv1dRunModeForward); - - mixed_qkv = reshape_projected_tokens_with_pad(attn_metadata, mixed_qkv); - mixed_qkv = mixed_qkv.transpose(1, 2); - } } else { - if (use_spec_verify) { - CHECK(input_params.num_accepted_tokens.defined()) - << "num_accepted_tokens must be populated for Qwen3.5 spec verify"; - } - torch::Tensor conv_input = reshape_qkvz_unpad(attn_metadata, mixed_qkv); - const auto& num_accepted = use_spec_verify - ? input_params.num_accepted_tokens_host - : std::vector(); - const std::vector linear_state_indices_host( - input_params.embedding.linear_state_ids.begin(), - input_params.embedding.linear_state_ids.end()); - if (register_conv1d_graph_update) { - if (use_spec_verify) { - const auto conv1d_branch = - xllm::npu::CausalConv1dGraphBranch::kSpecVerify; - mixed_qkv = run_causal_conv1d_graph_update( - graph_context, - conv_input, - conv_weight, - conv_cache, - std::optional(), - input_params.parallel.query_start_loc, - linear_state_indices_host, - num_accepted, - conv1d_branch); - } else { - auto conv_input_2d = conv_input.dim() == 3 - ? conv_input.reshape({-1, conv_input.size(-1)}) - : conv_input; - xllm::kernel::CausalConv1dUpdateParams conv1d_params; - conv1d_params.x = conv_input_2d; - conv1d_params.conv_state = conv_cache; - conv1d_params.weight = conv_weight; - conv1d_params.conv_state_indices = logical_state_indices; - conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens; - conv1d_params.max_query_len = attn_metadata.max_query_len; - mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params); - if (conv_input.dim() == 3) { - mixed_qkv = - mixed_qkv.view({conv_input.size(0), -1, mixed_qkv.size(-1)}); - } - } - } else { - if (use_spec_verify) { - torch::Tensor output = torch::empty_like(conv_input); - xllm::kernel::causal_conv1d_out( - output, - conv_input, - conv_weight, - conv_cache, - std::optional(), - torch::IntArrayRef(input_params.parallel.query_start_loc), - torch::IntArrayRef(linear_state_indices_host), - torch::IntArrayRef(std::vector()), - torch::IntArrayRef(num_accepted), - xllm::npu::kCausalConv1dActivationSilu, - xllm::npu::kCausalConv1dGraphPadSlotId, - xllm::npu::kCausalConv1dRunModeUpdate); - mixed_qkv = output; - } else { - auto conv_input_2d = conv_input.dim() == 3 - ? conv_input.reshape({-1, conv_input.size(-1)}) - : conv_input; - xllm::kernel::CausalConv1dUpdateParams conv1d_params; - conv1d_params.x = conv_input_2d; - conv1d_params.conv_state = conv_cache; - conv1d_params.weight = conv_weight; - conv1d_params.conv_state_indices = logical_state_indices; - conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens; - conv1d_params.max_query_len = attn_metadata.max_query_len; - mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params); - if (conv_input.dim() == 3) { - mixed_qkv = - mixed_qkv.view({conv_input.size(0), -1, mixed_qkv.size(-1)}); - } - } - } - mixed_qkv = reshape_projected_tokens_with_pad(attn_metadata, mixed_qkv); + xllm::kernel::CausalConv1dUpdateParams conv1d_params; + conv1d_params.x = mixed_qkv.reshape({-1, mixed_qkv.size(-1)}); + conv1d_params.conv_state = conv_cache; + conv1d_params.weight = conv_weight; + conv1d_params.conv_state_indices = linear_state_indices; + conv1d_params.block_idx_last_scheduled_token = + std::optional(); + conv1d_params.initial_state_idx = std::optional(); + conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens; + conv1d_params.max_query_len = attn_metadata.max_query_len; + mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params); + // Reshape back to 3D [batch_size, dim, seq_len] + mixed_qkv = + mixed_qkv.view({batch_size, -1, mixed_qkv.size(-1)}).contiguous(); mixed_qkv = mixed_qkv.transpose(1, 2); } - const bool use_fused_sigmoid_gdn_decode = - fla_ssm_state_layout && !use_spec_verify && !is_any_prefill && - checkpoint_stride == 1; - torch::Tensor g; - torch::Tensor beta; + // Compute gated delta net decay and beta terms. - if (use_spec_verify || attn_metadata.is_chunked_prefill || - checkpoint_stride > 1) { - beta = torch::sigmoid(b); - torch::Tensor A_log_exp = A_log_.exp(); - torch::Tensor a_float = a.to(torch::kFloat32); - torch::Tensor a_plus_dt = a_float + dt_bias_; - torch::Tensor softplus_out = torch::nn::functional::softplus( - a_plus_dt, - torch::nn::functional::SoftplusFuncOptions().beta(1.0).threshold(20.0)); - g = -A_log_exp * softplus_out; - g = g.to(a.dtype()).contiguous(); - } else if (attn_metadata.is_prefill) { + if (attn_metadata.is_prefill) { xllm::kernel::FusedGdnGatingParams gdn_params; gdn_params.A_log = A_log_; gdn_params.a = a.contiguous().view({-1, a.size(-1)}); @@ -801,7 +413,7 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params); g = g.squeeze(0).contiguous().view({batch_size, seq_len, a.size(-1)}); beta = beta.squeeze(0).contiguous().view({batch_size, seq_len, b.size(-1)}); - } else if (!use_fused_sigmoid_gdn_decode) { + } else { xllm::kernel::FusedGdnGatingParams gdn_params; gdn_params.A_log = A_log_; gdn_params.a = a.view({-1, a.size(-1)}); @@ -811,216 +423,57 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( gdn_params.threshold = 20.0f; std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params); } - if (!used_direct_prefill_qkv) { - std::tie(processed_q, processed_k, processed_v) = - process_mixed_qkv(mixed_qkv); - } - torch::Tensor core_attn_out; - torch::Tensor last_recurrent_state; + auto [processed_q, processed_k, processed_v] = process_mixed_qkv(mixed_qkv); // Apply chunked or recurrent gated-delta attention and update caches. - if (use_spec_verify) { - torch::Tensor spec_num_accepted_tokens = expand_sequence_tensor_to_batch( - input_params.num_accepted_tokens.to(device, torch::kInt32), - batch_size, - "num_accepted_tokens"); - torch::Tensor spec_linear_state_base_indices = - expand_sequence_tensor_to_batch( - linear_state_base_indices, batch_size, "linear_state_base_indices"); - torch::Tensor step_offsets = - torch::arange(seq_len, - torch::TensorOptions() - .dtype(spec_linear_state_base_indices.dtype()) - .device(device)); - torch::Tensor checkpoint_indices = - spec_linear_state_base_indices.unsqueeze(1) + step_offsets; - double scale = 1.0 / std::sqrt(static_cast(processed_q.size(-1))); - core_attn_out = - run_spec_verify_gated_delta_rule(processed_q, - processed_k, - processed_v, - g, - beta, - ssm_cache, - checkpoint_indices, - spec_num_accepted_tokens, - attn_metadata.q_cu_seq_lens, - attn_metadata.q_seq_lens_vec, - scale); - } else if (is_any_prefill) { - CHECK_GE(attn_metadata.q_seq_lens_vec.size(), - static_cast(batch_size)) - << "q_seq_lens_vec must be populated for Qwen3.5 prefill."; - const bool use_single_prefill_pack = - batch_size == 1 && attn_metadata.q_seq_lens_vec.size() == 1 && - attn_metadata.q_seq_lens_vec[0] == seq_len; - torch::Tensor packed_processed_q; - torch::Tensor packed_processed_k; - torch::Tensor packed_processed_v; - torch::Tensor packed_g_tensor; - torch::Tensor packed_beta_tensor; - if (use_single_prefill_pack) { - packed_processed_q = processed_q; - packed_processed_k = processed_k; - packed_processed_v = processed_v; - packed_g_tensor = g; - packed_beta_tensor = beta; - } else { - std::vector packed_q; - std::vector packed_k; - std::vector packed_v; - std::vector packed_g; - std::vector packed_beta; - packed_q.reserve(batch_size); - packed_k.reserve(batch_size); - packed_v.reserve(batch_size); - packed_g.reserve(batch_size); - packed_beta.reserve(batch_size); - for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { - const int64_t valid_len = attn_metadata.q_seq_lens_vec[batch_idx]; - if (!used_direct_prefill_qkv) { - packed_q.emplace_back(processed_q[batch_idx].narrow( - /*dim=*/0, /*start=*/0, valid_len)); - packed_k.emplace_back(processed_k[batch_idx].narrow( - /*dim=*/0, /*start=*/0, valid_len)); - packed_v.emplace_back(processed_v[batch_idx].narrow( - /*dim=*/0, /*start=*/0, valid_len)); - } - packed_g.emplace_back( - g[batch_idx].narrow(/*dim=*/0, /*start=*/0, valid_len)); - packed_beta.emplace_back( - beta[batch_idx].narrow(/*dim=*/0, /*start=*/0, valid_len)); - } - if (used_direct_prefill_qkv) { - packed_processed_q = processed_q; - packed_processed_k = processed_k; - packed_processed_v = processed_v; - } else { - packed_processed_q = torch::cat(packed_q, 0).unsqueeze(0); - packed_processed_k = torch::cat(packed_k, 0).unsqueeze(0); - packed_processed_v = torch::cat(packed_v, 0).unsqueeze(0); - } - packed_g_tensor = torch::cat(packed_g, 0).unsqueeze(0); - packed_beta_tensor = torch::cat(packed_beta, 0).unsqueeze(0); - } - - xllm::kernel::MegaChunkGdnParams mega_chunk_gdn_params; - mega_chunk_gdn_params.q = packed_processed_q; - mega_chunk_gdn_params.k = packed_processed_k; - mega_chunk_gdn_params.v = packed_processed_v; - mega_chunk_gdn_params.g = packed_g_tensor; - mega_chunk_gdn_params.beta = packed_beta_tensor; + if (attn_metadata.is_prefill) { + xllm::kernel::ChunkGatedDeltaRuleParams chunk_gated_delta_params; + chunk_gated_delta_params.q = processed_q; + chunk_gated_delta_params.k = processed_k; + chunk_gated_delta_params.v = processed_v; + chunk_gated_delta_params.g = g; + chunk_gated_delta_params.beta = beta; // Get initial state from ssm_cache for sequences with previous state // Shape: [batch_size, num_heads, head_k_dim, head_v_dim] torch::Tensor initial_state_tensor = - torch::index_select(ssm_cache, 0, linear_state_base_indices); - CHECK_EQ(input_params.linear_state_validity_mask.size(), - input_params.embedding.linear_state_ids.size()) - << "linear state validity mask must be sequence-scoped."; - for (size_t i = 0; i < input_params.linear_state_validity_mask.size(); - ++i) { - if (input_params.linear_state_validity_mask[i] == 0) { - initial_state_tensor.select(0, static_cast(i)).fill_(0.0); - } - } - if (!fla_ssm_state_layout && attn_metadata.is_chunked_prefill) { - initial_state_tensor = - initial_state_tensor.transpose(-1, -2).contiguous(); - } - mega_chunk_gdn_params.initial_state = initial_state_tensor; - mega_chunk_gdn_params.output_final_state = true; - mega_chunk_gdn_params.cu_seqlens = attn_metadata.q_cu_seq_lens; - mega_chunk_gdn_params.q_seq_lens = c10::ArrayRef( - attn_metadata.q_seq_lens_vec.data(), static_cast(batch_size)); - mega_chunk_gdn_params.use_qk_l2norm_in_kernel = !used_direct_prefill_qkv; - torch::Tensor packed_core_attn_out; - std::tie(packed_core_attn_out, last_recurrent_state) = - xllm::kernel::mega_chunk_gdn(mega_chunk_gdn_params); - if (use_single_prefill_pack) { - core_attn_out = packed_core_attn_out; - if (core_attn_out.scalar_type() != processed_v.scalar_type()) { - core_attn_out = core_attn_out.to(processed_v.scalar_type()); - } - } else { - core_attn_out = - used_direct_prefill_qkv - ? torch::zeros({batch_size, seq_len, local_v_heads, head_v_dim_}, - z.options()) - : torch::zeros_like(processed_v); - int64_t packed_offset = 0; - for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { - const int64_t valid_len = attn_metadata.q_seq_lens_vec[batch_idx]; - core_attn_out[batch_idx] - .narrow(/*dim=*/0, /*start=*/0, valid_len) - .copy_(packed_core_attn_out[0].narrow( - /*dim=*/0, packed_offset, valid_len)); - packed_offset += valid_len; - } - } - torch::Tensor state_to_store = fla_ssm_state_layout - ? last_recurrent_state - : last_recurrent_state.transpose(-1, -2); - ssm_cache.index_put_({linear_state_base_indices}, - state_to_store.to(ssm_cache.dtype())); - } else if (checkpoint_stride > 1) { - auto ssm_state = - torch::index_select(ssm_cache, 0, linear_state_base_indices); - if (!fla_ssm_state_layout) { - ssm_state = ssm_state.transpose(-1, -2); - } - ssm_state = ssm_state.contiguous(); + torch::index_select(ssm_cache, 0, linear_state_indices); + // Todo: chunked-prefill/prefix-cache use initial_state + initial_state_tensor.fill_(0.0); + chunk_gated_delta_params.initial_state = initial_state_tensor; + chunk_gated_delta_params.output_final_state = true; + chunk_gated_delta_params.cu_seqlens = attn_metadata.q_cu_seq_lens; + chunk_gated_delta_params.head_first = false; + chunk_gated_delta_params.use_qk_l2norm_in_kernel = true; std::tie(core_attn_out, last_recurrent_state) = - torch_recurrent_gated_delta_rule( - processed_q, processed_k, processed_v, g, beta, ssm_state); - torch::Tensor state_to_store = fla_ssm_state_layout - ? last_recurrent_state - : last_recurrent_state.transpose(-1, -2); - ssm_cache.index_put_({linear_state_base_indices}, - state_to_store.to(ssm_cache.dtype())); + xllm::kernel::chunk_gated_delta_rule(chunk_gated_delta_params); + ssm_cache.index_put_( + {linear_state_indices}, + last_recurrent_state.transpose(-1, -2).to(ssm_cache.dtype())); } else { + processed_q = xllm::kernel::l2_norm(processed_q, 1e-6); + processed_k = xllm::kernel::l2_norm(processed_k, 1e-6); + auto zero = torch::zeros({1}, attn_metadata.q_seq_lens.options()); + torch::Tensor actual_seq_lengths = + torch::cat({zero, attn_metadata.q_seq_lens}, 0); double scale = 1.0 / std::sqrt(static_cast(processed_q.size(-1))); - if (fla_ssm_state_layout) { - xllm::kernel::FusedSigmoidGatingDeltaRuleUpdateParams params; - params.A_log = A_log_.contiguous(); - params.a = a.contiguous(); - params.dt_bias = dt_bias_.contiguous(); - params.q = processed_q.contiguous(); - params.k = processed_k.contiguous(); - params.v = processed_v.contiguous(); - params.b = b.contiguous(); - params.initial_state_source = ssm_cache; - params.initial_state_indices = linear_state_base_indices.contiguous(); - params.cu_seqlens = attn_metadata.q_cu_seq_lens.contiguous(); - params.scale = static_cast(scale); - params.use_qk_l2norm_in_kernel = true; - params.softplus_beta = 1.0f; - params.softplus_threshold = 20.0f; - core_attn_out = - xllm::kernel::fused_sigmoid_gating_delta_rule_update(params); - } else { - processed_q = xllm::kernel::l2_norm(processed_q, /*eps=*/1e-6); - processed_k = xllm::kernel::l2_norm(processed_k, /*eps=*/1e-6); - auto zero = torch::zeros({1}, attn_metadata.q_seq_lens.options()); - torch::Tensor actual_seq_lengths = - torch::cat({zero, attn_metadata.q_seq_lens}, 0); - core_attn_out = xllm::kernel::recurrent_gated_delta_rule( - processed_q.reshape( - {-1, processed_q.size(-2), processed_q.size(-1)}), - processed_k.reshape( - {-1, processed_k.size(-2), processed_k.size(-1)}), - processed_v.reshape( - {-1, processed_v.size(-2), processed_v.size(-1)}), - ssm_cache, - beta.squeeze(0).contiguous(), - scale, - actual_seq_lengths, - logical_state_indices, - c10::nullopt, - g.squeeze(0).contiguous(), - c10::nullopt) - .unsqueeze(0) - .contiguous(); - } + core_attn_out = xllm::kernel::recurrent_gated_delta_rule( + processed_q.reshape( + {-1, processed_q.size(-2), processed_q.size(-1)}), + processed_k.reshape( + {-1, processed_k.size(-2), processed_k.size(-1)}), + processed_v.reshape( + {-1, processed_v.size(-2), processed_v.size(-1)}), + ssm_cache, + beta.squeeze(0).contiguous(), + scale, + actual_seq_lengths, + linear_state_indices, + c10::nullopt, + g.squeeze(0).contiguous(), + c10::nullopt) + .unsqueeze(0) + .contiguous(); } + auto z_reshaped = z.view({-1, z.size(-1)}); auto core_attn_out_reshaped = core_attn_out.view({-1, core_attn_out.size(-1)}); @@ -1033,47 +486,25 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( auto rearranged_norm = norm_out.reshape({norm_out.size(0), norm_out.size(1) * norm_out.size(2)}); rearranged_norm = reshape_qkvz_unpad(attn_metadata, rearranged_norm); - // For chunked prefill or spec verify, reshape_projected_tokens_with_pad may - // pad each batch to max_len, causing output tokens > original_num_tokens. We - // need to slice back to original_num_tokens to match the residual shape. - if (rearranged_norm.size(0) > original_num_tokens) { - // Slice excess padding tokens - rearranged_norm = - rearranged_norm.slice(0, 0, original_num_tokens).contiguous(); - } - if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { - return o_proj_->forward(rearranged_norm, - row_parallel_reduce_mode_for_fc1(*fc1_ctx)); - } - return o_proj_->forward(rearranged_norm); + auto attn_output = o_proj_->forward(rearranged_norm); + return attn_output; } torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_unpad( const AttentionMetadata& attn_metadata, const torch::Tensor& padded_qkvz) const { - const bool has_padded_queries = - attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; - if (!has_padded_queries) { + if (!attn_metadata.is_prefill) { return padded_qkvz; } std::vector valid_batches; - const bool has_host_lens = !attn_metadata.q_seq_lens_vec.empty(); - int64_t bs = has_host_lens - ? static_cast(attn_metadata.q_seq_lens_vec.size()) - : attn_metadata.q_seq_lens.size(0); - valid_batches.reserve(bs); + int64_t bs = attn_metadata.q_seq_lens.size(0); int64_t max_len = attn_metadata.max_query_len; const auto& ori_seq_lens = attn_metadata.q_seq_lens; auto reshaped_qkvz = padded_qkvz.view({bs, max_len, -1}); for (int64_t b = 0; b < bs; ++b) { - int64_t ori_len = has_host_lens ? attn_metadata.q_seq_lens_vec[b] - : ori_seq_lens[b].template item(); - torch::Tensor valid_batch = - reshaped_qkvz[b].slice(/*dim=*/0, /*start=*/0, ori_len); - valid_batches.emplace_back(valid_batch); - } - if (valid_batches.size() == 1) { - return valid_batches[0].contiguous(); + int64_t ori_len = ori_seq_lens[b].template item(); + torch::Tensor valid_batch = reshaped_qkvz[b].slice(0, 0, ori_len); + valid_batches.push_back(valid_batch); } return torch::cat(valid_batches, 0).contiguous(); } @@ -1081,60 +512,41 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_unpad( torch::Tensor Qwen3GatedDeltaNetBaseImpl::get_linear_state_indices( const ModelInputParams& input_params, const torch::Device& device) const { - CHECK(!input_params.embedding.linear_state_ids.empty()) + CHECK(!input_params.linear_state_ids.empty()) << "linear_state_ids must be populated for gated delta net"; - if (input_params.embedding.linear_state_indices.defined()) { - auto indices = input_params.embedding.linear_state_indices; - if (indices.device() != device || indices.scalar_type() != torch::kInt) { - indices = - indices.to(torch::TensorOptions().dtype(torch::kInt).device(device), - /*non_blocking=*/true, - /*copy=*/true); - } - return indices.contiguous(); + if (input_params.linear_state_indices.defined()) { + return input_params.linear_state_indices; } return torch::tensor( - input_params.embedding.linear_state_ids, + input_params.linear_state_ids, torch::TensorOptions().dtype(torch::kInt).device(device)); } -torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_projected_tokens_with_pad( +torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_with_pad( const AttentionMetadata& attn_metadata, - const torch::Tensor& projected_tokens) const { - const bool has_host_lens = !attn_metadata.q_seq_lens_vec.empty(); - int64_t bs = has_host_lens - ? static_cast(attn_metadata.q_seq_lens_vec.size()) - : attn_metadata.q_seq_lens.size(0); + const torch::Tensor& qkvz) const { + int64_t bs = attn_metadata.q_seq_lens.size(0); int64_t max_len = attn_metadata.max_query_len; const auto& start_loc = attn_metadata.q_seq_lens; - const bool need_padding = - attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; - if (!need_padding) { - return projected_tokens.view({bs, -1, projected_tokens.size(-1)}); - } - if (has_host_lens && bs == 1 && attn_metadata.q_seq_lens_vec[0] == max_len && - projected_tokens.dim() == 2 && projected_tokens.size(0) == max_len) { - return projected_tokens.view({1, max_len, projected_tokens.size(-1)}); + if (!attn_metadata.is_prefill) { + return qkvz.view({qkvz.size(0), -1, qkvz.size(-1)}); } std::vector batches; - batches.reserve(bs); int64_t idx = 0; for (int64_t b = 0; b < bs; ++b) { - int64_t cur_len = has_host_lens ? attn_metadata.q_seq_lens_vec[b] - : start_loc[b].template item(); - torch::Tensor batch = - projected_tokens.slice(/*dim=*/0, idx, idx + cur_len).contiguous(); + int64_t cur_len = start_loc[b].template item(); + torch::Tensor batch = qkvz.slice(0, idx, idx + cur_len).contiguous(); idx = idx + cur_len; if (batch.size(0) != max_len) { batch = batch.size(0) > max_len - ? batch.slice(/*dim=*/0, /*start=*/0, max_len).contiguous() + ? batch.slice(0, 0, max_len).contiguous() : torch::nn::functional::pad( batch, torch::nn::functional::PadFuncOptions( {0, 0, 0, max_len - batch.size(0)})) .contiguous(); } - batches.emplace_back(batch); + batches.push_back(batch); } auto ret = torch::stack(batches, 0).contiguous(); return ret; diff --git a/ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.h b/ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.h index fdc82b4d..2994f329 100644 --- a/ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.h +++ b/ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.h @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2026 The xLLM Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -17,7 +17,6 @@ limitations under the License. #include -#include #include #include #include @@ -52,40 +51,19 @@ class Qwen3GatedDeltaNetBaseImpl : public torch::nn::Module { const ModelInputParams& input_params); protected: - virtual std::pair project_decode_inputs( - const torch::Tensor& hidden_states) = 0; - virtual std::pair project_flat_inputs( - const torch::Tensor& hidden_states) = 0; - // Qwen3.5 overrides this to project and reshape its separate qkv/z/b/a - // weights in every forward mode. Qwen3Next keeps qkvz/ba packed and returns - // nullopt to select the fused-split fallback. - virtual std::optional< - std::tuple> - project_split_inputs(const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata) { - return std::nullopt; - } - virtual bool use_fla_ssm_state_layout() const { return false; } + virtual std::pair project_padded_inputs( + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata) = 0; void load_common_state_dict(const StateDict& state_dict); void verify_common_loaded_weights(const std::string& prefix) const; - torch::Tensor get_linear_state_indices(const ModelInputParams& input_params, - const torch::Device& device) const; - - std::pair project_padded_inputs( - const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata); - + torch::Tensor reshape_qkvz_with_pad(const AttentionMetadata& attn_metadata, + const torch::Tensor& qkvz) const; torch::Tensor reshape_qkvz_unpad(const AttentionMetadata& attn_metadata, const torch::Tensor& padded_qkvz) const; - - // Projection outputs are packed as [total_tokens, dim], while GDN kernels - // consume dense [batch, max_query_len, dim] tensors. Split the packed tokens - // by query length and pad each sequence before entering the kernels. - torch::Tensor reshape_projected_tokens_with_pad( - const AttentionMetadata& attn_metadata, - const torch::Tensor& projected_tokens) const; + torch::Tensor get_linear_state_indices(const ModelInputParams& input_params, + const torch::Device& device) const; std::tuple process_mixed_qkv( torch::Tensor& mixed_qkv) const;