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

完整部署快照

Browse Source
This commit is contained in:
root
2026-09-20 02:20:26 +00:00
commit ce553c0dc8
5836 changed files with 1631123 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

229
upstream_ref/tpu-inference/moe.py Normal file
View File

@@ -0,0 +1,229 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from enum import Enum
from typing import TYPE_CHECKING, Tuple, Union
import jax
import jax.numpy as jnp
from jax.sharding import Mesh
from vllm.model_executor.layers.fused_moe import RoutedExperts
from tpu_inference import envs
from tpu_inference.kernels.fused_moe.v1.kernel import fused_ep_moe
from tpu_inference.layers.common.fused_moe_gmm import fused_moe_func
from tpu_inference.logger import init_logger
from tpu_inference.utils import to_jax_dtype
if TYPE_CHECKING:
from tpu_inference.layers.common.process_weights.moe_weights import (
FusedMoEWeights, UnfusedMoEWeights)
from tpu_inference.layers.jax.moe.moe import JaxMoE, JaxRoutedExperts
else:
FusedMoEWeights = None
UnfusedMoEWeights = None
JaxMoE = None
JaxRoutedExperts = None
logger = init_logger(__name__)
class MoEBackend(Enum):
# FUSED_MOE is using the Fused MoE kernel found in tpu-inference/tpu_inference/kernels/fused_moe/v1/kernel.py
# and is production ready.
# NOTE: for FusedMOE on the JAX/Flax path, we expect the MoE quant method (e.g. UnquantizedFusedMoEMethod) to
# create a new `kernel_gating_upproj_E2DF` param that replaces `kernel_gating_EDF` and `kernel_up_proj_EDF`
# and thus runs the forward pass on the fused weight. `kernel_down_proj_EFD` is unchanged.
FUSED_MOE = "fused_moe"
# GMM_EP is using the GMM kernel found in tpu-inference/tpu_inference/layers/common/fused_moe_gmm.py
# and is `expert_sharded_gmm` the GMM calls. Production ready.
# NOTE: for GMM_EP JAX/Flax path, we expect the MoE quant method (e.g. UnquantizedFusedMoEMethod) to
# create a new `kernel_gating_upproj_EDF` param that replaces `kernel_gating_EDF` and `kernel_up_proj_EDF`
# and thus runs the forward pass on the fused weight. `kernel_down_proj_EFD` is unchanged.
GMM_EP = "gmm_ep"
# Same as GMM_EP but uses `tensor_sharded_gmm_*` for the GMM calls and is production ready
GMM_TP = "gmm_tp"
# DENSE_MAT uses a simple dense matmul for the MoE backend, is intended for testing, and is
# only used in the JAX path for now
# NOTE: for DENSE_MAT in the JAX/Flax path, there are no changes for weights for the unfused backends (DENSE_MAT and MEGABLOX_GMM).
# That is, `kernel_gating_EDF`, `kernel_up_proj_EDF`, and `kernel_down_proj_EFD` are unchanged.
DENSE_MAT = "dense_mat"
# Also only used in the JAX path for now
MEGABLX_GMM = "megablox_gmm"
@classmethod
def fused_moe_backends(cls):
"""Returns those backends that use fused weights"""
return {cls.FUSED_MOE, cls.GMM_EP, cls.GMM_TP}
def moe_apply(
layer: Union[RoutedExperts, JaxRoutedExperts, JaxMoE],
x: jax.Array,
gating_output: Union[jax.Array, Tuple[jax.Array, jax.Array]],
weights: Union[FusedMoEWeights, UnfusedMoEWeights],
moe_backend: MoEBackend,
mesh: Mesh,
extra_backend_kwargs: dict,
) -> jax.Array:
extra_backend_kwargs = dict(
extra_backend_kwargs) if extra_backend_kwargs else {}
scatter_results = extra_backend_kwargs.pop("scatter_results", False)
moe_chunk_size = extra_backend_kwargs.pop("moe_chunk_size", 0)
# When set, the GMM backends skip their tensor-/expert-parallel all-reduce
# and return per-shard partial sums, so the reduction can be deferred and
# fused with a later collective (e.g. a shared-expert all-reduce).
defer_all_reduce = extra_backend_kwargs.pop("defer_all_reduce", False)
if defer_all_reduce and moe_backend not in {
MoEBackend.GMM_EP, MoEBackend.GMM_TP
}:
raise ValueError(
"defer_all_reduce can only be True for GMM_EP and GMM_TP backends")
with jax.named_scope(layer._get_name()):
activation = layer.activation if isinstance(
layer.activation, str) else layer.activation.value
if activation == "silu":
swiglu_limit = getattr(layer, "swiglu_limit", None)
if swiglu_limit is not None and swiglu_limit > 0:
activation = "silu_and_mul_with_clamp"
match moe_backend:
case MoEBackend.FUSED_MOE:
subc_quant_w1_sz = None
subc_quant_w2_sz = None
if weights.w13_weight_scale is not None and weights.w2_weight_scale is not None:
padded_hidden_size = weights.w13_weight.shape[-2]
# NB: w13_weight_scale: (num_experts, 2, hidden_size // subc_quant_w1_sz, 1, intermediate_size)
assert padded_hidden_size % weights.w13_weight_scale.shape[
2] == 0
subc_quant_w1_sz = padded_hidden_size // weights.w13_weight_scale.shape[
2]
intermediate_size = weights.w13_weight.shape[-1]
# NB: w2_weight_scale: (num_experts, intermediate_size // subc_quant_w2_sz, 1, hidden_size)
assert intermediate_size % weights.w2_weight_scale.shape[
1] == 0
subc_quant_w2_sz = intermediate_size // weights.w2_weight_scale.shape[
1]
actual_hidden_size = x.shape[-1]
padding_size = weights.w13_weight.shape[-2] - actual_hidden_size
x = jnp.pad(x, ((0, 0), (0, padding_size)))
output = fused_ep_moe(
mesh=mesh,
tokens=x,
w1=weights.w13_weight,
w2=weights.w2_weight,
gating_output=gating_output,
top_k=layer.top_k,
renormalize_topk_logits=layer.renormalize,
act_fn=activation,
scoring_fn=layer.scoring_func,
subc_quant_w1_sz=subc_quant_w1_sz,
subc_quant_w2_sz=subc_quant_w2_sz,
w1_scale=weights.w13_weight_scale,
w2_scale=weights.w2_weight_scale,
b1=weights.w13_bias,
b2=weights.w2_bias,
**extra_backend_kwargs,
)[:, :actual_hidden_size]
case MoEBackend.GMM_EP | MoEBackend.GMM_TP:
# Check if activation_dtype was passed via kwargs or as an environment variable
activation_dtype = extra_backend_kwargs.get(
"activation_dtype", envs.MOE_ALL_GATHER_ACTIVATION_DTYPE)
all_gather_fp8 = (bool(activation_dtype)
and to_jax_dtype(activation_dtype)
== jnp.float8_e4m3fn)
output = fused_moe_func(
hidden_states=x,
w1=weights.w13_weight,
w2=weights.w2_weight,
w1_scale=weights.w13_weight_scale,
w2_scale=weights.w2_weight_scale,
w1_bias=weights.w13_bias,
w2_bias=weights.w2_bias,
gating_output=gating_output,
topk=layer.top_k,
renormalize=layer.renormalize,
mesh=mesh,
use_ep=layer.use_ep,
activation=activation,
scoring_fn=layer.scoring_func,
all_gather_fp8=all_gather_fp8,
enable_rs_kernel=envs.ENABLE_RS_KERNEL,
use_gmm_fused_rs_kernel=envs.USE_GMM_FUSED_RS_KERNEL,
onehot_moe_permute_threshold=envs.
ONEHOT_MOE_PERMUTE_THRESHOLD,
scatter_results=scatter_results,
defer_all_reduce=defer_all_reduce,
hash_based_topk_indices=extra_backend_kwargs.get(
"hash_based_topk_indices", None),
expert_score_correction_bias=extra_backend_kwargs.get(
"e_score_correction_bias", None),
moe_chunk_size=moe_chunk_size,
num_valid_tokens=extra_backend_kwargs.get(
"num_valid_tokens", None),
)
case MoEBackend.DENSE_MAT:
# NOTE: circular import avoidance
from tpu_inference.layers.jax.moe.dense_moe import \
dense_moe_func
assert isinstance(
gating_output,
tuple), "Expected the gating output to be a tuple"
assert len(
gating_output
) == 2, "Expected the gating output to be have 2 entries: weights and indices"
return dense_moe_func(
weights=weights,
x_TD=x,
gating_output=gating_output,
cast_dtype=layer.dtype,
num_local_experts=layer.num_local_experts,
apply_expert_weight_before_computation=layer.
apply_expert_weight_before_computation,
activation_ffw_ted=layer.activation_ffw_ted,
activation_ffw_td=layer.activation_ffw_td,
hidden_act=layer.hidden_act,
mesh=mesh)
case MoEBackend.MEGABLX_GMM:
# NOTE: circular import avoidance
from tpu_inference.layers.jax.moe.sparse_moe import \
sparse_moe_func
return sparse_moe_func(weights=weights,
x_TD=x,
gating_output=gating_output,
layer=layer,
mesh=mesh)
return output
# TODO(#3041): Inherit from vLLM's FusedMoEMethodBase, so it can take FusedMoeConfig
# as init arg, and unify more logic.
class FusedMoEMethodBase:
"""Base class that prepare TPU specific configs"""
def __init__(self, moe_backend: MoEBackend, ep_axis_name: str):
self.extra_backend_kwargs: dict = {
"moe_chunk_size": envs.VLLM_MOE_CHUNK_SIZE
}
if moe_backend == MoEBackend.FUSED_MOE:
self.extra_backend_kwargs["ep_axis_name"] = ep_axis_name