[omirosh]

Part 3 of 5 of the split of #16. This PR is stacked on #88. Until #83 and #88 merge, the diff also shows their three commits. Review only the last commit, 3a8fada.

Motivation

#83 runs the tail of each decoder layer as one K2 launch. In front of K2, each of the 69 Gated DeltaNet (GDN) layers still runs SGLang's input norm and GDN as separate launches: the residual add and the input norm, in_proj_qkvz and in_proj_ba, the conv update, the gated delta rule and the gated norm.

This PR runs that front as one K1 launch on each rank. K1 is the gdn_pre kernel of sgl-project/sglang#43516 (at 51b0c592b2), which adapts it from vllm-project/vllm#60833. As in K2, each wait in K1 has a deadline.

A decode step of 1 to 8 rows now runs each GDN layer as K1, then K2. The 23 full-attention layers run SGLang's attention, then K2, as in #83.

The verify steps of #88 do not use K1. K1 updates the SSM state in place and keeps no state for each draft token, so a rejected draft token cannot roll back. With NEXTN, the server builds no K1 kernel. The verify steps run SGLang's GDN, then K2.

Modifications

Area	Change
qwen38_a95b_mxfp4/gfx950/gdn.py	Adds K1. One launch for each GDN layer and rank runs the residual add and the input norm, in_proj_qkvz and in_proj_ba, the conv update, the gated delta rule and the gated norm. It writes the core output (out_proj's input), the residual, and the conv and SSM state rows in place. K1 computes the recurrence in f32 and stores the SSM state in f32 or bf16. A pad row (state index ≤ 0) writes no state and a zero core row.
qwen38_a95b_mxfp4/spec.py	Runs K1 in front of K2 on each GDN layer in decode steps. Builds two K1 keys for each width. One is for layer 0, which has no residual input. The other is for the other GDN layers. Reads SGLang's state indices and each layer's conv and SSM state. Adds the refusals below. With speculative decoding, it builds no K1 key.
qwen38_a95b_mxfp4/check.py	Check mode also runs SGLang's input norm and GDN next to each K1 launch, on the state rows as they were, and then restores the rows. It compares the core, the residual, and the conv and SSM state rows. NUMERICS adds k1.core.cosine, k1.conv.cosine and k1.state.cosine ≥ 0.9999.
test/kernels/qwen38_a95b_mxfp4/check_k1.py	Adds the K1 checks on one GPU, on random weights in SGLang's layouts (results below).
test/kernels/qwen38_a95b_mxfp4/check_port.py, inputs.py	The port test also compares K1 with the gdn_pre kernel of sgl-project/sglang#43516, bit for bit.
test/kernels/test_qwen38_a95b_mxfp4_kernels.py	Adds TestQwen38MoeK1.
test/e2e/test_qwen38_a95b_mxfp4_server.py	The check-mode test also needs K1 checks at 1, 3 and 8 rows. The count of rows routed elsewhere uses only the K2 checks.
test/unit/test_qwen38_a95b_mxfp4_spec.py	Checks the ReplaySSM refusal and the GDN weight refusals.
README.md	The model row and the design text name K1.

Refusals. This PR adds these startup refusals:

• --enable-linear-replayssm without speculative decoding. K1 updates the SSM state in place, and ReplaySSM defers the update to a ring. [Feat] [MI355X] qwen38_a95b_mxfp4: Run each decoder layer's tail as one K2 launch #83 accepted this flag.
• An SSM state that is not f32 or bf16, or a conv state that is not bf16.
• A GDN layer whose in_proj_qkvz, in_proj_ba, conv, dt_bias or norm is not bf16, or whose A_log is not f32.
• A GDN conv with a bias or without SiLU, a gated norm that is not over the head with SiLU, or a norm after the gate.

The spec still refuses the MXFP4 GDN projections. That check now covers in_proj_qkvz too.

With the test CUDA graphs, each rank builds 8 K1 and 4 K2 kernels and logs "Qwen3.8 monokernel: 8 K1 and 4 K2 launches (widths [1, 2, 4, 8], decode)". With NEXTN 3/1/4, it logs "0 K1 and 2 K2 launches (widths [4, 8], verify)".

This PR does not change the K2 kernels, common/ or the other models.

Accuracy Tests

All the results below use this setup:

Item	Value
GPUs	8× MI355X (gfx950), one node
TP	8
Image	rocm/sgl-dev:v0.5.22-rocm10-mi35x-20261028@sha256:b51edf074b567263afb19be9108756645ee00d0604374b6e4c5006caccaa17fe
SGLang	5bb7ecc2c247db05e3de41d21f36aed7aa1a2dfb, which has sgl-project/sglang#43839, with no unmerged SGLang PR
AITER, FlyDSL	285579b6211b0da30f35b91a890c29c823597cd1, 0.4.0
Model revision	13dc9676cdd9e2374c3efb50389d3392b975acae
Server flags	ARGS of test/e2e/qwen38_a95b_mxfp4_server.py (--mamba-ssm-dtype bfloat16, so K1 stores a bf16 SSM state). The monokernel-on servers add --monokernel-backend amd.
Branch commit	5bba9fc (tag q38v3-tested/3-k1). The head 3a8fada is rebased onto main at b10f3b7. Its model code and tests differ from that commit only in the text of a docstring and a skip message. The test runs used the old base 0aa9749, before #76 changed PeerBuffer.close().

The port test (check_port.py) compares K1 and K2 with the original kernels of sgl-project/sglang#43516 at 51b0c592b2. That PR is not merged, so the test run mounts its two kernel folders (k3_mono_flydsl and qwen3_5_mono_flydsl) read-only into SGLang's tree in the container. No SGLang code imports them. Only the port test does.

GSM8K, monokernels on and off

The test uses lm_eval 0.4.13 with local-chat-completions and the chat template. It runs all 1,319 questions, 5-shot, at temperature 0, with max_gen_toks 32768. It sends 8 concurrent requests, so the decode steps take the K1 and K2 path. Each arm starts a fresh server. Each question has the same prompt_hash in both arms.

Filter	Monokernels off	Monokernels on	Right only off	Right only on	McNemar's p (exact)
flexible-extract	1290 (0.9780)	1295 (0.9818)	6	11	0.332
strict-match	1150 (0.8719)	1140 (0.8643)	104	94	0.523

The arms differ by 5 questions on flexible-extract and by 10 on strict-match. McNemar's test finds no accuracy difference on either filter.

Strict-match needs a #### N line. Of the answers that flexible-extract finds right, 140 have no #### N line in the off arm, and 155 in the on arm. Strict-match fails exactly these answers. Most of them end in **Answer: N** instead (79 and 90), for example **Answer: 25**. So the strict-match gap is in the answer format, not in the answers.

Kernel checks (test/kernels)

Check	Result
Compile every K1 build key (widths 1 to 8, first layer or not, f32 or bf16 state)	All 32 keys compile
Pad rows at 1, 4, 7 and 8 rows. A pad row writes a zero core row and no state, and no other slot changes.	4 cases pass
bf16 state. K1 runs the f32 recurrence on the widened bf16 state.	The core and the conv state are bit for bit equal to the f32 kernel's. The new state is the f32 result, rounded to bf16.
Batch invariance. A probe row's outputs do not change with the other rows or the width.	57 cases pass
Graph replay of widths 8, 1, 3, 8, 2, 5 in one graph on one state	200 replays, 0 mismatching checks against eager mode
Fault: K1 stalls at its first wait	K1 sets the deadline error word of that wait, and the launch returns in 1.05 s
Port test: K1's core, residual and state slots against the sgl-project/sglang#43516 gdn_pre kernel	18 cases (widths 1 to 8 with f32 and bf16 state, and the first layer at 1 and 8 rows) are bit for bit equal

Each run also repeats the K2 checks of #83, with the same results. The kernel tests do not compare K1 with SGLang's ops. Check mode in the e2e test does that, with the real weights.

Three runs of all tests

Run	unit	kernels	e2e	GSM8K decode (200)	GSM8K NEXTN (200)	Highest mean KL (2, 3, 7, 8 rows)	Rows that check mode skipped	Lowest k1.core.cosine	Lowest k2.out.cosine
1	12 passed	6 passed	5 passed	0.975	0.985	0.0178	0 of 3956	0.999990	0.998749
2	12 passed	6 passed	5 passed	0.975	0.975	0.0222	0 of 3956	0.999990	0.999756
3	12 passed	6 passed	5 passed	0.980	0.985	0.0086	0 of 3956	0.999973	0.999360

Each run has 1104 K1 checks and 1472 K2 checks. In all three runs, k1.state.cosine is at least 0.999989 and k1.conv.cosine rounds to 1.000000. K1's residual is bit for bit equal to SGLang's (k1.res_out.max_abs is 0), and k2.res_out.cosine rounds to 1.000000.

The e2e tests are the five tests of #88:

• GSM8K decode: 200 questions at 12 concurrent requests, without NEXTN. The batch grows past 8 rows as it fills, and drops below 8 rows as it drains.
• GSM8K NEXTN (TestMonokernelVerify): 200 questions at 4 concurrent requests, with NEXTN 3/1/4. The verify steps run K2 only.
• Prefill match: at 2, 3, 7 and 8 rows, the KL divergence of the decode logprobs against prefill is below 0.06.
• Check mode: each K1 and K2 layer output is within NUMERICS, at 1, 3 and 8 rows. The test fails if check mode skips more than 1% of the K2 rows.
• Fault injection: after 100 steps, rank 1 stalls at its first wait. At the deadline, every rank stops, and the server returns no tokens from that step.

The tests do not cover TP4 or MI350X, because the spec serves only TP8 on MI355X.

Speed Tests and Profiling

The test uses sglang.bench_serve with random prompts of 8192 input and 1024 output tokens (range ratio 0.8). Each arm starts a fresh server, in the order off, on, on, off. The servers use the flags and environment of test/e2e/qwen38_a95b_mxfp4_server.py, plus --tool-call-parser qwen3_coder and --watchdog-timeout 12000.

Concurrency	Off median TPOT (ms), rounds 1 and 2	On median TPOT (ms), rounds 1 and 2	Per-user decode tok/s, off and on	Speedup	#83 on (ms)	Speedup over #83	Spread between rounds
1	11.01, 11.03	6.99, 7.00	90.7, 143.0	1.58×	7.46	1.07×	0.2%
2	11.59, 11.58	8.28, 8.28	86.3, 120.8	1.40×	8.69	1.05×	0.1%
4	13.77, 13.76	10.17, 10.17	72.6, 98.3	1.35×	10.41	1.02×	0.1%
8	17.33, 17.40	14.53, 14.54	57.6, 68.8	1.20×	14.90	1.03×	0.4%
16 (past the largest width)	23.31, 23.34	23.31, 23.34	42.9, 42.9	1.00×	23.37	1.00×	0.1%

• The speedup uses the mean of the two rounds.
• The "[Feat] [MI355X] qwen38_a95b_mxfp4: Run each decoder layer's tail as one K2 launch #83 on" column is the mean of the two monokernel-on rounds of [Feat] [MI355X] qwen38_a95b_mxfp4: Run each decoder layer's tail as one K2 launch #83, on the same node and image, on 2026-10-29. The off arms of the two runs agree within 0.3%.
• The spread is the largest difference between the two rounds of one arm.
• At concurrency 1 to 8, K1 makes the step 2% to 7% faster than [Feat] [MI355X] qwen38_a95b_mxfp4: Run each decoder layer's tail as one K2 launch #83's K2 alone. Each gain is larger than the spread.
• At concurrency 16, most steps have more than 8 rows and run SGLang's path, so the speed does not change.
• No concurrency is slower with the monokernels on.

Registers. The table comes from a compile-only build of the 32 K1 build keys at c3b16fb (the head before the rebase, with the same kernel code) for gfx950. Each key has a stop=True build (the spec) and a stop=False build (the tests). The numbers come from the AMDGPU metadata in FlyDSL's ISA dump. Where the builds of one width differ (first layer, bf16 state, stop), a cell gives the range.

Width	VGPR	VGPR spill	SGPR spill	LDS (bytes)	Scratch (bytes)
1	208	0	0	24656	0
2	208	0	0	41056	0
3	206 to 214	0	0	57456	0
4	208	0	0	73856	0
5	209 to 210	0	0 to 4	90256	0
6	208	0	0	106656	0
7	205 to 212	0	7	123056	0
8	207 to 214	0	4	139456	0

• No K1 build spills a VGPR or uses scratch. The SGPR spills go to VGPR lanes.
• No build uses AGPRs. The builds have 70 to 106 SGPRs.
• LDS grows with the width, to 139,456 of gfx950's 163,840 bytes at 8 rows.
• The K2 builds give the same numbers as in [Feat] [MI355X] qwen38_a95b_mxfp4: Run each decoder layer's tail as one K2 launch #83.

Kernels before and after. The table compares one decode step at batch 1 in two TP-0 profiles with the monokernels on: #83 at 3daea98 (K2 only) and this PR at 3a8fada (K1 and K2). Both servers ran on one node with 8× MI355X. The prompt has 8192 tokens, and the output has 100 tokens. The counts come from the step[DECODE bs=1] slices on the GPU stream. All 89 decode steps after the first 10 have the same counts. Each time is the mean of these 89 steps.

Part of the GDN layer	Before (#83): SGLang kernel	Calls per step	Time per call (µs)	After
The residual add and the input norm	_gemma_fused_add_rmsnorm_kernel; _gemma_rmsnorm_kernel on layer 0	68 + 1	4.5	One K1 launch for each GDN layer: qwen38_mono_k1_s1_f0_b1 (68 calls per step, 27.8 µs each) and qwen38_mono_k1_s1_f1_b1 (layer 0, 27.1 µs)
in_proj_qkvz and in_proj_ba	Cijk_Alik_Bljk_…_MT64x16x256_…	69	18.3	In K1
The split of the projection output	fused_qkvzba_split_reshape_cat_contiguous_kernel	69	4.5	In K1
The conv update	_causal_conv1d_update_kernel	69	4.5	In K1
The gated delta rule	fused_recurrent_gated_delta_rule_packed_decode_kernel	69	4.5	In K1
The gated norm	_layer_norm_fwd_1pass_kernel	69	4.2	In K1
out_proj to the MoE all-reduce	qwen38_mono_k2_s1_a0	92	39.8	The same K2 launch, 40.1 µs

• The 6 kernels of a GDN front take 40.5 µs together. K1 takes 27.8 µs.
• The step falls from 750 to 405 kernels, and from 8.11 to 7.25 ms of kernel time.
• In the K1 names, f1 is the key of layer 0, which has no residual input. b1 is the key of a bf16 SSM state. s1 is 1 row.
• The 23 attention layers do not change. They run SGLang's input norm and attention, then K2.
• Each screenshot selects one step[DECODE bs=1] slice and shows only its kernels (filter Category = kernel).

Before (#83), one decode step:

After (this PR), one decode step:

Checklist

• pre-commit run --all-files passes.
• Tests under test/ cover the change.
• 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. (N/A: main already has **/*qwen38_a95b_mxfp4* @omirosh @yichiche @chuyeh-amd.)

Review and Merge Process

1. The lint and Check PR title checks must pass.
2. Ask a code owner from .github/CODEOWNERS to review.
3. Merge with "Squash and merge": the PR title becomes the commit title on main.

[omirosh]

@yichiche Here are the test results. All three runs of each suite pass at 5bba9fc (tag q38v3-tested/3-k1). The head c3b16fb changes only the text of a docstring and a skip message.

The runs use the setup of the PR body: one node with 8× MI355X, TP8, the image rocm/sgl-dev:v0.5.22-rocm10-mi35x-20261028 and SGLang 5bb7ecc2c2. Each suite runs in that image, from /sgl-workspace/sgl-amd-mono:

python3 -m pytest -v -s -p no:cacheprovider test/unit/test_contract.py test/unit/test_qwen38_a95b_mxfp4_spec.py -k 'qwen38 or common'
python3 -m pytest -v -s -p no:cacheprovider test/kernels/test_qwen38_a95b_mxfp4_kernels.py
python3 -m pytest -v -s -p no:cacheprovider test/e2e/test_qwen38_a95b_mxfp4_server.py

The kernel run mounts the two kernel folders of sgl-project/sglang#43516 read-only for the port test. The e2e servers set SGLANG_ROCM_GDN_ONLINE_MXFP4=0 in test/e2e/qwen38_a95b_mxfp4_server.py.

Suite	Run 1	Run 2	Run 3
unit	12 passed in 18 s	12 passed in 19 s	12 passed in 18 s
kernels	6 passed in 1 min 53 s	6 passed in 1 min 52 s	6 passed in 1 min 53 s
e2e	5 passed in 17 min 28 s	5 passed in 18 min 9 s	5 passed in 18 min 47 s

The three runs have the same tests with the same results. This is the pytest output of run 1:

unit: 12 passed

collected 46 items / 34 deselected / 12 selected
test/unit/test_contract.py::test_every_sglang_import_is_declared[sgl_amd_mono/qwen38_a95b_mxfp4/check.py] PASSED
test/unit/test_contract.py::test_every_sglang_import_is_declared[sgl_amd_mono/qwen38_a95b_mxfp4/spec.py] PASSED
test/unit/test_contract.py::test_sglang_provides_what_the_plugin_reads[sgl_amd_mono/qwen38_a95b_mxfp4/check.py] PASSED
test/unit/test_contract.py::test_sglang_provides_what_the_plugin_reads[sgl_amd_mono/qwen38_a95b_mxfp4/spec.py] PASSED
test/unit/test_qwen38_a95b_mxfp4_spec.py::test_the_plugin_registers_the_spec PASSED
test/unit/test_qwen38_a95b_mxfp4_spec.py::test_config_reasons PASSED
test/unit/test_qwen38_a95b_mxfp4_spec.py::test_model_reasons[kw0] PASSED
test/unit/test_qwen38_a95b_mxfp4_spec.py::test_model_reasons[kw1] PASSED
test/unit/test_qwen38_a95b_mxfp4_spec.py::test_model_reasons[kw2] PASSED
test/unit/test_qwen38_a95b_mxfp4_spec.py::test_check_compares_the_rows_routed_alike PASSED
test/unit/test_qwen38_a95b_mxfp4_spec.py::test_model_reasons_name_the_layer PASSED
test/unit/test_qwen38_a95b_mxfp4_spec.py::test_aux_hidden_states_run_sglangs_forward PASSED
======== 12 passed, 34 deselected, 18 warnings in 17.93s ========

kernels: 6 passed

collected 6 items
test/kernels/test_qwen38_a95b_mxfp4_kernels.py::TestQwen38MoeLayout::test_scales PASSED
test/kernels/test_qwen38_a95b_mxfp4_kernels.py::TestQwen38MoeLayout::test_w13_tiles PASSED
test/kernels/test_qwen38_a95b_mxfp4_kernels.py::TestQwen38MoeLayout::test_w2_tiles PASSED
test/kernels/test_qwen38_a95b_mxfp4_kernels.py::TestQwen38MoeK1::test_k1 PASSED
test/kernels/test_qwen38_a95b_mxfp4_kernels.py::TestQwen38MoeK2::test_k2 PASSED
test/kernels/test_qwen38_a95b_mxfp4_kernels.py::TestQwen38MoePort::test_bitwise_against_sglang_43516 PASSED
======== 6 passed, 18 warnings in 113.77s (0:01:53) ========

e2e: 5 passed

collected 5 items
test/e2e/test_qwen38_a95b_mxfp4_server.py::TestMonokernelDecode::test_covered_steps_match_prefill PASSED
test/e2e/test_qwen38_a95b_mxfp4_server.py::TestMonokernelDecode::test_gsm8k_past_the_largest_width PASSED
test/e2e/test_qwen38_a95b_mxfp4_server.py::TestMonokernelVerify::test_gsm8k_past_the_largest_width PASSED
test/e2e/test_qwen38_a95b_mxfp4_server.py::TestMonokernelCheck::test_each_covered_layer_within_its_numerics PASSED
test/e2e/test_qwen38_a95b_mxfp4_server.py::TestMonokernelFault::test_a_stalled_rank_stops_the_server_without_tokens PASSED
======== 5 passed, 18 warnings in 1048.64s (0:17:28) ========

Each full e2e log is about 2 to 4 MB, because it holds the server output. I can attach the logs if you need them.