feat(vllm): support Decode Context Parallelism (DCP) for Qwen3.5 in MaxTextForCausalLM

[sierraisland]

Description

Adds support for Decode Context Parallelism (DCP, --decode-context-parallel-size) when serving hybrid/MoE models (such as Qwen3.5-397B-A17B) via MaxTextForCausalLM on vLLM / tpu-inference.
Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

For models with a small number of KV heads (e.g., Qwen3.5-397B-A17B with only 2 KV heads on 8 TPU chips), sharding KV heads beyond 2 ways requires duplicating KV heads across devices. Under Decode Context Parallelism (DCP, e.g., tp=1, ep=2, dcp=4), the 2 KV heads are sharded 2-way across ('model', 'expert') while the KV cache sequence/context dimension is sharded 4-way across 'dcp', and Q heads, GDN heads, MoE experts, and vocabulary remain sharded 8-way across ('model', 'expert', 'dcp') / ('attn_dp_expert', 'expert', 'dcp').
Previously, enabling --decode-context-parallel-size > 1 with MaxTextForCausalLM failed because:

1. vllm.yml logical axis rules did not include 'dcp' on Q/GDN/MoE/MLP/vocab axes.
2. Attention.forward_serve_vllm always called single-device rpa_ops instead of dispatching to tpu_inference's dcp_forward when mesh.shape["dcp"] > 1.
3. RoutedMoE.__init__ only included "attn_dp_expert" in _expert_parallelism_name when enable_dp_attention=True, causing an axis mismatch with tpu-inference's ShardingAxisName.EXPERT when enable_dp_attention=False.
4. Qwen3NextGatedDeltaNet hardcoded KV_HEAD (['model', 'expert'], 2-way under DCP) for activation PartitionSpecs while GDN weights and Mamba state caches were sharded 8-way across gdn_head (['model', 'expert', 'dcp']).
5. patch_kv_cache_manager in adapter.py used ShardingAxisName.ATTN_HEAD (which includes 'dcp') instead of ShardingAxisName.KV_HEAD to compute num_kv_heads padding for full-attention KV cache specs.

Tests

Verified end-to-end serving and multi-turn agentic benchmarking on tpu7x-8 with Qwen/Qwen3.5-397B-A17B-FP8 (MaxTextForCausalLM), comparing attndp4 (enable_dp_attention=True, tp=1, ep=8) vs. dcp4 (--decode-context-parallel-size 4, tp=4, ep=2):

Mode	Groups / Concurrency	Output Tok/s	Total Tok/s	Turns/s	TPOT Avg (p50)	Turns 2+ TTFT Avg (p50)
attndp4	g=1, c=1	344.98	48,460.07	2.32	43.37 ms (40.28 ms)	410.20 ms (322.66 ms)
dcp4	g=1, c=1	442.76 (+28.3%)	62,158.49 (+28.3%)	2.98 (+28.3%)	33.10 ms (30.35 ms)	302.54 ms (225.83 ms)
attndp4	g=2, c=2	472.53	62,059.79	3.10	63.07 ms (58.57 ms)	462.31 ms (376.22 ms)
dcp4	g=2, c=2	488.46 (+3.4%)	64,114.83 (+3.3%)	3.21 (+3.4%)	32.59 ms (28.75 ms)	4481.93 ms (4291.22 ms)

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 introduces support for context parallelism (dcp) in vLLM inference within MaxText. It updates the logical axis rules and data sharding configurations to include the dcp axis, and integrates dcp_forward in the attention layer's vLLM serving path. Additionally, it refines the expert parallelism name configuration in MoE layers, adjusts the KV cache specification for hybrid Qwen3 models to correctly handle attention layers alongside Mamba layers, and dynamically determines the GDN head axis based on the attention configuration in Qwen3. There are no review comments, and I have no feedback to provide.

[codecov]

Codecov Report

❌ Patch coverage is 82.35294% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/moe.py	50.00%	2 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!