data: port complete MoE + xllm layer call chains from upstream repos
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
214
ex_engine/moe/fused_moe_method_base.py
Normal file
214
ex_engine/moe/fused_moe_method_base.py
Normal 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
|
||||
Reference in New Issue
Block a user