dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

data: port complete MoE + xllm layer call chains from upstream repos

Browse Source 
MoE call chain from ds_vllm (vllm-project/vllm latest):
  ex_engine/moe/ — 20 files, 8736 lines
  - modular_kernel.py (1630 lines) — base classes for modular MoE
  - experts/fused_batched_moe.py (972 lines) — NaiveBatchedExperts
  - prepare_finalize/batched.py (171 lines) — token grouping by expert
  - topk_weight_and_reduce.py (176 lines) — scatter-add finalize
  - fused_moe.py (1740 lines) — main fused_moe dispatch
  - config.py (1407 lines) — FusedMoEQuantConfig
  - activation.py, utils.py, layer.py, etc.

xllm layer code (jd-opensource/xllm):
  ex_engine/xllm_layers/ — 39 files, 5859 lines
  - ilu/fused_moe.cpp (797 lines) — production ixformer 7-step MoE pipeline
  - ilu/attention.cpp (189 lines) — paged_attention + flash_attn bridge
  - npu_torch/qwen3_gated_delta_net_base.cpp (576 lines) — GDN reference
  - common/rms_norm.cpp, rotary_embedding.cpp, activation.cpp, dense_mlp.cpp

xllm ILU kernels — synced 10 files to upstream (diffs from prior edits)

These are reference implementations, NOT hand-written.
Source repos: vllm-project/vllm, jd-opensource/xllm
This commit is contained in:
Claude
2026-08-15 14:26:18 +00:00
parent 7aa5054574
commit 6415249693
47 changed files with 10894 additions and 851 deletions
Download Patch File Download Diff File Expand all files Collapse all files

166
ex_engine/moe/__init__.py
View File

@@ -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")

29
ex_engine/moe/activation.py
View File

@@ -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:

1407
ex_engine/moe/config.py Normal file
View File

File diff suppressed because it is too large Load Diff

0
ex_engine/moe/experts/__init__.py Normal file
View File

170
ex_engine/moe/experts/fallback.py Normal file
View File

@@ -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,
)

972
ex_engine/moe/experts/fused_batched_moe.py Normal file
View File

@@ -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,
)

1740
ex_engine/moe/fused_moe.py Normal file
View File

File diff suppressed because it is too large Load Diff

214
ex_engine/moe/fused_moe_method_base.py Normal file
View File

@@ -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

118
ex_engine/moe/fused_moe_modular_method.py Normal file
View File

@@ -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,
)

406
ex_engine/moe/layer.py Normal file
View File

@@ -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,
)

1630
ex_engine/moe/modular_kernel.py Normal file
View File

File diff suppressed because it is too large Load Diff

192
ex_engine/moe/moe_align_block_size.py Normal file
View File

@@ -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

202
ex_engine/moe/moe_fused_mul_sum.py Normal file
View File

@@ -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

283
ex_engine/moe/moe_permute_unpermute.py Normal file
View File

@@ -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()

29
ex_engine/moe/prepare_finalize/__init__.py Normal file
View File

@@ -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.
]

171
ex_engine/moe/prepare_finalize/batched.py Normal file
View File

@@ -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,
)

141
ex_engine/moe/prepare_finalize/no_dp_ep.py Normal file
View File

@@ -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()
)

176
ex_engine/moe/topk_weight_and_reduce.py Normal file
View File

@@ -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

441
ex_engine/moe/utils.py Normal file
View File

@@ -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)