[XPU] Support decode context parallelism (DCP) on Intel XPU

[AnuSajikumar6264]

Motivation

Enable --dcp-size > 1 (decode context parallelism) on Intel XPU in the intel_xpu attention backend, as requested in review. DCP shards the KV cache across the ranks of a TP group by token position, and merges each rank's partial attention output by log-sum-exp. It grows the servable context when tp_size exceeds the model's KV heads, which TP would otherwise replicate.

Modifications

DCP lives in xpu_backend.py and uses only XPU-native kernels (sgl-kernel-xpu and torch ops). There is no Triton on the DCP path.

Step	XPU implementation
KV write	The position's owner (pos % dcp_size) stores it at loc // dcp_size. Rows other ranks own get index -1, which store_cache_xpu skips, so every layer uses the native store. The slots are built once per batch.
Decode	All-gather Q across the DCP group, then run the paged FMHA (flash_attn_with_kvcache(return_softmax_lse=True)) over this rank's shard. Merge with one all_to_all_single plus sgl-kernel-xpu merge_state_v2. This is the new XPU default, --dcp-comm-backend a2a; ag_rs (cp_lse_ag_out_rs_mha) stays selectable.
Extend	New tokens run causal (flash_attn_varlen_func). The sharded prefix runs over the gathered Q and is merged with cp_lse_ag_out_rs_mha. merge_state_v2 combines the two.
Page table	A strided view of req_to_token, one gather and one divide.

Sharding and merge follow the shared DCP path the Triton backend runs on CUDA/ROCm, and a2a stays decode-only, as upstream's is.

Two XPU-specific details:

• sgl-kernel-xpu's varlen prefill hangs at seqlen_q == 1, so a 1-token chunk uses the closed form (output v, LSE the single score).
• The FMHA returns NaN for a row whose rank holds no keys for it, so those rows are zeroed before the merge.

Shared plumbing, all XPU-gated:

• parallel_state admits XPU.
• _alloc_page_size uses the physical allocator page.
• _dcp_comm_backend_default resolves a2a on XPU; the help text and DCP docs are updated to match.

Startup checks (XPU only)

attention_hook._check_xpu_dcp refuses each of these at startup, naming the flag to change:

• an attention backend other than intel_xpu;
• MLA models;
• unequal K and V head dims (those fall back to an indexed write that does not skip -1);
• --dcp-size not dividing --tp-size;
• models whose ranks would not share their DCP group's KV heads (everything but Qwen3.5 when total_num_kv_heads > tp / dcp);
• a decode GQA group above 16;
• speculative decoding;
• chunked local attention;
• encoder-decoder models.

xpu_backend refuses what only the loaded model settles: sliding-window layers, a quantized KV cache, and captured graphs.

No code path off XPU changes.

Changes since the previous revision

These address the review by @mingfeima:

• Non-owned KV rows: they now get index -1, which store_cache_xpu skips, instead of slot 0, which a live token can map to. The parity test now places live tokens at local slot 0: the old code fails 10 of its subtests, the fix passes.
• KV cache dtype: the DCP check reads the resolved dtype, so --kv-cache-dtype bf16 is accepted, and an fp8 cache (explicit, or an fp8 checkpoint under auto) is still refused.
• a2a buffers: send and recv are cached on the backend and reused while the shape holds, so the decode layers after the first skip the allocation.
• Kernel version check: removed, since main pins sglang-kernel-xpu 0.3.0.
• Naming: replicated is now kv_head_layout_supported.

Earlier, from the review by @rahulvijayaraghavan:

• Decode merge: XPU-native a2a, now the XPU default. Per layer on 2 GPUs (16 heads, batch 1-128) it takes 178-276 us, vs 411-462 us for ag_rs, at the same max error.
• KV write: the native store, with slots built once per batch from the padded positions. This replaces the per-layer mask and the Triton masked write: ~12 vs 44-47 us per layer at decode sizes.
• Extend merge: merge_state_v2, 2.3-5.4x faster than the torch merge. The fp32 casts now happen only where the merge needs them.
• Page table: rebuilt with torch, 23 us vs 69-76 us.
• Empty-row mask: built on the device, with no blocking copy. The host stall drops from ~16 ms to 0.2 ms.
• forward_batch_info.py: no longer changed.
• Tests: the parity test now writes the new tokens through a real KV pool, runs decode under both merges, and covers ranks without keys.

