Forward tpu-inference MoE kernel knobs and surface MoE padding

[khatwanimohit]

Description

Two fixes to the vLLM serving path for MoE models.

1. RoutedMoE.fused_moe_matmul calls tpu-inference's fused_moe_func without the
   environment-backed knobs that tpu-inference passes on its own serving path, so
   ENABLE_RS_KERNEL, USE_GMM_FUSED_RS_KERNEL, ONEHOT_MOE_PERMUTE_THRESHOLD and
   VLLM_MOE_CHUNK_SIZE are silently ignored when serving a MaxText model. Forward them.

   ONEHOT_MOE_PERMUTE_THRESHOLD matters beyond tuning: it picks the one-hot permute path
   over the SparseCore ragged_gather_reduce kernel, which doesn't fit v5p. Before this,
   any expert_parallelism > 1 died on the first request with
   AssertionError: num_row_partitions=16 must be <= num_simd_lanes=8.

2. The adapter pads moe_intermediate_size to satisfy GMM_v2's tiling constraint, which
   inflates every expert weight (moe_mlp_tp_size always equals --tensor-parallel-size,
   so TP=4 pads 512->1024 and TP=8 pads 512->2048). It was logged at absl INFO, which vLLM's
   logging doesn't surface, so serving Qwen3.5-35B-A3B quietly used 2x the MoE weights with
   nothing in the logs. Log it at WARNING with the cost.

No behaviour change with no env vars set, and the padding logic itself is untouched.

Tests

Qwen3.5-35B-A3B on v5p-8, bf16 (v5p has no fp8 MXU). Server:

export HF_HOME=... USE_MOE_EP_KERNEL=0 ATTN_BUCKETIZED_NUM_REQS=true \
  ATTN_CUSTOM_NUM_REQS_BUCKETS=4 ONEHOT_MOE_PERMUTE_THRESHOLD=32768 \
  VLLM_MOE_CHUNK_SIZE=256 SLICE_ROPE_CACHE=1 DP_SCHED_BATCH_PREFILL=false \
  NEW_MODEL_DESIGN=1 NUM_PRECOMPILE_WORKERS=8 SKIP_JAX_PRECOMPILE=1 \
  VLLM_ENABLE_V1_MULTIPROCESSING=0

vllm serve Qwen/Qwen3.5-35B-A3B \
  --max-model-len=65536 --max-num-batched-tokens=2048 --max-num-seqs=256 \
  --no-enable-prefix-caching --gpu-memory-utilization=0.9 --tensor-parallel-size=4 \
  --async-scheduling --port=8000 --language-model-only --enable-auto-tool-choice \
  --tool-call-parser=qwen3_coder --reasoning-parser=qwen3 \
  --default-chat-template-kwargs '{"enable_thinking": false}' \
  '--limit-mm-per-prompt={"image": 0, "video": 0}' \
  --hf-overrides '{"architectures": ["MaxTextForCausalLM"]}' \
  "--additional_config={\"sharding\": {\"sharding_strategy\": $SHARDING}, \"maxtext_config\": {\"model_name\": \"qwen3.5-35b-a3b\", \"load_parameters_path\": \"gs://maxtext-model-checkpoints/qwen3.5-35b-a3b/unscanned/0/items\", \"scan_layers\": false, \"weight_dtype\": \"bfloat16\", \"attention\": \"vllm_rpa\", \"enable_nnx\": true, \"pure_nnx_decoder\": true, \"allow_split_physical_axes\": true, \"use_multimodal\": false, \"prefuse_moe_weights\": true}}" \
  --block-size=256 --enable-chunked-prefill

with

# before
SHARDING='{"enable_dp_attention": true, "attn_dp_size": 2}'
# after  (crashed on first request before this PR)
SHARDING='{"enable_dp_attention": false, "tensor_parallelism": 1, "expert_parallelism": 4}'

Client, same for every run:

python tpu-inference/scripts/vllm/benchmarking/agentic_benchmark/benchmark_agentic.py \
  --model Qwen/Qwen3.5-35B-A3B --model-path-or-id Qwen/Qwen3.5-35B-A3B \
  --trace-file gs://wenxindong-vm/rl/mlperf2026/agentic_benchmark/gbs1024_trace_file.jsonl \
  --global-prefix-len 6476 --num-groups 1 --concurrency 1 --group-size 8

240/240 turns in every run, identical output tokens (37,194), so wall clock is comparable.
Native vLLM (MODEL_IMPL_TYPE=vllm, same flags) included for reference.

	before	after	native vLLM
weights/device	33.92 GiB	18.31 GiB	16.95 GiB
wall clock	1450.86 s	1056.98 s	877.64 s
tok/s/chip	838.80	1151.31	1322.14
output tok/s	25.64	35.19	42.38
TPOT avg	282.74 ms	160.73 ms	184.82 ms
TPOT p99	1125.54 ms	603.23 ms	893.00 ms
turn-1 prefill	112,390 ms	144,246 ms	99,449 ms
turns 2+ TTFT	4,923 ms	3,855 ms	4,699 ms

+37% tok/s/chip, and the gap to native goes from 1.58x to 1.15x. Decode is now ahead of
native; what's left is turn-1 prefill, since tensor_parallelism=1 leaves the prefill GEMMs
unsharded. TP=2 / EP=2 is also padding-free (18.12 GiB/device) and is the obvious follow-up
— it was equally unrunnable before this PR.

The padding warning on the before config, where previously nothing was logged:

WARNING:absl:Padding moe_intermediate_size from 512 to 1024 to match MLP MoE requirements
(moe_mlp_tp_size=4, 2*num_lanes=256). This multiplies the MoE weights and MoE FLOPs by 2x.

Note --enable-expert-parallel on the CLI doesn't affect the JAX mesh — tpu-inference reads
expert parallelism only from additional_config, which is why it's set there above.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request updates the vLLM adapter to log a warning when padding moe_intermediate_size due to its memory and FLOP costs. It also forwards environment-backed kernel configuration variables from tpu_inference to fused_moe_func in moe.py to ensure consistent kernel behavior when serving MaxText models through vLLM. The review feedback suggests using getattr with fallback defaults when accessing these environment variables to prevent potential AttributeError crashes with older versions of the tpu_inference package.

[codecov]

Codecov Report

❌ Patch coverage is 20.00000% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...t/integration/vllm/maxtext_vllm_adapter/adapter.py	0.00%	4 Missing ⚠️

📢 Thoughts on this report? Let us know!