[ChangLi811]

Motivation

With prefill context parallelism (--enable-prefill-cp, sgl-project/sglang#43465 on HIP), --enable-amd-monokernel refuses to start: the spec does not accept CONTEXT_PARALLEL, and the attention layers hold all 64 heads instead of a TP share. Long-context serving then has to choose between prefill CP and the monokernels at low concurrency.

Modifications

• Under prefill CP SGLang keeps every attention head on every rank (attention TP 1) and the whole KV and index caches on every rank, but decode and verify steps never run CP. K1 and K2 therefore run this rank's TP share as without CP: wq_b's and wo_a's rows are views of the replicated weights, wo_b's columns (in aiter's (16, 16) preshuffle too) and its block codes a copy. The kernels are unchanged.
• _config_disable_reasons accepts CP only as prefill CP over the whole TP group; decode CP or a narrower CP group is refused.
• draft.py accepts the DSpark draft's replicated heads the same way.
• Unit tests: the CP gate, and the head shards of row-major and preshuffled weights.
• The README lists [Feat] [AMD] DeepSeek-v4.1: Support prefill context parallel sglang#43465 among the requirements.

Accuracy Tests

DeepSeek-V4.1-Flash + DSpark (block size 5), one MI355X node, rocm/sgl-dev:v0.5.22-rocm10-mi35x-20261025, SGLang: a local merge of #43465 (15e93db9e), #44730, #44031 and #46036 (aa6e51982). GSM8K: sgl-eval, 1319 questions, max tokens 4096. Monokernels on and off on fresh servers of the same configuration (prefill CP + --enable-cp-decode-attn-tp + decoder SWA bounded replay).

TP / CP	monokernels	GSM8K	DSpark accept length
TP4 / CP4 (4 runs)	on	0.960–0.967	3.51
TP4 / CP4 (4 runs)	off	0.959–0.963	3.48–3.56
TP2 / CP2	on	0.966	3.55
TP2 / CP2	off	0.959	3.52

Decode-vs-prefill KL stays within the spec's bounds; long-context needles (to 100k), CP's zero-row request groups, c16 / c32 stress and image prompts (MMMU-Pro) match the monokernels-off server. These were taken on an earlier base; after rebasing onto main the kernel checks pass at TP2 and TP4 and a TP4 / CP4 server (FP8 unified KV and the page8 index pool on, #79 and #80) scores 0.964.

GSM8K above runs 64 threads, so most verify steps are wider than 8 requests and take SGLang's path. With AgentX's c1 recipe (max-running-requests 2, its 19 environment variables, SGLANG_HIP_DSV41_INDEX_K_PAGE8=1), every verify step runs on the monokernels; on #70 (now in main) plus this PR, #79 and #80 (DeepSeek-V4.1-Flash-1017 weights, fresh servers):

config	on	off
TP4, paged	0.967 / 3.41	0.967 / —
TP4 + prefill CP4, FP8 unified	0.967 / 3.41	0.963 / —
TP2, paged	0.967 / 3.39	0.968 / 3.40
TP2 + prefill CP2, FP8 unified	0.968 / 3.41	0.967 / 3.39

(GSM8K / DSpark accept length. TP4 "off" accept lengths are left out: with the recipe's --disable-shared-experts-fusion, SGLang's own path accepts less at batch 1–3, which --enforce-shared-experts-fusion, as the README uses, restores.)

Speed Tests and Profiling

TP4 / CP4, ShareGPT through the chat API, 500 output tokens: output throughput (tok/s) / mean TPOT (ms).

concurrency	off	on	gain
1	336 / 2.73	535 / 1.62	+59%
2	524 / 3.55	794 / 2.27	+52%
4	913 / 3.87	1132 / 3.02	+24%
8	1347 / 5.43	1522 / 4.84	+13%
16	1674 / 7.10	1774 / 6.70	+6%

Checklist

• pre-commit run --all-files passes (ruff, isort, codespell).
• Tests under test/ cover the change.
  • CPU, against SGLang aa6e51982: pytest test/unit passes 264 (262 on main), including the new CP tests and the deepseek_v41 contract, which also passes against SGLang main (c6c753395). The 7 failures and 3 collection errors (GLM-5.3, MiniMax and DeepSeek-V4-Pro tests) are the same on main and unrelated to this PR.
• The SGLang commit or PR this was tested against is named above.
• Accuracy and speed results are attached when kernel or forward code changes.
• A new model adds one line to .github/CODEOWNERS that names its owner. (No new model.)