From 76cd098e665ecd7f570c70b5dd0df9894fe57d73 Mon Sep 17 00:00:00 2001 From: khatwanimohit Date: Tue, 1 Sep 2026 16:18:29 +0000 Subject: [PATCH] Forward tpu-inference MoE kernel knobs and surface MoE padding When a MaxText model is served through vLLM, RoutedMoE.fused_moe_matmul calls tpu-inference's fused_moe_func with only the required arguments. tpu-inference's own serving path passes several environment-backed knobs as well, so the two paths run the same kernel with different configurations, and env vars that operators set on the vLLM command line are silently ignored on the MaxText path. Forward the four environment-backed knobs: ENABLE_RS_KERNEL, USE_GMM_FUSED_RS_KERNEL, ONEHOT_MOE_PERMUTE_THRESHOLD and VLLM_MOE_CHUNK_SIZE. ONEHOT_MOE_PERMUTE_THRESHOLD matters beyond tuning: it selects the one-hot permute path over the SparseCore ragged_gather_reduce kernel, which is what lets expert parallelism run at all on TPU generations whose SparseCore has fewer SIMD lanes than that kernel needs. On v5p, expert_parallelism > 1 previously failed with "num_row_partitions=16 must be <= num_simd_lanes=8". Also make two silent sharding surprises visible in the vLLM adapter. Raise the MoE padding message from absl INFO to WARNING and include its cost. vLLM's logging configuration does not surface absl INFO records, so the padding was invisible: serving Qwen3.5-35B-A3B at moe_mlp_tp_size=4 pads moe_intermediate_size 512 -> 1024 and doubles the MoE weights (measured 33.92 GiB/device instead of 16.95 GiB/device), and at moe_mlp_tp_size=8 it quadruples them, with no indication in the logs. Warn when --enable-expert-parallel is passed but the mesh has no expert shards. The native tpu-inference model paths derive use_ep from parallel_config.enable_expert_parallel, while the mesh takes expert parallelism only from additional_config's sharding_strategy and MaxText derives use_ep from the mesh. The same command line therefore runs expert-parallel MoE natively and tensor-parallel MoE under MaxText, which is the configuration that triggers the padding above. The warning points at expert_parallelism as the fix. Bump the post-training pins to pick these up. USE_GMM_FUSED_RS_KERNEL does not exist in the previously pinned tpu-inference (7ecc401e, Jul 27), and neither does the use_gmm_fused_rs_kernel argument of fused_moe_func, so the change above needs a newer commit to import at all: tpu-inference 7ecc401e -> b67ae5f8 (main) vllm 0ba2aa35 -> d626108b (tpu-inference .buildkite/vllm_lkg.version) Both the old and new tpu-inference pin jax==0.11.0, jaxlib==0.11.0 and libtpu==0.0.44, which already match tpu_post_train_overrides.txt, so the jax version in CI is unchanged. The post-training lock is regenerated with seed-env using the same JAX seed commit recorded for it in docs (52d5cb38). Regeneration needs one fix to run: base requirements ask for an unpinned llguidance, which now resolves to 1.8.0, while vLLM requires >=1.7.0,<1.8.0 in both the old and new commits. Bound it to <1.8.0. The functionally required outcome of the regeneration is tokamax>=0.0.13 (up from 0.0.12, where the GMM_v2 kernel lives) and huggingface-hub>=1.29.0; the lock is installed with --resolution=lowest and the GitHub deps with --no-deps, so these floors are exactly what CI installs. --- .../extra_deps/post_train_github_deps.txt | 4 +-- .../vllm/maxtext_vllm_adapter/adapter.py | 28 +++++++++++++++++-- src/maxtext/layers/moe.py | 13 +++++++++ 3 files changed, 41 insertions(+), 4 deletions(-) diff --git a/src/dependencies/extra_deps/post_train_github_deps.txt b/src/dependencies/extra_deps/post_train_github_deps.txt index cdc56f4053..437d934bbe 100644 --- a/src/dependencies/extra_deps/post_train_github_deps.txt +++ b/src/dependencies/extra_deps/post_train_github_deps.txt @@ -1,3 +1,3 @@ google-tunix @ https://github.com/google/tunix/archive/c4ec573d29e4c3a3955b348256d464b119c8a6d1.zip -tpu-inference @ https://github.com/vllm-project/tpu-inference/archive/7ecc401e6faefe3f793f3e4a8e6e2dbc49eca868.zip -vllm @ git+https://github.com/vllm-project/vllm@0ba2aa35a81dcc3246b26291368b53fa2389c7d7 +tpu-inference @ https://github.com/vllm-project/tpu-inference/archive/b67ae5f8f234fd559cf4b376840e7bb9ae6d3275.zip +vllm @ git+https://github.com/vllm-project/vllm@d626108b1841888ec90aced33367149a6bbc7e4b diff --git a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py index cd842eccac..ce062c11f8 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py +++ b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py @@ -128,6 +128,23 @@ def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters f"num_lanes={num_lanes}, tp={tp}, attn_dp={attn_dp}, ep={ep}, moe_mlp_tp_size={moe_mlp_tp_size}" ) + # The native tpu-inference model paths derive use_ep from + # parallel_config.enable_expert_parallel, but the mesh built by + # ShardingConfigManager.from_vllm_config takes expert parallelism only from + # additional_config's sharding_strategy. MaxText derives use_ep from the mesh, so + # --enable-expert-parallel alone leaves the experts unsharded here while the native + # implementation of the same model runs expert-parallel. Warn rather than fail, since + # the run is still correct, only sharded differently than the flag suggests. + expert_shard_degree = ep * sharding_config.attn_dp_expert_size + if vllm_config.parallel_config.enable_expert_parallel and expert_shard_degree == 1: + max_logging.warning( + "--enable-expert-parallel was requested but the mesh has no expert shards " + f"(expert={ep}, attn_dp_expert={sharding_config.attn_dp_expert_size}), so MaxText will shard the MoE over " + "the MLP dimension instead of the expert dimension. The vLLM flag does not reach the JAX mesh; set " + 'expert_parallelism in the vLLM additional_config sharding_strategy, e.g. \'{"sharding": ' + '{"sharding_strategy": {"expert_parallelism": }}}\', to actually shard experts.' + ) + # Replicate the number of KV heads if its less than the total degree of model parallelism if kv_tp_size % num_kv_heads == 0 and num_kv_heads < kv_tp_size: max_logging.log( @@ -152,8 +169,15 @@ def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters while (padded_hidden_size // moe_mlp_tp_size) < (2 * num_lanes): padded_hidden_size = next_power_of_two(padded_hidden_size + 1) - max_logging.log( - f"Padding moe_intermediate_size from {hidden_size} to {padded_hidden_size} to match MLP MoE requirements." + # This inflates every expert weight, so it is a real memory/FLOP cost rather than a + # cosmetic reshape: at moe_mlp_tp_size=4 a 512-wide MoE is padded to 1024 (2x the MoE + # weights), and at moe_mlp_tp_size=8 to 2048 (4x). Log it at WARNING so it is visible + # under vLLM's logging configuration, which does not surface absl INFO records. + max_logging.warning( + f"Padding moe_intermediate_size from {hidden_size} to {padded_hidden_size} to match MLP MoE requirements " + f"(moe_mlp_tp_size={moe_mlp_tp_size}, 2*num_lanes={2 * num_lanes}). This multiplies the MoE weights and " + f"MoE FLOPs by {padded_hidden_size / hidden_size:g}x. Consider sharding the MoE over the expert axis " + f"instead, by setting expert_parallelism in the vLLM additional_config sharding_strategy." ) overrides["padded_base_moe_mlp_dim"] = padded_hidden_size diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 82d4b01e47..de70870da2 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -3290,6 +3290,7 @@ def fused_moe_matmul( # pylint: disable=import-outside-toplevel # pytype: disable=import-error from tpu_inference.layers.common.fused_moe_gmm import fused_moe_func + from tpu_inference import envs as tpu_inference_envs except ImportError as e: raise ImportError("fused_moe_matmul requires the tpu-inference package.") from e @@ -3333,6 +3334,18 @@ def fused_moe_matmul( use_ep=use_ep, activation=activation, scoring_fn=scoring_fn, + # Forward the same environment-backed kernel knobs that tpu-inference passes on its + # own serving path (tpu_inference/layers/common/moe.py). Without these, env vars such + # as ONEHOT_MOE_PERMUTE_THRESHOLD and VLLM_MOE_CHUNK_SIZE are silently ignored when a + # MaxText model is served through vLLM, so the two paths run the same kernel with + # different configurations. In particular, ONEHOT_MOE_PERMUTE_THRESHOLD selects the + # one-hot permute path instead of the SparseCore ragged_gather_reduce kernel, which + # is what makes expert parallelism usable on TPU generations whose SparseCore has + # fewer SIMD lanes than the kernel requires (e.g. v5p). + enable_rs_kernel=tpu_inference_envs.ENABLE_RS_KERNEL, + use_gmm_fused_rs_kernel=tpu_inference_envs.USE_GMM_FUSED_RS_KERNEL, + onehot_moe_permute_threshold=tpu_inference_envs.ONEHOT_MOE_PERMUTE_THRESHOLD, + moe_chunk_size=tpu_inference_envs.VLLM_MOE_CHUNK_SIZE, ) # Reshape output 2D [T, D] -> 3D [B, S, D]