Enable context parallel in HIP

[PiDinosauR2804]

Purpose

This PR adds an opt-in Indexer M-split for ROCm, behind SGLANG_DSA_INDEXER_M_SPLIT (default off), that shards the indexer's query rows (the M dimension) across attn-TP ranks:

1. Query rows are cut into stripes of SGLANG_DSA_INDEXER_M_SPLIT_STRIPE rows (default 512), and stripe i is owned by rank i % tp_size. Interleaving rather than contiguous blocks matters because the workload is causal: late rows score against far more keys than early rows. With interleaved stripes every rank's total stays at roughly M^2 / 2 / tp_size.
2. Each rank computes logits and top-k only for the stripes it owns, using the same aiter fp8_mqa_logits kernel and the same topk_transform per-row-range contract the existing chunked path already uses, and writes its indices into the shared topk_result buffer, which is pre-filled with -1.
3. One AllReduce(MAX) over topk_result[:M] merges the partial results. Rows a rank did not compute are -1, so MAX recovers the owner's indices and every rank ends up with the full top-k, bit-identical to the replicated result.

Per rank, the indexer logits GEMM and top-k shrink by ~tp_size; the added cost is one M x 2048 x int32 all-reduce per layer (256 MiB at M = 32768, on the same collective path as the rest of the TP layer). The K-cache store, the K-cache gather, the decode path (_get_topk_paged), and every non-ROCm platform are unchanged.

What changed

