[FEAT] Context parallel indexer for GLM-5.x

[minhpham-duc]

Purpose

This PR splits the DSA sparse-attention indexer prefill work across tensor-parallel (TP) ranks on ROCm/AITER. It adds two new env vars to control it.

The problem. In rocm_aiter_sparse_attn_indexer, every TP rank does the same full work. Each rank computes rocm_fp8_mqa_logits and top-k over all M query rows. The work is copied on every rank, so TP does not make the indexer faster.

The fix. We split the rows across ranks. Rank r takes rows [block + r*stripe, block + (r+1)*stripe) for each block of tp_world_size * stripe rows. So each rank does only part of the work.

Why stripes and not simple chunks. With a causal mask, later rows cost more than earlier rows. If we gave each rank one big chunk, rank 0 would get all the cheap rows and rank 7 would get all the expensive rows. That is not balanced. With stripes, every rank gets a mix of cheap and expensive rows. So each rank does about M²/2/tp work.

How ranks combine results. We use all_reduce(MAX) on topk_indices_buffer. Rows that a rank did not write are still -1, because the buffer is reset to -1 at the start of the call. So MAX picks the real value and gives the full result. The all_reduce runs once per prefill, outside the chunk loop.

New env vars:

var	default	meaning
VLLM_ROCM_USE_AITER_CP_INDEXER	False	turn the split on
VLLM_ROCM_USE_AITER_CP_INDEXER_STRIPE_SIZE	512	rows per stripe

We skip the split and use the old path when tp_world_size == 1 or chunk_m < tp_world_size.

Related work

vLLM #46103 [ROCm][Perf] TP-shard Lightning indexer prefill

The same problem on the vLLM side: on ROCm/gfx950 the DSA (Lightning) indexer prefill, rocm_fp8_mqa_logits plus per-row top-k, runs replicated on every TP rank and becomes the dominant, TP-non-scaling term at long context. vllm-project#46103 shards it along the query-token (M) dimension as well, gated on tp_size > 1 and num_prefill_tokens >= 512 * tp_size, behind VLLM_ROCM_USE_AITER_MLA_SPARSE_TP_SHARD (default on) with an optional ..._FOLD mode. It reports, on 8x MI355X TP8, a 2.1x throughput gain and ~60% lower TTFT for GLM-5.1-FP8 at 199K tokens, +40% throughput for DeepSeek-V4-Pro at 512K, and a neutral ±2-3% at ≤8K context. The PR is still open (needs rebase) at the time of writing.

Test Plan

Hardware: 8x MI355X (gfx950, 288 GB each), ROCm 7.2.3. Model: GLM-5.2-MXFP4 (GlmMoeDsaForCausalLM, index_topk=2048), TP=8.

Three configurations, each a fresh server boot:

• before = main @ 8c34a3723. The env var does not exist on that branch.
• after (stripe 512) = this branch, VLLM_ROCM_USE_AITER_CP_INDEXER=1, VLLM_ROCM_USE_AITER_CP_INDEXER_STRIPE_SIZE=512.
• after (stripe 1024) = same, with VLLM_ROCM_USE_AITER_CP_INDEXER_STRIPE_SIZE=1024.

We benchmarked two _STRIPE_SIZE values, 512 and 1024, because different stripe sizes can yield different performance results.

We ran the baseline on main rather than on this branch with the flag set to 0, so the baseline cannot accidentally include any of this code.

All three use the same vllm serve command:

vllm serve amd/GLM-5.2-MXFP4 \
  --kv-cache-dtype fp8_e4m3 --trust-remote-code --tensor-parallel-size 8 \
  --linear-backend aiter --moe-backend aiter \
  --tool-call-parser glm47 --enable-auto-tool-choice --reasoning-parser glm45 \
  --no-enable-prefix-caching --port 8586

Env vars for all runs: VLLM_USE_V2_MODEL_RUNNER=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1 VLLM_ROCM_USE_AITER_FP8BMM=0 VLLM_ROCM_USE_AITER_FP4BMM=0.

We turned prefix caching off with --no-enable-prefix-caching. So every request does a full prefill. We checked this two ways: the engine log shows enable_prefix_caching=False, and every result has total_input_tokens == num_prompts * ISL.

Benchmark command for each point:

vllm bench serve --backend openai --base-url http://127.0.0.1:8586 \
  --endpoint /v1/completions --model amd/GLM-5.2-MXFP4 \
  --tokenizer amd/GLM-5.2-MXFP4 --trust-remote-code \
  --dataset-name random --random-input-len $ISL --random-output-len 1024 \
  --max-concurrency $CONC --num-prompts $((CONC*10)) --num-warmups $((CONC*2)) \
  --request-rate inf --ignore-eos --save-result \
  --percentile-metrics ttft,tpot,e2el,itl --metric-percentiles 0,75,90,99,100

