[ChangLi811]

Needs sgl-project/sglang#44031 (V4.1's FP8 unified KV with the DSpark draft).

Motivation

The monokernels read only SGLang's paged KV, so SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton with SGLANG_DSV4_UNIFIED_KV_FP8=1 is refused, while the unified KV is the cache V4.1 serving moves to (with #44031, including the DSpark draft).

Modifications

• K1 writes the step's rows at SGLang's ring rows in the two-pool layout (a 512 B nope row of e4m3 with each E8M0 code twice, a 128 B bf16 rope row); for the DSpark draft's bf16 ring it writes the bf16 pair, rounded before RoPE as SGLang's draft writer does.
• K2 reads the window keys from SGLang's ragged SWA stream and the compressed rows past the layer's base as the paged record's words, so the attention core is the paged one; on the bf16 ring it quantizes each key's nope tiles as K1 quantizes a record.
• One field, kv in ("paged", "fp8", "bf16"), names the layout in the builds and the runner; the paged path is unchanged. The bf16 unified KV without the FP8 pools is still refused.
• Kernel checks: check_k1 --unified / --bf16-ring against SGLang's own writers of the two layouts; check_k2 --unified --bf16-ring on the unified pools holding the paged keys, bit for bit against the paged K2.
• The README lists [Feat] [AMD] DeepSeek-V4.1: FP8 unified KV for the ratio-1/2 layers, with DSpark, PD and the UMBP linker sglang#44031 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. FP8 unified KV + decoder replay; on and off on fresh servers.

config	monokernels	GSM8K	accept length
TP4	on	0.967	3.39
TP4	off	0.967	3.37
TP2	on	0.967	3.29
TP2	off	0.966	3.29
TP4 + prefill CP4	on	0.964	3.34
TP4 + prefill CP4	off	0.967	3.32

Needles, zero-row groups and c16 / c32 stress pass; KL within bounds. Kernel checks at TP2 and TP4 (widths 1–48): check_k1 / check_k2 paged, --unified, --bf16-ring all pass after rebasing onto main.

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, #78 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.)

PD

MoRI PD (enabled by #70) over two MI355X nodes: P = TP4 + prefill CP4 + DSpark without the monokernels, D = TP4 + DSpark; SGLANG_HIP_DSV41_INDEX_K_PAGE8=1 on both peers; P and D started afresh for every row; GSM8K through the PD router.

KV	D's monokernels	GSM8K, 64 threads	GSM8K, 8 threads	accept length, 8 threads
FP8 unified	on	0.964	0.966	3.54
FP8 unified	off	0.962	0.962	3.52

At 8 threads every verify step of the D runs on the monokernels when they are on. Greedy prompts (short, 15k–100k) and 50 long needles pass on every row; no GPU faults on either peer.

Speed Tests and Profiling

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

concurrency	off	on	gain
1	355 / 2.60	520 / 1.70	+47%
2	580 / 3.18	780 / 2.35	+35%
4	844 / 3.95	1133 / 3.08	+34%
8	1315 / 5.48	1426 / 5.03	+8%
16	1902 / 7.06	1917 / 7.00	+1%

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 unified-KV 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 pool's get_unified_compress_row_base comes with [Feat] [AMD] DeepSeek-V4.1: FP8 unified KV for the ratio-1/2 layers, with DSpark, PD and the UMBP linker sglang#44031, not in SGLang main yet. It is read by name (PENDING_POOL_READS), as minimax_m3 does for its pending switches: on SGLang main the FP8 unified KV is refused at startup with that reason, and the paged KV is unaffected; on SGLang with #44031 nothing changes. Once #44031 merges, a follow-up moves the name into SGLANG_READS and drops the check.
• 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.)