230 lines
10 KiB
Python
230 lines
10 KiB
Python
# 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
|