[billishyahao]

This PR co-work with mono patch: sgl-project/sgl-amd-mono#13

Motivation

The AMD monokernel backend runs a Kimi-K3 decode or DSpark verify step as two persistent launches a layer. They run in this order.

• K1, the KDA front launch, covers the attention side of a KDA layer: the attention-side residual seam, input_layernorm, the in-projection, f_b, the conv, the recurrence and the gated output norm. It covers the 69 KDA layers. The 24 full-attention layers keep stock SGLang's attention instead.
• K2, the layer launch, takes the TP-local attention output from there and runs o_proj, the attention-TP all-reduce, the MLP attention-residual seam and the MoE. It covers all 92 MoE layers, whichever attention produced its input.

That backend is sgl-project/sgl-amd-mono. It loads through the sglang.srt.plugins entry point and is selected with --monokernel-backend amd. SGLang needs nothing else for it: install_monokernel_forward swaps the whole model forward (srt/monokernels/runtime.py), so the kernels, the spec and its registration all live in the plugin. This method is the one thing the plugin cannot reach from outside, which is why it is the only change here. In short, this change follows the sgl mono kernel RFC #43949

Modifications

1. The KDA in-projection gate, which K1 reads. Take the merge when do_fuse_qkvbfg allows it, as before, or when a monokernel backend is selected:

if not self.use_full_rank_gate:
    return False
if not (self.do_fuse_qkvbfg or get_exec().kernel.monokernel_backend):
    return False

With no backend selected this reduces to use_full_rank_gate and do_fuse_qkvbfg, which is what the predicate tested before. That identity is deliberate. The fused GEMM is a different kernel from the split pair. A wider gate on its own would move decode output by one bf16 rounding step for every ROCm Kimi-K3 user. The file's own docstring records that the two paths agree to bf16 rounding and not bitwise.

The weight checks below the gate do not change, and they still decide. The backend changes the policy. It does not remove the safety checks. _is_unquantized_mergeable requires the three weights to share one dtype, and that dtype to be bf16 or fp16. A per-channel FP8 weight or a packed MXFP4 weight still refuses, because _merge_weights_as_views concatenates .weight alone.

Nothing else changes. do_fuse_qkvbfg also gates a weight-loader branch and the low-rank gate's projection layout. This commit touches neither.

2. forward_pre_o_proj() on KimiK3DeltaAttention and KimiK3MLAAttention, where K2 starts. It returns the TP-local attention output after the output gate and before o_proj. The caller owns o_proj and the attention-TP reduction that o_proj carries, so o_proj(forward_pre_o_proj(x)) == forward(x).

• KimiK3DeltaAttention: a plain extract. forward() now calls the new method and keeps the o_proj tail. The moved body does not change.
• KimiK3MLAAttention: the MLA cores reach o_proj from several places. The split therefore uses the instance-level wrapper that already applies the output gate. forward_pre_o_proj() sets _return_pre_o_proj, runs forward(), and takes the tensor that the wrapper puts in _pre_o_proj_out instead of running o_proj.

Four guards hold the MLA path together. It has no in-tree caller to exercise it, and a fault on it gives wrong numbers rather than an error.

• use_output_gate must be on. Without it the wrapper does not exist.
• The method is not re-entrant. A second entry while the flag is set raises.
• _pre_o_proj_out holds one tensor. If o_proj is reached twice in one attention forward, the second call raises instead of overwriting the first. The message names the layer.
• A quantized o_proj input raises. On ROCm the MLA core quantizes what it hands to o_proj when o_proj's weights are MXFP4 or block-scale fp8 (forward_mla_rocm.py:325,351). The output gate is skipped with it, so there is no gated bf16 tensor to hand back. Every rank quantizes alike, so this raises on all of them. It does not split the TP group.

The standard path remains no change. forward() produces what it produced before if no mono kernel server args are specified.

Accuracy Tests

8 x MI355X, TP8 / DCP1, FP8 KV cache, DSpark block size 7. /data-models/Kimi-K3, GSM8K 200 questions, --parallel 2:

	Accuracy
This branch, no monokernel backend	0.990
This branch, --monokernel-backend amd	0.985

Speed Tests and Profiling

--monokernel-backend amd against the standard path, both launches on. Two
server instances an arm, four probe runs each, the first run of each server
dropped. Decode steps/s rather than tok/s, because the speculative accept
length swamps tok/s on this recipe.

	decode steps/s	sd	range
Monokernels off	41.63	0.62	40.28 - 42.10
Monokernels on	55.26 (+32.8%)	0.38	54.72 - 55.94

