Enable PCP for GDN models

[wenxindongwork]

Summary

Makes Qwen3.5-397B (hybrid GDN) run end-to-end under PCP (MODEL_IMPL_TYPE=vllm).

Design: PCP replicates tokens throughout MoE, GDN, and splits the tokens dimension in Full Attention. I.e. PCP acts just like TP for all layers other than full attention where it shares the tokens instead of heads.

Performance

Qwen3.5-397B-A17B-FP8, fp8 KV cache, single request, 4k chunk, v7x-8. Prefill latency via examples/tpu_profiling.py (2 warmup iters + 1 profiled pass = pure execution).

Context	pcp4 x tp2	tp8	attn_dp4 x tp2
256k	15.1 s	15.7 s	43.4 s
1M	109.3 s	123.1 s	—

Non-PCP parity vs main

To confirm this stack does not regress non-PCP shardings, an online-serving benchmark on tp8 + attn_dp + EP (enable_dp_attention:true), 512 prompts, 8k in / 1k out, concurrency 256, was run on this branch vs origin/main.

Metric	branch	main	Δ
Total token throughput (tok/s)	25,418	25,350	+0.3%
Output token throughput (tok/s)	2,824	2,817	+0.3%
Request throughput (req/s)	2.76	2.75	+0.4%
Median TTFT (ms)	4,446	4,449	~0
P99 TTFT (ms)	48,221	48,267	~0
Median TPOT (ms)	77.96	78.11	−0.2%
Median ITL (ms)	38.50	38.73	−0.6%
Median E2EL (ms)	83,444	83,738	−0.4%

Quality

Qwen3.5-35B-A3B-FP8, gsm8k (full 1319 unless noted), fp8 KV, 8 chips

Config	gsm8k
tp8 / attn_dp8	0.96
attn_dp4 x tp2	0.97
pcp8	0.80

pcp8 = 0.80 is pre-existing bug before this PR, which will be resolved in a separate PR.

Repro

The full commands each config was run with. All need PYTHONPATH pointing at the PCP-patched vLLM (SingleTypeKVCacheManager.block_size *= pcp_world_size + the scheduler passing the real pcp_world_size).

1. Prefill latency (Performance table)

examples/tpu_profiling.py — SKIP_JAX_PRECOMPILE=1 + 2 warmup iters ⇒ the profiled pass is pure execution. 256k (shown for pcp4xtp2):

SKIP_JAX_PRECOMPILE=1 PROFILE_SINGLE_DEVICE=1 PYTHON_TRACER_LEVEL=0 \
MODEL_IMPL_TYPE=vllm NEW_MODEL_DESIGN=1 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 NUM_PRECOMPILE_WORKERS=8 VLLM_ENGINE_READY_TIMEOUT_S=7200 \
LIBTPU_INIT_ARGS=' --xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=false --xla_tpu_ars_combiner_threshold_in_bytes=0 --xla_tpu_enable_async_collective_merger=false --xla_tpu_check_legacy_constraints_in_reduce_scatter_legalizer=false' \
python examples/tpu_profiling.py \
  --model=Qwen/Qwen3.5-397B-A17B-FP8 \
  --max-model-len=262144 --max-num-batched-tokens=4096 --max-num-seqs=1 \
  --gpu-memory-utilization=0.9 --language-model-only --kv-cache-dtype=fp8 \
  --enable-expert-parallel --block-size=256 --enable-chunked-prefill \
  --input-len 262134 --output-len 10 --batch-size 1 --num-iters-warmup 2 --num-iters 1 \
  --profile-result-dir ./prof \
  --tensor-parallel-size=2 --prefill-context-parallel-size=4

Swap the last line for the other configs:

tp8            --tensor-parallel-size=8
attn_dp4xtp2   --tensor-parallel-size=8 --additional_config='{"sharding":{"sharding_strategy":{"enable_dp_attention":true,"attn_dp_size":4}}}'
pcp8           --tensor-parallel-size=8 --prefill-context-parallel-size=8

1M: add VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 to the env, set --gpu-memory-utilization=0.95 --max-model-len=1048576 --input-len 1048566, and add --hf-overrides='{"rope_scaling":{"rope_type":"yarn","factor":4.0,"original_max_position_embeddings":262144},"max_position_embeddings":1048576}' (YaRN ×4 makes 1M positions valid; latency-only).

2. Quality (gsm8k)

