Repository navigation
[AMD] integrate kimi k3 monokernel - #43481
Closed
billishyahao wants to merge 1 commit into
Closed
billishyahao wants to merge 1 commit into
billishyahao wants to merge 1 commit into
Conversation
5 tasks done
6 of 7 tasks
billishyahao
force-pushed
the
k3_mono_v2
branch
from
October 14, 2026 11:49
1ebfeb2 to
a99d535
Compare
Closed
5 tasks
Collaborator
|
Thanks! The monokernel framework has landed in SGLang (#44353):
|
5 tasks done
4 of 5 tasks
Collaborator
Author
|
Close this PR as it has been replaced by #45202 and sgl-project/sgl-amd-mono#103 |
HaiShaw
pushed a commit
to sgl-project/sgl-amd-mono
that referenced
this pull request
Nov 12, 2026
…ne K2 launch (#83) Part 1 of 5 of the split of #16. ## Motivation Qwen3.8-2.4T-A95B (`amd/Qwen3.8-2.4T-A95B-Quark-MXFP4`) has 92 decoder layers. Each layer has attention or Gated DeltaNet (GDN, a linear attention), then a sparse MoE. The MoE has a router that picks 10 of 512 MXFP4 experts, and one bf16 shared expert. In a decode step of 1 to 8 rows, SGLang runs the layer tail after the attention as many small launches and two TP all-reduces. At these batch sizes, the launches and the waits take most of the step time. This PR runs the tail of each layer as one persistent launch, K2, on each of the 8 ranks. The kernels are the K2 kernels of sgl-project/sglang#43516 (at `51b0c592b2`), which adapts them from vllm-project/vllm#60833. This PR moves them into the plugin (RFC sgl-project/sglang#43949), with two differences: - Each wait on another workgroup or GPU has a deadline. A failed wait stops the launch on every rank, and the server stops without tokens from that step. - This PR leaves out the fused shared expert (`fse`) and MXFP4 out_proj (`oq`) variants. The spec refuses those configurations at startup. Since sgl-project/sglang#43839, SGLang quantizes the GDN projections to MXFP4 by default. K2 reads out_proj in bf16, so start the server with `SGLANG_ROCM_GDN_ONLINE_MXFP4=0`. Later PRs add the MTP verify steps, K1 (the attention and GDN front of a layer), the fused shared expert and the MXFP4 GDN projections. ## Modifications | Area | Change | | --- | --- | | `qwen38_a95b_mxfp4/spec.py` | Adds the spec, its forward and its refusals. The spec covers the decode steps of 1 to 8 rows at TP8 on gfx950. SGLang runs the input norm and the attention or GDN, up to the projection input. Then one K2 launch runs the rest of the layer: out_proj (o_proj on the attention layers) and its all-reduce, the post-attention norm, the router, the shared and routed experts, and the MoE all-reduce. Other steps run SGLang's path. | | `qwen38_a95b_mxfp4/gfx950/` | Adds the K2 kernels (`layer.py`), the Kimi-K3 helpers that they use (`_k3.py`, by way of sgl-project/sglang#43481) and the MoE weight layout (`layout.py`). Each of the 8 build keys (widths 1 to 8) builds one kernel. A server builds only the widths of its CUDA graphs. With the test flags, these are 1, 2, 4 and 8. | | `qwen38_a95b_mxfp4/check.py` | Adds check mode (`SGLANG_AMD_MONO_CHECK=1`). Check mode runs SGLang's own layer code next to each K2 launch and compares the outputs. The spec's `NUMERICS` threshold is `k2.out.cosine` ≥ 0.99. | | `sgl_amd_mono/__init__.py` | Registers the spec in `SPECS`. | | Tests | Adds the model's tests in `test/unit`, `test/kernels` and `test/e2e` (results below). | | `README.md`, `NOTICE` | `README.md` adds the model row, the design text, the `SGLANG_ROCM_GDN_ONLINE_MXFP4=0` requirement and a usage command. `NOTICE` names the origins of the adapted code. | **Refusals.** The server stops at startup, with the reason in the log, for each of these configurations: - the fused shared expert (`--enforce-shared-experts-fusion`); - the MXFP4 GDN projections (`SGLANG_ROCM_GDN_ONLINE_MXFP4` not `0`); - speculative decoding, such as MTP (the spec does not accept the feature, so SGLang's monokernel guard refuses it); - a model config that does not match Qwen3.8 in `HF_CONFIG` (the hidden size, the GDN heads, the experts, the top-k, the expert sizes or SiLU); - a config without `norm_topk_prob`; - more layers than the mailbox tags hold; - a layer whose norm, router, shared expert, o_proj or out_proj is not bf16; - routed experts that are not MXFP4 in AITER's shuffled layout; - an image without FlyDSL. The kernels read SGLang's weights in place and make no weight copies. Each rank allocates one peer buffer, and one scratch buffer for each width. K2 does not replace SGLang's attention, GDN, embedding, final norm or logits processor. This PR does not change `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 | `test/e2e/qwen38_a95b_mxfp4_server.py`. The monokernel-on servers add `--monokernel-backend amd`. | | Branch commit | `b0c3818` (tag `q38v3-tested/1-k2`). The head `3daea98` is rebased onto `main` at `b10f3b7`. Its model code and tests differ from that commit only in the text of a comment, 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 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 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 | 1291 (0.9788) | 1289 (0.9773) | 10 | 8 | 0.815 | | strict-match | 1141 (0.8650) | 1143 (0.8666) | 103 | 105 | 0.945 | The arms differ by 2 questions on each filter, and McNemar's test finds no accuracy difference. Strict-match scores lower in both arms because of the answer format. In the off arm, all 151 questions that pass flexible-extract and fail strict-match have no `#### N` line. Their answers end in other forms, such as `**Answer: 45**`. ### GPQA Diamond and SWE-bench Pro, monokernels on and off These runs use `3daea98` in one image, with the flags of the GSM8K and speed servers plus `--default-chat-template-kwargs '{"reasoning_effort": "xhigh"}'`. Claude Code sends the effort `high`, which the chat template refuses. With this flag, the template renders its default effort, xhigh, for every request. The SWE-bench servers leave out `--disable-radix-cache`, because the turns of a task share long prefixes. Each arm starts a fresh server on its own node with 8× MI355X. #### GPQA Diamond The test uses sgl-eval 0.1.3 on all 198 questions, with 4 samples of each question. It samples at temperature 1.0, top-p 0.95, top-k 20 and min-p 0, with 131072 max tokens, and it sends 8 concurrent requests. | | Monokernels off | Monokernels on | Model card | | --- | --- | --- | --- | | avg@4, with the 95% interval over the questions | 93.43 ± 3.20 | 92.68 ± 3.24 | 92.6 | Over the 198 paired questions, on minus off is -0.76 points, with a 95% interval of ±1.10. Five questions score higher with the monokernels on, and nine with them off. The sign test gives p 0.424 and finds no difference. The model card value is the score in the Qwen3.8-Max column of [Qwen/Qwen3.8-2.4T-A95B](https://huggingface.co/Qwen/Qwen3.8-2.4T-A95B), and it is inside the 95% interval of both arms. With the monokernels off, 3 of the 792 samples stop at the 131072-token cap, and 5 do with them on. #### SWE-bench Pro The test runs 50 tasks of [SWE-bench Pro V2](https://github.com/scaleapi/SWE-bench_Pro-os/tree/66f92766bba642462d4bbe5479e83f91f9211862/v2), Scale's validated set of 642 tasks. The 50 tasks come from 11 repos. `random.Random(20261105).sample(sorted(ids), 50)` picks them from the 642 task names in `v2/tasks`. Harbor 0.23.0 runs Claude Code with the `locked_claude_code` agent of V2, 8 tasks at a time. Each task gets 50 minutes, network access only to the model server, and no WebFetch or WebSearch. Each turn can have 32768 output tokens and 31999 thinking tokens. The server samples with the defaults of the checkpoint (temperature 1.0, top-p 0.95, top-k 20). Each decode batch in the server logs has 1 to 8 running requests, so the decode steps take the K2 path. After the run, `patch_replay` grades each patch again on a fresh task image, and the table gives that grade. | | Monokernels off | Monokernels on | | --- | --- | --- | | Resolved | 46 of 50 (92%) | 48 of 50 (96%) | | Agent timeouts | 3 | 0 | | Median agent time per task | 18.2 min | 15.1 min | Over the 50 paired tasks, 3 pass only with the monokernels on and 1 only with them off. McNemar's exact test gives p 0.625 and finds no difference. One of the 3 is an infrastructure failure in the off arm. There, `docker compose` failed in the environment setup, and the agent did not start. In the task that passes only with the monokernels off, the `claude` process of the on arm was killed ("Killed") and left no patch. Fifty tasks find only a large change. The 95% interval of the resolved rate is ±7.6 points with the monokernels off and ±5.5 points with them on. We did not run the other 592 tasks of V2. ### Kernel checks (`test/kernels`, 8 GPUs, TP8) | Check | Result | | --- | --- | | The MoE weight layout (the w13 and w2 tiles and the scales) | Equal to SGLang's layout | | Compile every build key (widths 1 to 8) on every rank | All 8 keys compile | | Batch invariance on random TP8 shards. A row's `out` and `res_out` do not change with the other rows or the width. | 57 cases pass, and the ranks agree | | Graph replay of widths 8, 1, 3, 8, 2, 5 in one graph on one peer buffer | 200 replays are bit for bit equal to eager mode, and the ranks agree | | Fault: rank 3 stalls at its first wait | Rank 3 sets its deadline flag, the flag reaches every rank, and the launch returns in 0.52 s | | Port test: `out` and `res_out` against the sgl-project/sglang#43516 kernels, at widths 1 to 8 on every rank | Bit for bit equal | The kernel tests do not compare K2 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 (200, the e2e test) | Highest mean KL (2, 3, 7, 8 rows) | Rows that check mode skipped | Lowest `k2.out.cosine` | | --- | --- | --- | --- | --- | --- | --- | --- | | 1 | 12 passed | 5 passed | 4 passed | 0.970 | 0.0118 | 1 of 3956 | 0.999897 | | 2 | 12 passed | 5 passed | 4 passed | 0.975 | 0.0179 | 0 of 3956 | 0.998428 | | 3 | 12 passed | 5 passed | 4 passed | 0.975 | 0.0107 | 0 of 3956 | 0.999738 | Each run gave the kernel check results above. `k2.res_out.cosine` is 1 in all 4416 checks (3 runs of 1472). The e2e tests are: - GSM8K: 200 questions at 12 concurrent requests, with thinking, at temperature 1.0. The batch grows past 8 rows as it fills, and drops below 8 rows as it drains. - 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 K2 layer output is within `NUMERICS`. - 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 check-mode test runs without CUDA graphs, at 1, 3 and 8 rows, with 1472 checks in each run. In rare rows, K2 and SGLang pick a different 10th expert, because two experts score within one bf16 step. Check mode skips those rows for the layer output and counts them. The test fails if it skips more than 1% of the rows. 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 GSM8K and speed servers use the flags and environment of `test/e2e/qwen38_a95b_mxfp4_server.py`, plus `--tool-call-parser qwen3_coder` and `--watchdog-timeout 12000`. On each rank, the monokernel-on servers log "Qwen3.8 monokernel: 4 K2 launches (widths [1, 2, 4, 8])". | 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 | Spread between rounds | | --- | --- | --- | --- | --- | --- | | 1 | 11.06, 11.04 | 7.46, 7.46 | 90.5, 134.0 | 1.48× | 0.2% | | 2 | 11.61, 11.58 | 8.69, 8.69 | 86.2, 115.1 | 1.33× | 0.3% | | 4 | 13.79, 13.76 | 10.41, 10.41 | 72.6, 96.1 | 1.32× | 0.2% | | 8 | 17.38, 17.34 | 14.91, 14.89 | 57.6, 67.1 | 1.17× | 0.2% | | 16 (past the largest width) | 23.39, 23.35 | 23.34, 23.39 | 42.8, 42.8 | 1.00× | 0.2% | - The speedup uses the mean of the two rounds. - The spread is the largest difference between the two rounds of one arm. - 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 8 K2 build keys at `7c038cb` (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 two builds differ, a cell gives `stop=True / stop=False`. | Width | VGPR | VGPR spill | SGPR spill | LDS (bytes) | Scratch (bytes) | | --- | --- | --- | --- | --- | --- | | 1 | 203 | 0 | 22 / 123 | 27616 | 0 | | 2 | 203 / 205 | 0 | 65 / 26 | 46176 | 0 | | 3 | 209 | 0 | 23 / 27 | 62656 | 0 | | 4 | 213 | 0 | 26 / 29 | 79136 | 0 | | 5 | 209 | 0 | 22 / 27 | 95616 | 0 | | 6 | 219 | 0 | 23 / 27 | 112096 | 0 | | 7 | 213 | 0 | 23 / 28 | 128576 | 0 | | 8 | 221 | 0 | 26 | 145056 | 0 | - No build spills a VGPR or uses scratch. The SGPR spills go to VGPR lanes. - No build uses AGPRs. Each build has 106 SGPRs. - LDS grows with the width, to 145,056 of gfx950's 163,840 bytes at 8 rows. **Kernels before and after.** The table compares one decode step at batch 1 in two TP-0 profiles: monokernels off and monokernels on. Both servers ran the head `3daea98` 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. ### Kernel Replacement (layer tail, decode bs=1) | Part of the layer | Before: SGLang kernel | Calls per step | Time per call (µs) | After | | --- | --- | --- | --- | --- | | out_proj (69 GDN layers) or o_proj (23 attention layers) | `_gemm_a16_w16_kernel` (`SPLITK_BLOCK_SIZE_2048`) | 92 | 8.0 | One `qwen38_mono_k2_s1_a0` launch for each layer: 92 calls per step, 39.8 µs each | | The all-reduce after out_proj | `aiter::ar_ll_gfx9` | 92 | 7.6 | In K2 | | The post-attention norm | `_gemma_fused_add_rmsnorm_kernel` | 92 | 4.3 | In K2 | | The shared expert gate_up | `aiter::wv_splitk_small_fp16_bf16_kernel` | 92 | 8.0 | In K2 | | The shared expert SiLU and mul | `sgl_hip::activation::act_and_mul_kernel` | 92 | 4.3 | In K2 | | The shared expert down | `_gemm_a16_w16_kernel` (`SPLITK_BLOCK_SIZE_256`) | 92 | 4.1 | In K2 | | The router | `aiter::wv_splitk_small_fp16_bf16_kernel` | 92 | 8.0 | In K2 | | The top-k | `vllm::moe::topkGatingSoftmax` | 92 | 9.8 | In K2 | | The routed experts | `aiter::fmoe_bf16_pertokenMXfp4_g1u1_flat_vs_silu_16x32` | 92 | 19.2 | In K2 | | The shared expert gate, sigmoid and add | `_fused_gate_sigmoid_mul_add_kernel` | 92 | 4.2 | In K2 | | The MoE all-reduce | `aiter::ar_ll_gfx9` | 92 | 6.8 | In K2 | | The step epoch of the K2 waits | None | | | One `CUDAFunctorOnSelf_add<int>` per step (`self.epoch[:1].add_(1)` in `spec.py`) | | Before | Time (us) | After | Time (us) | | --- | ---: | --- | ---: | | `_gemm_a16_w16_kernel` (SPLITK_BLOCK_SIZE_2048, out_proj / o_proj, ×92) | 8.0 | `qwen38_mono_k2_s1_a0` (×92) | 39.8 | | `aiter::ar_ll_gfx9` (all-reduce after out_proj, ×92) | 7.6 | (in K2) | -- | | `_gemma_fused_add_rmsnorm_kernel` (×92) | 4.3 | (in K2) | -- | | `aiter::wv_splitk_small_fp16_bf16_kernel` (shared expert gate_up, ×92) | 8.0 | (in K2) | -- | | `sgl_hip::activation::act_and_mul_kernel` (×92) | 4.3 | (in K2) | -- | | `_gemm_a16_w16_kernel` (SPLITK_BLOCK_SIZE_256, shared expert down, ×92) | 4.1 | (in K2) | -- | | `aiter::wv_splitk_small_fp16_bf16_kernel` (router, ×92) | 8.0 | (in K2) | -- | | `vllm::moe::topkGatingSoftmax` (×92) | 9.8 | (in K2) | -- | | `aiter::fmoe_bf16_pertokenMXfp4_g1u1_flat_vs_silu_16x32` (×92) | 19.2 | (in K2) | -- | | `_fused_gate_sigmoid_mul_add_kernel` (×92) | 4.2 | (in K2) | -- | | `aiter::ar_ll_gfx9` (MoE all-reduce, ×92) | 6.8 | (in K2) | -- | | | | `CUDAFunctorOnSelf_add<int>` (×1 per step) | -- | | **Total** | **84.3** | **Total** | **39.8 (52.8% reduce)** | - The 11 kernels of a layer tail take 84.3 µs together. K2 takes 39.8 µs. - The step falls from 1669 to 750 kernels, and from 12.14 to 8.11 ms of kernel time. - SGLang still runs the embedding and its all-reduce, the input norms, the GDN or attention, the final norm, the LM head and the sampling. - Each screenshot selects one `step[DECODE bs=1]` slice and shows only its kernels (filter `Category = kernel`). Before (monokernels off), one decode step: <img width="1928" height="1172" alt="image" src="https://github.com/user-attachments/assets/977c336a-1228-4c3a-b711-e49b525828ce" /> After (monokernels on), one decode step: <img width="1928" height="1172" alt="image" src="https://github.com/user-attachments/assets/ccac7f6b-a326-4b25-b354-cf741bc899f8" /> E2E: 12.14 ms => 8.103 ms，33.3% reduce ## Checklist - [x] `pre-commit run --all-files` passes. - [x] Tests under `test/` cover the change. - [x] The SGLang commit or PR this was tested against is named above. - [x] 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`. --------- Co-authored-by: Cursor <cursoragent@cursor.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
This will serve as a PoC of mono kernels integration of kimi k3 on AMD. cc @HaiShaw @kkHuang-amd
Modifications
Accuracy Tests
Speed Tests and Profiling
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #38559256507
Latest PR Test (Extra): ❌ Run #38559256566
Latest PR Test (AMD ROCm 10): ❌ Run #38559256522