• python/sglang/srt/environ.py: SGLANG_DSA_INDEXER_M_SPLIT (EnvBool, default False) and SGLANG_DSA_INDEXER_M_SPLIT_STRIPE (EnvInt, default 512), grouped under the existing DSA backend section.
• python/sglang/srt/layers/attention/dsa/dsa_indexer.py:
  • Indexer.__init__ reads the two flags (ROCm-only gate) and logs once at layer 0 when active.
  • _get_topk_ragged gets a dedicated branch right after the OOM-chunk decision. It calls the new _get_topk_ragged_m_split, which returns None when there is nothing to split (attn-TP size 1, or fewer rows than ranks); in that case the existing replicated paths run untouched.
  • _get_topk_ragged_m_split implements the stripe loop. Each stripe reuses the exact per-row-range contract the existing OOM-chunk path already uses (ks[start:end], ke_offset=lengths, topk_indices_offset_override=global_offset[start:end] for the RAGGED top-k method; per-token cu_seqlens_q + token_to_batch_idx for the PAGED method), so top-k semantics are identical to the replicated path. When the OOM-chunk check fires, the stripe is additionally capped by the per-row logits budget (aiter's 2 GiB buffer_store limit on ROCm).
  • _m_split_all_reduce_max issues the MAX all-reduce. GroupCoordinator.all_reduce is SUM-only, so this calls torch.distributed.all_reduce(op=MAX) on the attn-TP device_group, using PyNCCL when a graph is being captured, mirroring _broadcast_indexer_topk_from_rank0_impl.

When it engages

All of the following must hold, otherwise the existing replicated indexer runs:

• SGLANG_DSA_INDEXER_M_SPLIT=1
• ROCm (is_hip()): the branch only wires the aiter Triton fp8_mqa_logits; CUDA (DeepGEMM) and XPU keep their paths.
• extend/prefill batches routed through _get_topk_ragged (decode, target-verify, and draft-extend stay on _get_topk_paged)
• attn-TP size > 1 and at least tp_size query rows in the batch

Test Plan

• Image: lmsysorg/sglang:v0.5.19-rocm720-mi35x (ROCm 7.2.0, torch 2.9.1+rocm7.2.0, triton 3.7.0, aiter c16d44b93), SGLang editable at this branch on top of 2f7393f0d.
• 8x MI355X (gfx950), TP8, GLM-5.2-MXFP4 (GlmMoeDsaForCausalLM, index_topk=2048, index_n_heads=32, index_head_dim=128, 78 layers).
• Serve configuration (the InferenceX recipes below launch the server themselves with exactly these arguments; for the accuracy runs the same command is used standalone, plus --reasoning-parser glm45 so the <think> block lands in reasoning_content and the graders see only the final answer in content):

SGLANG_DSA_INDEXER_M_SPLIT=1 sglang serve \   # =0 for the Before arm
  --trust-remote-code \
  --model-path /share-mv/amd/GLM-5.2-MXFP4 \
  --tp 8 \
  --kv-cache-dtype fp8_e4m3 \
  --dsa-prefill-backend tilelang \
  --dsa-decode-backend tilelang \
  --chunked-prefill-size 32768 \
  --mem-fraction-static 0.85 \
  --cuda-graph-max-bs-decode 256 \
  --max-running-requests 256 \
  --watchdog-timeout 1200 \
  --disable-radix-cache \
  --host 0.0.0.0 --port 8080

Activation check. Each TP rank must log this once at startup; its absence means the flag did not reach the schedulers or the platform gate failed:

DSA indexer M-split enabled: prefill indexer rows are striped across attn-TP ranks (stripe=512).

Accuracy gates follow the AMD "GLM-5.2 MXFP4 - vLLM status and validation" deck (GSM8K ~92%, GPQA-Diamond ~92%+, RULER/NIAH ~90%+), run against the served endpoint:

• Accuracy — GSM8K

lm_eval --model local-completions --tasks gsm8k \
  --model_args "model=<model>,base_url=http://127.0.0.1:$PORT/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False,timeout=600"

• Accuracy — GPQA-Diamond, 198 questions, zero-shot, against the running server:

.venv/bin/python -m sglang.test.run_eval --eval-name gpqa --model <model> \
  --host 127.0.0.1 --port $PORT --num-threads 32 --max-tokens 100000 \
  --temperature 1.0 --top-p 0.95 \
  --chat-template-kwargs '{"enable_thinking": true, "think_end_token": "</think>"}'

• Accuracy — RULER / NIAH, 500 samples per context length, offline engine:

lm_eval --model vllm \
  --model_args 'pretrained=<model>,tensor_parallel_size=8,max_model_len=262144,gpu_memory_utilization=0.85,kv_cache_dtype=fp8,trust_remote_code=True,enable_thinking=True,think_end_token=</think>,max_gen_toks=8192,block_size=64' \
  --tasks niah_single_2 --metadata '{"max_seq_lengths":[65536,131072]}' \
  --batch_size auto --seed 1234 --log_samples --output_path ./niah

Test Result

8x MI355X, TP8, GLM-5.2-MXFP4 throughout.

Accuracy summary (all three deck gates, Before vs After)

Gate	Before	After	Deck expectation
GSM8K strict-match (flexible-extract), 1319 q, 5-shot	97.12% ±0.46 (97.12%)	97.12% ±0.46 (97.50%)	~92%
GPQA Diamond, 198 q, 100k-token budget	90.40% ±2.09 (179/198)	91.92% ±1.93 (182/198)	~92%+
RULER niah_single_2 @ 64k, 500 samples	100%	100%	~90%+
RULER niah_single_2 @ 128k, 500 samples	100%	100%	~90%+

E2E serving throughput, ISL 8192 / OSL 1024

num_prompts = 10 x CONC.

Conc	Out tok/s Before	Out tok/s After	Mean TTFT ms Before → After	P99 TTFT ms Before → After	Mean TPOT ms Before → After
4	306.6	306.9 (+0.1%)	1043 → 1067 (+2.3%)	1517 → 1619	12.04 → 12.00
8	529.1	529.2 (+0.0%)	1812 → 1820 (+0.4%)	2960 → 2963	13.36 → 13.35
16	810.1	806.9 (-0.4%)	3267 → 3296 (+0.9%)	5839 → 5889	16.57 → 16.62
32	1182.6	1187.1 (+0.4%)	6172 → 6190 (+0.3%)	11599 → 11619	21.04 → 20.92
64	1462.4	1457.6 (-0.3%)	11889 → 11995 (+0.9%)	22945 → 23145	32.17 → 32.21
128	1810.7	1808.7 (-0.1%)	23427 → 23481 (+0.2%)	45701 → 45854	47.83 → 47.86
256	2048.3	2045.7 (-0.1%)	46313 → 46586 (+0.6%)	91372 → 91885	79.77 → 79.66

E2E serving throughput, ISL 50000 / OSL 1024

num_prompts = 10 x CONC.

Conc	Out tok/s Before	Out tok/s After	Mean TTFT ms Before → After	P99 TTFT ms Before → After	Mean TPOT ms Before → After
4	183.8	191.1 (+3.9%)	7063 → 6532 (-7.5%)	10219 → 9381	14.87 → 14.56 (-2.1%)
8	248.2	261.6 (+5.4%)	12156 → 11165 (-8.1%)	20089 → 18444	20.37 → 19.69 (-3.4%)
16	297.5	316.7 (+6.5%)	22064 → 20256 (-8.2%)	40013 → 36664	32.26 → 30.75 (-4.7%)
32	337.9	362.7 (+7.3%)	42044 → 38595 (-8.2%)	79816 → 73189	53.67 → 50.57 (-5.8%)
64	357.9	384.4 (+7.4%)	81852 → 75405 (-7.9%)	159320 → 146346	98.94 → 92.91 (-6.1%)

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR is described, including related work.
• Test commands are provided.
• Test and model evaluation results are provided with their tested base (activation check, GSM8K, GPQA Diamond, RULER/NIAH 64k+128k, and the two E2E sweeps over conc 4-256, Before vs After on the same branch).
• AI assistance is disclosed.

[dungthai414]

LGTM