The ranges do not overlap. The win is on 8-row verify steps; wider steps and
all prefill take the standard path.

Fixed-length decode latency at long context

The table above is short prompts at the draft's real acceptance. This one is the interactive shape the backend is for: a cached long prefix, a small new turn, and a short answer. It also reports the inter-token latency the backend moves, rather than a step rate that a reader has to convert.

ISL	ITL P50 off	ITL P50 on	Δ	ITL P90 off	ITL P90 on	Δ	steps/s off	steps/s on	Δ
64k	6.17 ms	4.62 ms	−25.2%	8.11 ms	6.17 ms	−24.0%	40.4	52.7	+30.5%
128k	6.67 ms	5.29 ms	−20.6%	8.56 ms	6.88 ms	−19.7%	37.3	47.1	+26.3%
256k	7.70 ms	6.19 ms	−19.6%	9.86 ms	8.10 ms	−17.9%	32.2	40.3	+24.9%

off / on is --monokernel-backend amd. Lower ITL is better, so Δ is the reduction; as a speedup the same numbers are +33.7% / +25.9% / +24.3% at P50.

P90 moves with P50 (−24.0% against −25.2% at 64k) and the P90/P50 ratio is unchanged (1.31 off, 1.33 on). The launch shifts the whole distribution rather than clipping a tail. The three columns agree because acceptance is pinned: ITL ~= 1000 / (steps/s * accept).

The gain falls as ISL grows (−25.2% to −19.6%). Attention over a longer KV takes a larger share of the step, and that part is outside the launch.

TTFT is unchanged -- 0.24/0.25 s, 0.30/0.29 s and 0.41/0.41 s at 64k/128k/256k. The launch covers decode only.

End-to-end agentic replay

The two tables above are synthetic prompts at a pinned acceptance. This one is a full AgentX agentic-coding replay, one hour per arm. It therefore carries real session content, real prefix-cache behaviour and the batching of the scheduler itself.

conc 1, 8 rows a step	off	--monokernel-backend amd	Δ
TPOT mean	6.28 ms	4.94 ms	−21.3%
TPOT P50	6.23 ms	4.91 ms	−21.2%
TPOT P75	6.60 ms	5.26 ms	−20.3%
TPOT P90	6.74 ms	5.37 ms	−20.3%
TPOT P95	6.77 ms	5.43 ms	−19.8%
interactivity P50	160.4 tok/s	203.7 tok/s	+27.0%
interactivity P90	148.5 tok/s	186.1 tok/s	+25.3%
TTFT mean	2.70 s	2.10 s	−22.4%
e2e P90	37.3 s	30.0 s	−19.5%
output	97.5 tok/s	112.0 tok/s	+14.9%
requests completed	205	222	+8.3%

Every percentile moves by the same amount, −19.8% to −21.3%, as in the fixed-length table. Throughput rises less than latency falls, +14.9% against −21.3%. AgentX replays a closed loop, so a faster answer shortens only the part of each session's cycle that the engine owns.

The number of steps the backend can take bounds the gain. An earlier pair of the same runs admitted two requests instead of one. That pair held the SSM state in bf16, so read its percentages against its own baseline and not against the table above. There 6 to 7% of decode windows held two running requests. That is 2 x 8 = 16 rows, which is outside the spec, and the result split: P50 −16.5%, but P90 −2.6% and P95 −2.1%, with the mean diluted to −9.4%. The server log shows it directly when read by #running-req: the one-request steps ran 9.8% faster and the two-request steps 2.2%. Admitting one request removes the steps that are not covered, and the whole distribution then moves together. A caller that wants the full gain must keep every verify step inside the supported width.

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.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ⏳ Run #41977697297
Latest PR Test (Extra): ❌ Run #41977697285
Latest PR Test (AMD ROCm 10): ⏳ Run #41977697327

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

[billishyahao]

base test all passed. https://github.com/sgl-project/sglang/actions/runs/39700405113 cc @HaiShaw

[billishyahao]

@HaiShaw @kkHuang-amd new round of basic test is passed. It shows no side effect on common path: https://github.com/sgl-project/sglang/actions/runs/41613665770

[HaiShaw]

@billishyahao can you bump the sgl-amd-mono commit# reference.
cc @chien-an-chen @kkHuang-amd @zijiecode

[billishyahao]

@billishyahao can you bump the sgl-amd-mono commit# reference. cc @chien-an-chen @kkHuang-amd @zijiecode

Thanks @HaiShaw for the reminding! I bump the sgl-amd-mono commit through 254ce94