Rebased onto main 8ad24fb560.

Accuracy Tests

Setup: Qwen3.5-4B, 2x Arc Pro B60, TP=2, intel_xpu, sglang-kernel-xpu 0.3.0, --disable-radix-cache --chunked-prefill-size 256.

GSM8K, all 1319 questions, 5-shot, same workload:

	dcp=1	dcp=2
This revision	0.8828	0.8858
Previous revision		0.8782

Greedy parity vs dcp=1 (12 requests, 64 new tokens each). The mean logprob gap over the agreeing tokens is 0.0028 with a2a and 0.0025 with ag_rs, vs 0.0027 before. Every divergence is a near-tie between two tokens, e.g. 0.29 vs 0.30.

Speed Tests and Profiling

Decode tok/s for Qwen3.5-4B, TP=2, 4096-token prompts, 64 new tokens; median of 5 interleaved runs:

batch	dcp=1	dcp=2, previous revision	dcp=2, this revision
1	15.16	13.03	13.49 (+3.5%)
4	65.42	52.50	54.82 (+4.4%)
16	253.16	209.41	213.88 (+2.1%)

Prefill is 0.2-0.8% faster. DCP trades decode speed for context. Qwen3.5-4B has 4 KV heads, so TP=2 does not replicate them: at dcp=2 each rank holds all 4 heads for half the positions, and the per-rank KV pool goes from 340,864 to 170,112 tokens. The total context is unchanged, so this model checks correctness, not the context gain. The gain shows when heads are replicated: GLM-4.1V-9B (2 KV heads) at TP=4 kept the same per-rank capacity at dcp=2, which doubles the servable context (measured on an earlier revision).

Tests

• test/registered/xpu/test_xpu_dcp.py (1 XPU, stage-b-test-1-gpu-xpu). It runs decode (under ag_rs and a2a) and extend on threaded ranks against full attention at dcp 2 and 4, including ranks that hold no keys. The new tokens are written through the backend into a real MHATokenToKVPool, so a wrong KV slot shows up as wrong attention output. The file also covers the startup checks.
• test/registered/xpu/test_xpu_dcp_gsm8k.py (2 XPUs, nightly-xpu-2-gpu): a GSM8K gate at 0.80 on Qwen3.5-4B, TP=2/DCP=2. Local run: 0.885, passed.
• test_server_args gains an XPU comm-backend default case, and its platform-faked cases now pin is_xpu.
• test_encoder_decoder_varlen_gather sets dcp_size on its stubbed backend.
• test_xpu_int4_dense sets tp_group on its layer. Since [Refactor] Retain process groups in parallel linear layers #42588 the GPTQ XPU kernel reads tp_size from it, so the split-group test was failing XPU CI on main.

How to test

# one XPU
python -m pytest -q test/registered/xpu/test_xpu_dcp.py

# two XPUs
python -m pytest -q test/registered/xpu/test_xpu_dcp_gsm8k.py

Or by hand, launching once with --dcp-size 1 and once with --dcp-size 2:

python -m sglang.launch_server --model-path Qwen/Qwen3.5-4B --device xpu \
  --tp-size 2 --dcp-size 2 --attention-backend intel_xpu \
  --disable-radix-cache --chunked-prefill-size 256 --mem-fraction-static 0.60 --port 30000

python -m sglang.test.run_eval --base-url http://127.0.0.1:30000 --eval-name gsm8k \
  --api completion --num-examples 1319 --num-shots 5 --max-tokens 512 --num-threads 32

Not covered here

• Graph mode with DCP: refused for now. The capture path doesn't build the DCP metadata, and all_to_all_single can't be captured on XPU (all_gather and all_reduce can). Planned as a follow-up.
• Radix cache: Qwen3.5 on XPU needs --disable-radix-cache, independently of DCP.
• dcp=4: covered by the unit tests and a 4-rank oneCCL check, but not on 4 separate GPUs.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): 🚫 Run #38197433978
Latest PR Test (Extra): ❌ Run #38197433448
Latest PR Test (AMD ROCm 10): ❌ Run #38197433838

[siju-samuel]

/tag-run-ci-label

[mingfeima]

also please rebase to master.

[rahulvijayaraghavan]

@AnuSajikumar6264
Nice work! Thanks for the contribution, addressing comments!