789 lines
29 KiB
Python
789 lines
29 KiB
Python
|
# Copyright 2026 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.
|
||
|
|
|
||
|
|
import functools
|
||
|
|
from typing import Literal
|
||
|
|
|
||
|
|
import jax
|
||
|
|
from jax import numpy as jnp
|
||
|
|
from jax.sharding import Mesh, NamedSharding
|
||
|
|
from jax.sharding import PartitionSpec as P
|
||
|
|
from tokamax._src.ops.experimental.gmm_v2.gmm_v2 import gmm_v2
|
||
|
|
|
||
|
|
import tpu_inference.envs as envs
|
||
|
|
from tpu_inference.kernels.collectives.hierrs_sc import wrapper as hier_rs_sc
|
||
|
|
from tpu_inference.kernels.sparse_core.dense_gather_reduce import \
|
||
|
|
dense_gather_reduce
|
||
|
|
from tpu_inference.kernels.sparse_core.ragged_gather_reduce_v2 import \
|
||
|
|
ragged_gather_reduce
|
||
|
|
from tpu_inference.kernels.sparse_core.ragged_gather_v2 import \
|
||
|
|
ragged_gather_v2 as ragged_gather
|
||
|
|
from tpu_inference.layers.common.quantization import quantize_tensor
|
||
|
|
from tpu_inference.layers.common.sharding import ShardingAxisName
|
||
|
|
from tpu_inference.logger import init_logger
|
||
|
|
from tpu_inference.utils import get_mesh_shape_product
|
||
|
|
|
||
|
|
logger = init_logger(__name__)
|
||
|
|
|
||
|
|
# Target chunk size of 2048 slots was found empirically to be optimal
|
||
|
|
# for MoE workloads (e.g., Qwen) to hide ICI/DMA latency during AllReduce.
|
||
|
|
TARGET_SLOT_CHUNK_SIZE = 2048
|
||
|
|
|
||
|
|
|
||
|
|
def _override_token_indices_for_random_routing(
|
||
|
|
topk_indices: jax.Array, global_num_experts: int) -> jax.Array:
|
||
|
|
logger.warning(
|
||
|
|
"Forcing random routing should be used for performance testing only.")
|
||
|
|
original_topk_indices = topk_indices
|
||
|
|
num_tokens, topk = original_topk_indices.shape
|
||
|
|
# Forcing random routing is useful to get rid of the effect
|
||
|
|
# of routing imbalance during performance debugging.
|
||
|
|
# (original_topk_indices // global_num_experts) is just zero, but we keep it so that
|
||
|
|
# the all-gather of topk_indices won't be skipped so that the performance comparison between
|
||
|
|
# with and without random routing is fair.
|
||
|
|
rng_key = jax.random.PRNGKey(42)
|
||
|
|
topk_indices = jax.vmap(lambda key: jax.random.choice(
|
||
|
|
key, global_num_experts, shape=(topk, ), replace=False))(
|
||
|
|
jax.random.split(rng_key, num_tokens)) + (original_topk_indices //
|
||
|
|
global_num_experts)
|
||
|
|
return topk_indices
|
||
|
|
|
||
|
|
|
||
|
|
def all_gather_topk_indices_and_weights(
|
||
|
|
topk_indices: jax.Array, topk_weights: jax.Array, dtype: jnp.dtype,
|
||
|
|
mesh: Mesh) -> tuple[jax.Array, jax.Array]:
|
||
|
|
# `topk_indices` and `topk_weights` are relatively small (and last dimension is top-k),
|
||
|
|
# directly all-gather them is inefficient. We use reshape, bitcast to convert the data into one array,
|
||
|
|
# all gather, then unpack.
|
||
|
|
top_k = topk_indices.shape[-1]
|
||
|
|
topk_indices = topk_indices.astype(jnp.int32).reshape(-1)
|
||
|
|
topk_weights = topk_weights.astype(jnp.float32).reshape(-1)
|
||
|
|
topk_weights = jax.lax.bitcast_convert_type(topk_weights,
|
||
|
|
topk_indices.dtype)
|
||
|
|
|
||
|
|
blob = jnp.stack([topk_indices, topk_weights])
|
||
|
|
# The optimization barrier here is to prevent the compiler from reordering the all-gather the operations above.
|
||
|
|
blob = jax.lax.optimization_barrier(blob)
|
||
|
|
gathered_blob = jax.lax.with_sharding_constraint(
|
||
|
|
blob, NamedSharding(mesh, P(None, ShardingAxisName.MLP_DATA)))
|
||
|
|
|
||
|
|
topk_indices = gathered_blob[0]
|
||
|
|
topk_weights = gathered_blob[1]
|
||
|
|
topk_indices = topk_indices.reshape(-1, top_k)
|
||
|
|
topk_weights = jax.lax.bitcast_convert_type(topk_weights, jnp.float32)
|
||
|
|
topk_weights = topk_weights.reshape(-1, top_k).astype(dtype)
|
||
|
|
|
||
|
|
return topk_indices, topk_weights
|
||
|
|
|
||
|
|
|
||
|
|
def apply_scoring_fn(scoring_fn: str, x: jax.Array) -> jax.Array:
|
||
|
|
match scoring_fn:
|
||
|
|
case "softmax":
|
||
|
|
return jax.nn.softmax(x, axis=-1)
|
||
|
|
case "sigmoid":
|
||
|
|
return jax.nn.sigmoid(x)
|
||
|
|
case "sqrtsoftplus":
|
||
|
|
return jnp.sqrt(jax.nn.softplus(x))
|
||
|
|
case _:
|
||
|
|
raise NotImplementedError(
|
||
|
|
f"FusedMoE does not support {scoring_fn} scoring function")
|
||
|
|
|
||
|
|
|
||
|
|
def gmm_wrapper(lhs,
|
||
|
|
rhs,
|
||
|
|
rhs_scale,
|
||
|
|
rhs_bias,
|
||
|
|
group_sizes,
|
||
|
|
group_offset,
|
||
|
|
fuse_act=None,
|
||
|
|
preferred_element_type=None):
|
||
|
|
gmm_res = gmm_v2(
|
||
|
|
lhs=lhs,
|
||
|
|
rhs=rhs,
|
||
|
|
rhs_scale=rhs_scale,
|
||
|
|
rhs_bias=rhs_bias,
|
||
|
|
group_sizes=group_sizes,
|
||
|
|
group_offset=group_offset[0],
|
||
|
|
zero_initialize=False,
|
||
|
|
fuse_act=fuse_act,
|
||
|
|
preferred_element_type=preferred_element_type,
|
||
|
|
)
|
||
|
|
return gmm_res
|
||
|
|
|
||
|
|
|
||
|
|
def valid_rows_mask(batch_size: int, group_sizes: jax.Array,
|
||
|
|
group_start: jax.Array, group_end: jax.Array) -> jax.Array:
|
||
|
|
"""Mask indicating rows processed by current shard."""
|
||
|
|
|
||
|
|
group_sizes_sum = jnp.cumulative_sum(group_sizes, include_initial=True)
|
||
|
|
|
||
|
|
token_start = group_sizes_sum[group_start]
|
||
|
|
token_end = group_sizes_sum[group_end]
|
||
|
|
|
||
|
|
index = jnp.arange(batch_size)
|
||
|
|
return jnp.where(jnp.logical_and(token_start <= index, index < token_end),
|
||
|
|
True, False)
|
||
|
|
|
||
|
|
|
||
|
|
def _permute_tokens_for_chunked_rs(x: jax.Array, dp_size: int,
|
||
|
|
num_chunks: int) -> jax.Array:
|
||
|
|
"""Permutes tokens to correct for chunked reduce-scatter.
|
||
|
|
|
||
|
|
This reordering ensures that when we slice the tokens into chunks and run
|
||
|
|
reduce-scatter on each chunk locally, the resulting scattered outputs
|
||
|
|
end up in the correct non-interleaved global sequence order after concatenation.
|
||
|
|
"""
|
||
|
|
num_tokens = x.shape[0]
|
||
|
|
chunk_size_per_device = num_tokens // (num_chunks * dp_size)
|
||
|
|
|
||
|
|
# Reshape to split the token dimension into (dp_size, num_chunks, chunk_size_per_device)
|
||
|
|
new_shape = (dp_size, num_chunks, chunk_size_per_device) + x.shape[1:]
|
||
|
|
x_reshaped = x.reshape(new_shape)
|
||
|
|
|
||
|
|
# Transpose the dp_size and num_chunks axes (axes 0 and 1)
|
||
|
|
transpose_axes = (1, 0, 2) + tuple(range(3, x_reshaped.ndim))
|
||
|
|
x_transposed = jnp.transpose(x_reshaped, transpose_axes)
|
||
|
|
|
||
|
|
# Flatten the first 3 dimensions back to num_tokens
|
||
|
|
return x_transposed.reshape(x.shape)
|
||
|
|
|
||
|
|
|
||
|
|
def moe_gmm_local(x: jax.Array,
|
||
|
|
w1: jax.Array,
|
||
|
|
w1_scale: jax.Array | None,
|
||
|
|
w1_bias: jax.Array | None,
|
||
|
|
w2: jax.Array,
|
||
|
|
w2_scale: jax.Array | None,
|
||
|
|
w2_bias: jax.Array | None,
|
||
|
|
group_sizes: jax.Array,
|
||
|
|
group_offset: jax.Array,
|
||
|
|
topk_argsort_revert_indices: jax.Array,
|
||
|
|
topk_weights: jax.Array,
|
||
|
|
*,
|
||
|
|
activation: str,
|
||
|
|
topk: int,
|
||
|
|
parallelism: Literal["tp", "ep"],
|
||
|
|
enable_rs_kernel: bool = False,
|
||
|
|
onehot_moe_permute_threshold: int = 0,
|
||
|
|
scatter_results: bool = False,
|
||
|
|
moe_chunk_size: int = 0,
|
||
|
|
defer_all_reduce: bool = False) -> jax.Array:
|
||
|
|
"""Main MoE logic on a local shard can run in TP or EP mode.
|
||
|
|
|
||
|
|
Set parallelism for "tp" or "ep"
|
||
|
|
"""
|
||
|
|
|
||
|
|
assert parallelism in ["tp", "ep"]
|
||
|
|
|
||
|
|
# GMM1 computes x @ (W_up | W_gate) together and activation, output is [tokens,padded_intermediate_size]
|
||
|
|
gmm1_res = gmm_wrapper(
|
||
|
|
x,
|
||
|
|
w1,
|
||
|
|
w1_scale,
|
||
|
|
w1_bias,
|
||
|
|
group_sizes,
|
||
|
|
group_offset,
|
||
|
|
fuse_act=activation,
|
||
|
|
preferred_element_type=x.dtype,
|
||
|
|
)
|
||
|
|
|
||
|
|
# When the parallelism is TP since w2_bias is not sharded, we should only apply bias
|
||
|
|
# once, not applying to every shard. So we set w2_bias to 0 to all shards other than
|
||
|
|
# shard 0. For EP, it is not needed since bias is sharded on leading expert axis.
|
||
|
|
if parallelism == "tp" and w2_bias is not None:
|
||
|
|
shard_id = jax.lax.axis_index(ShardingAxisName.MLP_TENSOR).sum()
|
||
|
|
w2_bias = jnp.where(shard_id == 0, w2_bias, 0)
|
||
|
|
gmm1_res = gmm1_res[:, :w2.shape[1]] # trim to hidden size if padded
|
||
|
|
gmm2_res = gmm_wrapper(gmm1_res, w2, w2_scale, w2_bias, group_sizes,
|
||
|
|
group_offset)
|
||
|
|
|
||
|
|
batch_size = gmm2_res.shape[0]
|
||
|
|
local_group_size = w1.shape[0]
|
||
|
|
|
||
|
|
reduction_axis = (ShardingAxisName.MLP_TENSOR
|
||
|
|
if parallelism == "tp" else ShardingAxisName.EXPERT)
|
||
|
|
|
||
|
|
scatter_axis_size = 1
|
||
|
|
reduce_axes = ()
|
||
|
|
scatter_axes = ()
|
||
|
|
|
||
|
|
if enable_rs_kernel:
|
||
|
|
reduction_axes = reduction_axis if isinstance(
|
||
|
|
reduction_axis, tuple) else (reduction_axis, )
|
||
|
|
for axis in reduction_axes:
|
||
|
|
scatter_axis_size *= jax.lax.axis_size(axis)
|
||
|
|
elif scatter_results:
|
||
|
|
dp_axes = ShardingAxisName.ATTN_DATA
|
||
|
|
if not isinstance(dp_axes, tuple):
|
||
|
|
dp_axes = (dp_axes, )
|
||
|
|
if isinstance(reduction_axis, tuple):
|
||
|
|
reduce_axes = tuple(a for a in reduction_axis if a not in dp_axes)
|
||
|
|
scatter_axes = tuple(a for a in reduction_axis if a in dp_axes)
|
||
|
|
else:
|
||
|
|
reduce_axes = () if reduction_axis in dp_axes else (
|
||
|
|
reduction_axis, )
|
||
|
|
scatter_axes = (
|
||
|
|
reduction_axis, ) if reduction_axis in dp_axes else ()
|
||
|
|
for a in scatter_axes:
|
||
|
|
scatter_axis_size *= jax.lax.axis_size(a)
|
||
|
|
|
||
|
|
# Chunk Size Resolution
|
||
|
|
num_tokens = batch_size // topk
|
||
|
|
is_onehot = (local_group_size < group_sizes.size) and (
|
||
|
|
onehot_moe_permute_threshold > 0
|
||
|
|
and batch_size <= onehot_moe_permute_threshold)
|
||
|
|
if moe_chunk_size > 0 and num_tokens > moe_chunk_size * scatter_axis_size and not is_onehot:
|
||
|
|
actual_chunk_size = moe_chunk_size * scatter_axis_size
|
||
|
|
else:
|
||
|
|
actual_chunk_size = num_tokens
|
||
|
|
|
||
|
|
is_pipelined = actual_chunk_size < num_tokens
|
||
|
|
if is_pipelined and scatter_axis_size > 1:
|
||
|
|
num_chunks = num_tokens // actual_chunk_size
|
||
|
|
topk_weights = _permute_tokens_for_chunked_rs(topk_weights,
|
||
|
|
scatter_axis_size,
|
||
|
|
num_chunks)
|
||
|
|
|
||
|
|
revert_indices_2d = topk_argsort_revert_indices.reshape(
|
||
|
|
num_tokens, topk)
|
||
|
|
revert_indices_2d = _permute_tokens_for_chunked_rs(
|
||
|
|
revert_indices_2d, scatter_axis_size, num_chunks)
|
||
|
|
topk_argsort_revert_indices = revert_indices_2d.flatten()
|
||
|
|
|
||
|
|
if local_group_size < group_sizes.size:
|
||
|
|
mask = valid_rows_mask(
|
||
|
|
gmm1_res.shape[0],
|
||
|
|
group_sizes,
|
||
|
|
group_offset,
|
||
|
|
group_offset + local_group_size,
|
||
|
|
)[topk_argsort_revert_indices].reshape(-1, topk, 1)
|
||
|
|
else:
|
||
|
|
mask = jnp.full((batch_size, ), True).reshape(-1, topk, 1)
|
||
|
|
|
||
|
|
out_list = []
|
||
|
|
|
||
|
|
# Pipelining Loop
|
||
|
|
for start_tok in range(0, num_tokens, actual_chunk_size):
|
||
|
|
end_tok = min(num_tokens, start_tok + actual_chunk_size)
|
||
|
|
start_idx = start_tok * topk
|
||
|
|
end_idx = end_tok * topk
|
||
|
|
|
||
|
|
cur_indices = topk_argsort_revert_indices[start_idx:end_idx]
|
||
|
|
cur_weights = topk_weights[start_tok:end_tok]
|
||
|
|
cur_mask = mask[start_tok:end_tok]
|
||
|
|
|
||
|
|
if local_group_size < group_sizes.size:
|
||
|
|
if onehot_moe_permute_threshold > 0 and batch_size <= onehot_moe_permute_threshold:
|
||
|
|
revert_indices = cur_indices.reshape(-1, topk)
|
||
|
|
onehot = jax.nn.one_hot(revert_indices,
|
||
|
|
batch_size,
|
||
|
|
dtype=gmm2_res.dtype)
|
||
|
|
combine = (onehot * cur_weights[..., None] *
|
||
|
|
cur_mask).sum(axis=1)
|
||
|
|
chunk_hidden = combine @ gmm2_res
|
||
|
|
else:
|
||
|
|
chunk_hidden = ragged_gather_reduce(gmm2_res, cur_indices,
|
||
|
|
cur_weights.reshape(-1),
|
||
|
|
cur_mask.reshape(-1), topk)
|
||
|
|
else:
|
||
|
|
chunk_hidden = dense_gather_reduce(
|
||
|
|
gmm2_res,
|
||
|
|
topk_argsort_revert_indices,
|
||
|
|
topk_weights,
|
||
|
|
topk,
|
||
|
|
)
|
||
|
|
if enable_rs_kernel:
|
||
|
|
rs_out = hier_rs_sc.hierarchical_reduce_scatter_local(
|
||
|
|
chunk_hidden,
|
||
|
|
num_devices=scatter_axis_size,
|
||
|
|
axis_name=reduction_axis)
|
||
|
|
out = rs_out.astype(x.dtype)
|
||
|
|
elif scatter_results:
|
||
|
|
if reduce_axes:
|
||
|
|
chunk_hidden = jax.lax.psum(chunk_hidden,
|
||
|
|
axis_name=reduce_axes)
|
||
|
|
if scatter_axes:
|
||
|
|
out = jax.lax.psum_scatter(chunk_hidden,
|
||
|
|
axis_name=scatter_axes,
|
||
|
|
scatter_dimension=0,
|
||
|
|
tiled=True).astype(x.dtype)
|
||
|
|
else:
|
||
|
|
out = chunk_hidden.astype(x.dtype)
|
||
|
|
else:
|
||
|
|
if not defer_all_reduce:
|
||
|
|
out = jax.lax.psum(chunk_hidden,
|
||
|
|
axis_name=reduction_axis).astype(x.dtype)
|
||
|
|
else:
|
||
|
|
out = chunk_hidden.astype(x.dtype)
|
||
|
|
out_list.append(out)
|
||
|
|
|
||
|
|
return jnp.concatenate(out_list,
|
||
|
|
axis=0) if len(out_list) > 1 else out_list[0]
|
||
|
|
|
||
|
|
|
||
|
|
def tensor_parallel_gmm(
|
||
|
|
x: jax.Array,
|
||
|
|
w1: jax.Array,
|
||
|
|
w1_scale: jax.Array | None,
|
||
|
|
w1_bias: jax.Array | None,
|
||
|
|
w2: jax.Array,
|
||
|
|
w2_scale: jax.Array | None,
|
||
|
|
w2_bias: jax.Array | None,
|
||
|
|
group_sizes: jax.Array,
|
||
|
|
topk_argsort_revert_indices: jax.Array,
|
||
|
|
topk_weights: jax.Array,
|
||
|
|
*,
|
||
|
|
activation: str,
|
||
|
|
topk: int,
|
||
|
|
mesh: Mesh,
|
||
|
|
enable_rs_kernel: bool = False,
|
||
|
|
onehot_moe_permute_threshold: int = 0,
|
||
|
|
scatter_results: bool = False,
|
||
|
|
moe_chunk_size: int = 0,
|
||
|
|
defer_all_reduce: bool = False,
|
||
|
|
) -> jax.Array:
|
||
|
|
data_p_spec = P(ShardingAxisName.MLP_DATA)
|
||
|
|
attn_data_p_spec = P(ShardingAxisName.ATTN_DATA)
|
||
|
|
group_offset = jnp.array([0])
|
||
|
|
|
||
|
|
w1_spec = P(None, None, ShardingAxisName.MLP_TENSOR)
|
||
|
|
w2_spec = P(None, ShardingAxisName.MLP_TENSOR, None)
|
||
|
|
|
||
|
|
w1_scale_spec = (None if w1_scale is None else P(
|
||
|
|
None, None, None, ShardingAxisName.MLP_TENSOR))
|
||
|
|
w1_bias_spec = (None if w1_bias is None else P(
|
||
|
|
None, None, ShardingAxisName.MLP_TENSOR))
|
||
|
|
|
||
|
|
num_blocks = 1 if w2_scale is None else w2_scale.shape[1]
|
||
|
|
w2_scale_spec = (None if num_blocks == 1 else P(
|
||
|
|
None, ShardingAxisName.MLP_TENSOR, None, None))
|
||
|
|
w2_bias_spec = None if w2_bias is None else P(None, None, None)
|
||
|
|
|
||
|
|
if scatter_results:
|
||
|
|
final_out_specs = attn_data_p_spec
|
||
|
|
else:
|
||
|
|
final_out_specs = data_p_spec
|
||
|
|
|
||
|
|
return jax.shard_map(
|
||
|
|
functools.partial(
|
||
|
|
moe_gmm_local,
|
||
|
|
activation=activation,
|
||
|
|
topk=topk,
|
||
|
|
parallelism="tp",
|
||
|
|
enable_rs_kernel=False,
|
||
|
|
onehot_moe_permute_threshold=onehot_moe_permute_threshold,
|
||
|
|
scatter_results=scatter_results,
|
||
|
|
moe_chunk_size=moe_chunk_size,
|
||
|
|
defer_all_reduce=defer_all_reduce,
|
||
|
|
),
|
||
|
|
mesh=mesh,
|
||
|
|
in_specs=(
|
||
|
|
data_p_spec,
|
||
|
|
w1_spec,
|
||
|
|
w1_scale_spec,
|
||
|
|
w1_bias_spec,
|
||
|
|
w2_spec,
|
||
|
|
w2_scale_spec,
|
||
|
|
w2_bias_spec,
|
||
|
|
data_p_spec,
|
||
|
|
P(),
|
||
|
|
data_p_spec,
|
||
|
|
data_p_spec,
|
||
|
|
),
|
||
|
|
out_specs=(final_out_specs),
|
||
|
|
check_vma=False,
|
||
|
|
)(
|
||
|
|
x,
|
||
|
|
w1,
|
||
|
|
w1_scale,
|
||
|
|
w1_bias,
|
||
|
|
w2,
|
||
|
|
w2_scale,
|
||
|
|
w2_bias,
|
||
|
|
group_sizes,
|
||
|
|
group_offset,
|
||
|
|
topk_argsort_revert_indices,
|
||
|
|
topk_weights,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def expert_parallel_gmm(
|
||
|
|
x: jax.Array,
|
||
|
|
w1: jax.Array,
|
||
|
|
w1_scale: jax.Array | None,
|
||
|
|
w1_bias: jax.Array | None,
|
||
|
|
w2: jax.Array,
|
||
|
|
w2_scale: jax.Array | None,
|
||
|
|
w2_bias: jax.Array | None,
|
||
|
|
group_sizes: jax.Array,
|
||
|
|
topk_argsort_revert_indices: jax.Array,
|
||
|
|
topk_weights: jax.Array,
|
||
|
|
*,
|
||
|
|
activation: str,
|
||
|
|
topk: int,
|
||
|
|
mesh: Mesh,
|
||
|
|
enable_rs_kernel: bool = False,
|
||
|
|
onehot_moe_permute_threshold: int = 0,
|
||
|
|
moe_chunk_size: int = 0,
|
||
|
|
scatter_results: bool = False,
|
||
|
|
defer_all_reduce: bool = False,
|
||
|
|
) -> jax.Array:
|
||
|
|
ep_size = get_mesh_shape_product(mesh, ShardingAxisName.EXPERT)
|
||
|
|
ep_p_spec = P(ShardingAxisName.EXPERT)
|
||
|
|
data_p_spec = P(ShardingAxisName.MLP_DATA)
|
||
|
|
ep_data_p_spec = P(ShardingAxisName.EXPERT_DATA)
|
||
|
|
attn_data_p_spec = P(ShardingAxisName.ATTN_DATA)
|
||
|
|
num_experts = w1.shape[0]
|
||
|
|
num_experts_per_shard = num_experts // ep_size
|
||
|
|
group_offset = jnp.arange(0, num_experts, num_experts_per_shard)
|
||
|
|
|
||
|
|
w1_scale_spec = None if w1_scale is None else ep_p_spec
|
||
|
|
w1_bias_spec = None if w1_bias is None else ep_p_spec
|
||
|
|
w2_scale_spec = None if w2_scale is None else ep_p_spec
|
||
|
|
w2_bias_spec = None if w2_bias is None else ep_p_spec
|
||
|
|
|
||
|
|
if scatter_results:
|
||
|
|
final_out_specs = attn_data_p_spec
|
||
|
|
elif enable_rs_kernel:
|
||
|
|
final_out_specs = ep_data_p_spec
|
||
|
|
else:
|
||
|
|
final_out_specs = data_p_spec
|
||
|
|
|
||
|
|
return jax.shard_map(
|
||
|
|
functools.partial(
|
||
|
|
moe_gmm_local,
|
||
|
|
activation=activation,
|
||
|
|
topk=topk,
|
||
|
|
parallelism="ep",
|
||
|
|
onehot_moe_permute_threshold=onehot_moe_permute_threshold,
|
||
|
|
enable_rs_kernel=enable_rs_kernel,
|
||
|
|
scatter_results=scatter_results,
|
||
|
|
moe_chunk_size=moe_chunk_size,
|
||
|
|
defer_all_reduce=defer_all_reduce,
|
||
|
|
),
|
||
|
|
mesh=mesh,
|
||
|
|
in_specs=(
|
||
|
|
data_p_spec,
|
||
|
|
ep_p_spec,
|
||
|
|
w1_scale_spec,
|
||
|
|
w1_bias_spec,
|
||
|
|
ep_p_spec,
|
||
|
|
w2_scale_spec,
|
||
|
|
w2_bias_spec,
|
||
|
|
data_p_spec,
|
||
|
|
ep_p_spec,
|
||
|
|
data_p_spec,
|
||
|
|
data_p_spec,
|
||
|
|
),
|
||
|
|
out_specs=(final_out_specs),
|
||
|
|
check_vma=False,
|
||
|
|
)(
|
||
|
|
x,
|
||
|
|
w1,
|
||
|
|
w1_scale,
|
||
|
|
w1_bias,
|
||
|
|
w2,
|
||
|
|
w2_scale,
|
||
|
|
w2_bias,
|
||
|
|
group_sizes,
|
||
|
|
group_offset,
|
||
|
|
topk_argsort_revert_indices,
|
||
|
|
topk_weights,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _apply_all_gather_fp8(hidden_states: jax.Array, mesh: Mesh,
|
||
|
|
dtype: jnp.dtype) -> jax.Array:
|
||
|
|
logger.info("Apply FP8 all-gather on input of MOE")
|
||
|
|
hidden_states_q, scale = quantize_tensor(
|
||
|
|
jnp.float8_e4m3fn,
|
||
|
|
hidden_states,
|
||
|
|
axis=-1,
|
||
|
|
)
|
||
|
|
# quantize_tensor squeezes the scale if axis is int. We need to expand it back.
|
||
|
|
scale = jnp.expand_dims(scale, -1)
|
||
|
|
|
||
|
|
# Dequantize if needed
|
||
|
|
return jax.shard_map(
|
||
|
|
lambda x, s: (x.astype(jnp.float32) * s).astype(dtype),
|
||
|
|
mesh=mesh,
|
||
|
|
in_specs=(
|
||
|
|
P(ShardingAxisName.MLP_DATA, None),
|
||
|
|
P(ShardingAxisName.MLP_DATA, None),
|
||
|
|
),
|
||
|
|
out_specs=P(ShardingAxisName.MLP_DATA, None),
|
||
|
|
)(hidden_states_q, scale)
|
||
|
|
|
||
|
|
|
||
|
|
@jax.jit(static_argnames=(
|
||
|
|
"topk",
|
||
|
|
"renormalize",
|
||
|
|
"mesh",
|
||
|
|
"use_ep",
|
||
|
|
"activation",
|
||
|
|
"scoring_fn",
|
||
|
|
"all_gather_fp8",
|
||
|
|
"enable_rs_kernel",
|
||
|
|
"use_gmm_fused_rs_kernel",
|
||
|
|
"onehot_moe_permute_threshold",
|
||
|
|
"scatter_results",
|
||
|
|
"moe_chunk_size",
|
||
|
|
"defer_all_reduce",
|
||
|
|
))
|
||
|
|
def fused_moe_func(
|
||
|
|
hidden_states: jax.Array,
|
||
|
|
w1: jax.Array,
|
||
|
|
w2: jax.Array,
|
||
|
|
w1_scale: jax.Array | None,
|
||
|
|
w2_scale: jax.Array | None,
|
||
|
|
w1_bias: jax.Array | None,
|
||
|
|
w2_bias: jax.Array | None,
|
||
|
|
gating_output: jax.Array,
|
||
|
|
topk: int,
|
||
|
|
renormalize: bool,
|
||
|
|
mesh: Mesh,
|
||
|
|
use_ep: bool,
|
||
|
|
activation: str,
|
||
|
|
scoring_fn: str,
|
||
|
|
all_gather_fp8: bool = False,
|
||
|
|
enable_rs_kernel: bool = False,
|
||
|
|
use_gmm_fused_rs_kernel: bool = False,
|
||
|
|
onehot_moe_permute_threshold: int = 0,
|
||
|
|
scatter_results: bool = False,
|
||
|
|
defer_all_reduce: bool = False,
|
||
|
|
hash_based_topk_indices: jax.Array | None = None,
|
||
|
|
expert_score_correction_bias: jax.Array | None = None,
|
||
|
|
moe_chunk_size: int = 0,
|
||
|
|
num_valid_tokens: jax.Array | None = None,
|
||
|
|
) -> jax.Array:
|
||
|
|
"""Route tokens in hidden_states into each experts based on routing.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
hidden_states: [num_tokens, hidden_size]
|
||
|
|
w1: first moe weights [num_experts, hidden_size, intermediate_size * 2]
|
||
|
|
w2: second moe weights [num_experts, intermediate_size, hidden_size]
|
||
|
|
w1_scale: w1 scale [num_experts, num_blocks, 1, intermediate_size * 2]
|
||
|
|
w2_scale: w2 scale [num_experts, num_blocks, 1, hidden_size]
|
||
|
|
w1_bias: optional bias of w1 [num_experts, 1, intermediate_size * 2]
|
||
|
|
w2_bias: optional bias of w2 [num_experts, 1, hidden_size]
|
||
|
|
gating_output: routing information of tokens [num_tokens, num_experts]
|
||
|
|
topk: number of experts to choose per token.
|
||
|
|
renormalize: normalize gating_output.
|
||
|
|
mesh: mesh to perform moe.
|
||
|
|
use_ep: use expert parallelism.
|
||
|
|
activation: activation function to perform on the output of w1.
|
||
|
|
scoring_fn: scoring function to apply on gating_output.
|
||
|
|
enable_rs_kernel: enable custom Hierarchical Reduce-Scatter kernel.
|
||
|
|
use_gmm_fused_rs_kernel: enable fused GMM reduce-scatter kernel.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Output of moe operation [num_tokens, hidden_size]
|
||
|
|
"""
|
||
|
|
|
||
|
|
if use_ep and use_gmm_fused_rs_kernel:
|
||
|
|
from tpu_inference.kernels.experimental.fused_moe.fused_moe_rs import \
|
||
|
|
fused_moe_func_rs
|
||
|
|
logger.info("fused_moe_rs kernel in use")
|
||
|
|
return fused_moe_func_rs(hidden_states=hidden_states,
|
||
|
|
w1=w1,
|
||
|
|
w2=w2,
|
||
|
|
w1_scale=w1_scale,
|
||
|
|
w2_scale=w2_scale,
|
||
|
|
w1_bias=w1_bias,
|
||
|
|
w2_bias=w2_bias,
|
||
|
|
gating_output=gating_output,
|
||
|
|
topk=topk,
|
||
|
|
renormalize=renormalize,
|
||
|
|
mesh=mesh,
|
||
|
|
activation=activation,
|
||
|
|
scoring_fn=scoring_fn,
|
||
|
|
fp8_post_gather=all_gather_fp8)
|
||
|
|
|
||
|
|
num_tokens, hidden_size = hidden_states.shape
|
||
|
|
global_num_experts, padded_hidden_size, _ = w1.shape
|
||
|
|
dtype = hidden_states.dtype
|
||
|
|
|
||
|
|
assert (num_tokens * topk) % 16 == 0, (
|
||
|
|
"The kernel requires num_tokens * topk to be a multiple of "
|
||
|
|
f"16 but got {num_tokens}*{topk}={num_tokens*topk}")
|
||
|
|
|
||
|
|
assert gating_output.shape == (num_tokens, global_num_experts)
|
||
|
|
|
||
|
|
topk_weights = apply_scoring_fn(scoring_fn, gating_output)
|
||
|
|
if hash_based_topk_indices is not None:
|
||
|
|
topk_indices = hash_based_topk_indices
|
||
|
|
topk_weights = jnp.take_along_axis(topk_weights, topk_indices, axis=-1)
|
||
|
|
elif envs.MOE_APPROX_TOPK:
|
||
|
|
topk_weights, topk_indices = jax.lax.approx_max_k(
|
||
|
|
topk_weights,
|
||
|
|
k=topk,
|
||
|
|
recall_target=envs.MOE_APPROX_TOPK_RECALL_TARGET)
|
||
|
|
else:
|
||
|
|
if expert_score_correction_bias is not None:
|
||
|
|
_, topk_indices = jax.lax.top_k(
|
||
|
|
topk_weights + expert_score_correction_bias[None, :], k=topk)
|
||
|
|
topk_weights = jnp.take_along_axis(topk_weights,
|
||
|
|
topk_indices,
|
||
|
|
axis=-1)
|
||
|
|
else:
|
||
|
|
topk_weights, topk_indices = jax.lax.top_k(topk_weights, k=topk)
|
||
|
|
if renormalize:
|
||
|
|
topk_weights = topk_weights / topk_weights.sum(axis=-1, keepdims=True)
|
||
|
|
# Route padding tokens to expert 0 instead of picking a selected expert. This
|
||
|
|
# is especially useful when we have a low number of tokens (e.g. low
|
||
|
|
# concurrency), where padding tokens may activate unnecessary expert weights
|
||
|
|
# and slow down the gmm kernel.
|
||
|
|
if num_valid_tokens is not None:
|
||
|
|
token_valid = (jnp.arange(num_tokens) < num_valid_tokens)[:, None]
|
||
|
|
topk_indices = jnp.where(token_valid, topk_indices, 0)
|
||
|
|
topk_weights = jnp.where(token_valid, topk_weights, 0.0)
|
||
|
|
# All gathering topk_indices and topk_weights if attention dp is used.
|
||
|
|
if get_mesh_shape_product(mesh, ShardingAxisName.ATTN_DATA) > 1:
|
||
|
|
topk_indices, topk_weights = all_gather_topk_indices_and_weights(
|
||
|
|
topk_indices, topk_weights, dtype, mesh)
|
||
|
|
topk_weights = topk_weights.astype(dtype)
|
||
|
|
topk_weights = jax.lax.with_sharding_constraint(
|
||
|
|
topk_weights, NamedSharding(mesh, P(ShardingAxisName.MLP_DATA, None)))
|
||
|
|
|
||
|
|
# Only enable Reduce-Scatter if flag is on and Attention is pure DP
|
||
|
|
total_num_devices = mesh.devices.size
|
||
|
|
is_attn_dp = get_mesh_shape_product(
|
||
|
|
mesh, ShardingAxisName.ATTN_DATA) == total_num_devices
|
||
|
|
actual_enable_rs_kernel = enable_rs_kernel and is_attn_dp
|
||
|
|
|
||
|
|
if envs.FORCE_MOE_RANDOM_ROUTING:
|
||
|
|
logger.warning(
|
||
|
|
"Forcing random routing should be used for performance testing only."
|
||
|
|
)
|
||
|
|
# Forcing random routing is useful to get rid of the effect
|
||
|
|
# of routing imbalance during performance debugging.
|
||
|
|
topk_indices = _override_token_indices_for_random_routing(
|
||
|
|
topk_indices, global_num_experts)
|
||
|
|
|
||
|
|
def _process_tokens_locally(hidden_states_local, topk_indices_local):
|
||
|
|
num_tokens_local = hidden_states_local.shape[0]
|
||
|
|
topk_indices_flat = topk_indices_local.flatten()
|
||
|
|
topk_argsort_indices = jnp.argsort(topk_indices_flat)
|
||
|
|
token_indices = jnp.arange(num_tokens_local,
|
||
|
|
dtype=jnp.int32).repeat(topk)
|
||
|
|
token_indices_sorted = token_indices[topk_argsort_indices]
|
||
|
|
# Below one_hot is equivalent to jnp.bincount(topk_indices_flat,
|
||
|
|
# length=global_num_experts) but is more performant.
|
||
|
|
group_sizes_local = jax.nn.one_hot(topk_indices_flat,
|
||
|
|
global_num_experts,
|
||
|
|
dtype=jnp.int32).sum(axis=0)
|
||
|
|
topk_argsort_revert_indices = jnp.argsort(topk_argsort_indices)
|
||
|
|
|
||
|
|
if use_ep:
|
||
|
|
num_ep_shard = get_mesh_shape_product(mesh,
|
||
|
|
ShardingAxisName.EXPERT)
|
||
|
|
local_num_experts = global_num_experts // num_ep_shard
|
||
|
|
shard_idx = jax.lax.axis_index(ShardingAxisName.EXPERT)
|
||
|
|
|
||
|
|
experts_start = shard_idx * local_num_experts
|
||
|
|
experts_end = experts_start + local_num_experts
|
||
|
|
group_offsets = jnp.cumulative_sum(group_sizes_local,
|
||
|
|
include_initial=True)
|
||
|
|
shard_output_start = group_offsets[experts_start]
|
||
|
|
shard_output_end = group_offsets[experts_end]
|
||
|
|
num_tokens = token_indices_sorted.shape[0]
|
||
|
|
if num_tokens <= onehot_moe_permute_threshold:
|
||
|
|
# Use one-hot matmul for permutation, which can be faster
|
||
|
|
# for small batch size
|
||
|
|
onehot = jax.nn.one_hot(token_indices_sorted,
|
||
|
|
hidden_states_local.shape[0],
|
||
|
|
dtype=hidden_states_local.dtype)
|
||
|
|
x = onehot @ hidden_states_local
|
||
|
|
else:
|
||
|
|
x = ragged_gather(
|
||
|
|
hidden_states_local,
|
||
|
|
token_indices_sorted,
|
||
|
|
shard_output_start,
|
||
|
|
shard_output_end,
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
x = hidden_states_local[token_indices_sorted]
|
||
|
|
|
||
|
|
return x, group_sizes_local, topk_argsort_revert_indices
|
||
|
|
|
||
|
|
if all_gather_fp8:
|
||
|
|
hidden_states = _apply_all_gather_fp8(hidden_states, mesh, dtype)
|
||
|
|
|
||
|
|
x, group_sizes, topk_argsort_revert_indices = jax.shard_map(
|
||
|
|
_process_tokens_locally,
|
||
|
|
mesh=mesh,
|
||
|
|
in_specs=(
|
||
|
|
P(ShardingAxisName.MLP_DATA, None),
|
||
|
|
P(ShardingAxisName.MLP_DATA, None),
|
||
|
|
),
|
||
|
|
out_specs=(
|
||
|
|
P(ShardingAxisName.MLP_DATA),
|
||
|
|
P(ShardingAxisName.MLP_DATA),
|
||
|
|
P(ShardingAxisName.MLP_DATA),
|
||
|
|
),
|
||
|
|
check_vma=False,
|
||
|
|
)(hidden_states, topk_indices)
|
||
|
|
|
||
|
|
try:
|
||
|
|
x = jnp.pad(x, ((0, 0), (0, padded_hidden_size - hidden_size)))
|
||
|
|
except Exception as e:
|
||
|
|
raise ValueError(
|
||
|
|
f"Error when padding input hidden states from {hidden_size} to {padded_hidden_size}."
|
||
|
|
) from e
|
||
|
|
|
||
|
|
if use_ep:
|
||
|
|
x = expert_parallel_gmm(
|
||
|
|
x,
|
||
|
|
w1,
|
||
|
|
w1_scale,
|
||
|
|
w1_bias,
|
||
|
|
w2,
|
||
|
|
w2_scale,
|
||
|
|
w2_bias,
|
||
|
|
group_sizes,
|
||
|
|
topk_argsort_revert_indices,
|
||
|
|
topk_weights,
|
||
|
|
activation=activation,
|
||
|
|
topk=topk,
|
||
|
|
mesh=mesh,
|
||
|
|
enable_rs_kernel=actual_enable_rs_kernel,
|
||
|
|
onehot_moe_permute_threshold=onehot_moe_permute_threshold,
|
||
|
|
scatter_results=scatter_results,
|
||
|
|
moe_chunk_size=moe_chunk_size,
|
||
|
|
defer_all_reduce=defer_all_reduce,
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
x = tensor_parallel_gmm(
|
||
|
|
x,
|
||
|
|
w1,
|
||
|
|
w1_scale,
|
||
|
|
w1_bias,
|
||
|
|
w2,
|
||
|
|
w2_scale,
|
||
|
|
w2_bias,
|
||
|
|
group_sizes,
|
||
|
|
topk_argsort_revert_indices,
|
||
|
|
topk_weights,
|
||
|
|
activation=activation,
|
||
|
|
topk=topk,
|
||
|
|
mesh=mesh,
|
||
|
|
enable_rs_kernel=actual_enable_rs_kernel,
|
||
|
|
onehot_moe_permute_threshold=onehot_moe_permute_threshold,
|
||
|
|
scatter_results=scatter_results,
|
||
|
|
moe_chunk_size=moe_chunk_size,
|
||
|
|
defer_all_reduce=defer_all_reduce,
|
||
|
|
)
|
||
|
|
|
||
|
|
return x[:num_tokens, :hidden_size]
|