Test points: ISL/OSL 8192/1024 and 51200/1024, each at concurrency 4, 8, 16, 32, 64, plus 128 and 256 for 8192. 12 points per configuration, 36 runs in total.

Three correctness evals were run against the served model (the server config above is unchanged; these are all HTTP clients).

gsm8k — full 1319-question test set via lm-eval, 5-shot, greedy, thinking off:

lm_eval --model local-chat-completions \
  --model_args '{"model":"amd/GLM-5.2-MXFP4","base_url":"http://localhost:8586/v1/chat/completions","trust_remote_code":true,"eos_string":"</s>","max_retries":1,"num_concurrent":64,"timeout":1800,"tokenized_requests":false,"max_length":262144}' \
  --tasks gsm8k --include_path ./gsm8k --apply_chat_template \
  --gen_kwargs '{"max_tokens":16384,"temperature":0,"top_p":1,"chat_template_kwargs":{"thinking":false}}'

GPQA-Diamond — 198 questions via sglang's run_eval (NeMo-Skills mcq prompt + eval_mcq grader), thinking on:

python3 -m sglang.test.run_eval --eval-name gpqa \
  --model amd/GLM-5.2-MXFP4 \
  --host 127.0.0.1 --port 8586 --num-threads 32 --max-tokens 100000 \
  --chat-template-kwargs '{"enable_thinking": true, "think_end_token": "</think>"}'

RULER / NIAH — niah_single_2 at 64k and 128k via lm-eval, thinking on:

lm_eval --model local-chat-completions \
  --model_args '{"model":"amd/GLM-5.2-MXFP4","base_url":"http://localhost:8586/v1/chat/completions","tokenizer":"amd/GLM-5.2-MXFP4","trust_remote_code":true,"num_concurrent":16,"timeout":36000,"max_retries":1,"tokenized_requests":false,"max_length":262144}' \
  --tasks niah_single_2 \
  --metadata '{"max_seq_lengths":[65536,131072]}' \
  --gen_kwargs '{"max_tokens":8192,"temperature":0,"chat_template_kwargs":{"enable_thinking":true,"think_end_token":"</think>"}}' \
  --apply_chat_template --batch_size auto --seed 1234 --log_samples --output_path ./niah

RULER needs tokenizer= because it sizes its own prompts to exact token counts. timeout is aiohttp's total per-request budget including time queued behind num_concurrent, so it must exceed the whole run's wall time, not just one request.

Test Result

Correctness

All three evals were run on both configurations.

GPQA-Diamond (198 questions, sglang.test.run_eval, thinking enabled, --max-tokens 100000):

config	score	truncated_rate	stop_rate	error_rate
baseline (main)	0.621	0.343	0.657	0.0
context parallel indexer (stripe 512)	0.596	0.379	0.621	0.0

The 0.025 gap is 5 questions out of 198 and is not a regression. Truncation moved by a similar amount in the same direction: 68 baseline responses against 75 indexer responses were cut off at the token cap and scored zero, a 7-question difference. At a ~35% truncation rate this eval is partly measuring whether the reasoning fits in the budget, so treat the table as a like-for-like comparison at a fixed budget rather than an absolute GPQA score.

gsm8k (1319 questions, 5-shot, greedy, thinking off):

filter	baseline (main)	context parallel indexer (stripe 512)
exact_match, strict-match	96.82% (1277/1319)	97.19% (1282/1319)
exact_match, flexible-extract	96.89% (1278/1319)	97.42% (1285/1319)

RULER / NIAH (niah_single_2, 500 documents per length, thinking on):

sequence length	baseline (main)	context parallel indexer (stripe 512)
65536	0.990	0.996
131072	0.992	0.982

Performance

"before" = baseline (main), "after" = this branch with the context parallel indexer on. In every table below, Δ is signed so that positive always means the context parallel indexer is better — for throughput (higher is better) and for TTFT/TPOT (lower is better) alike.

Results with _STRIPE_SIZE = 512

ISL 8192 / OSL 1024 — Output throughput (tok/s)

Conc	before	after	Δ
4	273.6	273.5	-0.04%
8	467.6	466.5	-0.25%
16	737.3	737.2	-0.01%
32	1129.0	1132.1	+0.27%
64	1511.5	1515.3	+0.25%
128	1978.1	1980.9	+0.14%
256	2359.3	2359.8	+0.02%