Serve Qwen3.5-35B-A3B-FP8, then lm_eval. Clear ~/.cache/vllm/xla_cache first after any libtpu change (a stale cache silently garbles fp8-KV). Shown for pcp8 (--tensor-parallel-size=1 --prefill-context-parallel-size=8):

SKIP_JAX_PRECOMPILE=1 MODEL_IMPL_TYPE=vllm NEW_MODEL_DESIGN=1 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 NUM_PRECOMPILE_WORKERS=8 VLLM_ENGINE_READY_TIMEOUT_S=7200 \
LIBTPU_INIT_ARGS=' --xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=false --xla_tpu_ars_combiner_threshold_in_bytes=0 --xla_tpu_enable_async_collective_merger=false --xla_tpu_check_legacy_constraints_in_reduce_scatter_legalizer=false' \
vllm serve Qwen/Qwen3.5-35B-A3B-FP8 \
  --max-model-len=65536 --max-num-batched-tokens=2048 --max-num-seqs=1 \
  --enable-prefix-caching --gpu-memory-utilization=0.9 --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}' --kv-cache-dtype=fp8 \
  --enable-expert-parallel --block-size=256 --enable-chunked-prefill \
  --tensor-parallel-size=1 --prefill-context-parallel-size=8

lm_eval --model local-chat-completions \
  --model_args base_url="http://localhost:8000/v1/chat/completions",model="Qwen/Qwen3.5-35B-A3B-FP8",eos_string="<|im_end|>",num_concurrent=1,max_retries=3,tokenized_requests=False \
  --tasks gsm8k --gen_kwargs max_gen_toks=2048 --apply_chat_template --fewshot_as_multiturn

Other configs: tp8 = --tensor-parallel-size=8; attn_dp8 = --tensor-parallel-size=8 --additional_config='{"sharding":{"sharding_strategy":{"enable_dp_attention":true,"attn_dp_size":8}}}' (bump num_concurrent for the DP configs).

3. Non-PCP parity (online serving)

Serve Qwen3.5-397B-FP8 (tp8 + attn_dp + EP), precompiled at startup (no SKIP_JAX_PRECOMPILE), then benchmark_serving.py. Run once per checkout (branch vs origin/main) with vLLM held constant:

MODEL_IMPL_TYPE=vllm NEW_MODEL_DESIGN=1 USE_GMM_FUSED_RS_KERNEL=1 USE_BATCHED_RPA_KERNEL=1 \
ATTN_BUCKETIZED_NUM_REQS=true ATTN_CUSTOM_NUM_REQS_BUCKETS=4,8,16,32,64 ATTN_BUCKETS=4,8,16,32,64 \
MIN_TOKEN_BUCKET=8 ONEHOT_MOE_PERMUTE_THRESHOLD=32768 MOE_ROUTE_PADDING_TO_EXPERT0=1 \
SLICE_ROPE_CACHE=1 DP_SCHED_BATCH_PREFILL=true VLLM_ENGINE_READY_TIMEOUT_S=7200 \
vllm serve Qwen/Qwen3.5-397B-A17B-FP8 \
  --max-model-len=9236 --max-num-batched-tokens=2048 --max-num-seqs=64 \
  --no-enable-prefix-caching --gpu-memory-utilization=0.9 --tensor-parallel-size=8 \
  --async-scheduling --port=8000 --language-model-only \
  --enable-auto-tool-choice --tool-call-parser=qwen3_coder --reasoning-parser=qwen3 \
  '--limit-mm-per-prompt={"image": 0, "video": 0}' --kv-cache-dtype=fp8 \
  --enable-expert-parallel --block-size=256 \
  '--additional_config={"sharding": {"sharding_strategy": {"enable_dp_attention": true}}}'

python benchmark_serving.py --model Qwen/Qwen3.5-397B-A17B-FP8 --backend vllm \
  --dataset-name random --random-input-len=8192 --random-output-len=1024 \
  --num-prompts=512 --random-range-ratio=1.0 --ignore-eos --max-concurrency=256 \
  --num-warmups=0 --percentile-metrics=ttft,tpot,itl,e2el --use-chat-template \
  --save-result --result-dir ./result

[github-actions]

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a Github issue, please include a link, e.g.,:
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[bhuvanpkaruturi]

Can you please run the quality/accuracy test at pcp=4,tp=2 (where rest of the performance was benchmarked) to ensure we are not regressing even from the 0.80 score on pcp=8

[wenxindongwork]

pcp=4,tp=2 also gives 0.832 lm_eval result