6415249693
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
119 lines
3.6 KiB
Python
119 lines
3.6 KiB
Python
# 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,
|
|
)
|