ISL 8192 / OSL 1024 — Mean TTFT (ms)

Conc	before	after	Δ
4	923	924	-0.10%
8	1514	1459	+3.67%
16	2054	1737	+15.41%
32	2486	2473	+0.50%
64	3084	3110	-0.83%
128	4111	4098	+0.30%
256	6224	6367	-2.30%

ISL 8192 / OSL 1024 — Mean TPOT (ms)

Conc	before	after	Δ
4	13.7	13.7	-0.03%
8	15.6	15.7	-0.61%
16	19.7	20.0	-1.56%
32	25.9	25.9	+0.25%
64	39.3	39.2	+0.33%
128	60.7	60.6	+0.13%
256	102.4	102.2	+0.16%

ISL 51200 / OSL 1024 — Output throughput (tok/s)

Conc	before	after	Δ
4	178.7	181.3	+1.51%
8	249.3	254.9	+2.25%
16	311.2	319.7	+2.74%
32	363.4	375.0	+3.18%
64	375.3	389.0	+3.66%

ISL 51200 / OSL 1024 — Mean TTFT (ms)

Conc	before	after	Δ
4	5862	5420	+7.54%
8	7405	6702	+9.50%
16	8382	7502	+10.50%
32	10337	9854	+4.68%
64	27511	26473	+3.77%

ISL 51200 / OSL 1024 — Mean TPOT (ms)

Conc	before	after	Δ
4	16.7	16.8	-0.60%
8	24.9	24.8	+0.08%
16	43.2	42.7	+1.18%
32	77.9	75.7	+2.87%
64	142.2	137.2	+3.51%

Results with _STRIPE_SIZE = 1024

ISL 8192 / OSL 1024 — Output throughput (tok/s)

Conc	before	after	Δ
4	273.6	275.4	+0.66%
8	467.6	472.3	+1.00%
16	737.3	739.3	+0.28%
32	1129.0	1139.0	+0.88%
64	1511.5	1530.5	+1.26%
128	1978.1	1992.2	+0.71%
256	2359.3	2379.5	+0.85%

ISL 8192 / OSL 1024 — Mean TTFT (ms)

Conc	before	after	Δ
4	923	926	-0.30%
8	1514	1456	+3.82%
16	2054	2051	+0.12%
32	2486	2453	+1.30%
64	3084	3091	-0.22%
128	4111	4064	+1.14%
256	6224	6091	+2.14%

ISL 8192 / OSL 1024 — Mean TPOT (ms)

Conc	before	after	Δ
4	13.7	13.6	+0.72%
8	15.6	15.5	+0.72%
16	19.7	19.7	+0.29%
32	25.9	25.7	+0.83%
64	39.3	38.8	+1.35%
128	60.7	60.3	+0.67%
256	102.4	101.6	+0.77%

ISL 51200 / OSL 1024 — Output throughput (tok/s)

Conc	before	after	Δ
4	178.7	183.6	+2.76%
8	249.3	256.2	+2.75%
16	311.2	321.5	+3.32%
32	363.4	377.7	+3.93%
64	375.3	391.2	+4.24%

ISL 51200 / OSL 1024 — Mean TTFT (ms)

Conc	before	after	Δ
4	5862	5441	+7.18%
8	7405	7018	+5.23%
16	8382	7654	+8.69%
32	10337	9395	+9.12%
64	27511	26341	+4.25%

ISL 51200 / OSL 1024 — Mean TPOT (ms)

Conc	before	after	Δ
4	16.7	16.5	+1.14%
8	24.9	24.4	+1.94%
16	43.2	42.3	+2.18%
32	77.9	75.5	+3.09%
64	142.2	136.5	+4.06%

Stripe-size comparison, ISL 51200 / OSL 1024 / conc 4

_STRIPE_SIZE	Out tok/s	Δ vs baseline	Mean TTFT ms	Mean TPOT ms
baseline (off)	178.7	—	5862	16.68
512	182.5	+2.13%	5600	16.46
1024	183.7	+2.85%	5572	16.34

Summary:

• ISL 8k: no change. Within about ±0.3% at stripe 512, and +0.3% to +1.3% at stripe 1024. 8k prefills are not indexer-bound, so there is nothing to win here.
• ISL 50k: faster at every point. +1.5% to +3.7% at stripe 512, +2.7% to +4.2% at stripe 1024. TTFT and TPOT improve too.
• Stripe 1024 is better than 512 at every 50k point.

Treat differences under about 1% as noise. The 50k gain holds across five concurrency levels and both stripe sizes.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing (anything written below this line will be removed by GitHub Actions)