[Perf] [AMD] Kimi-K3: fused MoE routing, sort, quant and epilogue by aryatsc · Pull Request #46652 · sgl-project/sglang

aryatsc · 2026-10-26T20:00:55Z

Motivation

At K3 decode batch sizes, MoE routing, sorting, activation quant and the shared-expert epilogue were separate small launches, and launch-bound.

Modifications

Generic sorter hook. moe_sorting_small.py stays model-agnostic. It exposes SortExtension, a hook for the sort AITER runs inside fused_moe, scoped with the sort_extension(...) context manager. With no extension set, behavior is unchanged.
K3 plan in a k3_* module. k3_sort_plan.py holds K3SortPlan(SortExtension): the route-state handoff and the deferred shared-expert GEMM. It calls the K3 Triton kernels in k3_route_sort, k3_route_sort_mid and k3_moe_epi.
The K3 model drives it. kimi_k3.py routes with k3_route_sort_mid.route_quant when the shape gate passes, then runs the experts under k3_sort_plan(plan).
AITER imports are direct; there are no try/except fallbacks.
Tuning tables live in AITER. There is no runtime patch of AITER's tuned table. The DFlash-shaped decode rows are in [Config] [MoE] Kimi-K3 A8W4: retune the token 8/16 decode rows for DFlash verify batches ROCm/aiter#7481; without them K3 uses AITER's current rows and stays correct.

Gates are on shapes and gfx950, not on the model name. Kimi-K3.1 has the same MoE shapes at TP8 (896 experts, top-k 16, hidden 3584, per-rank intermediate 384, SiTU beta 4 / linear beta 25), so these paths apply to it too. Validated on a Kimi-K3.1 Pro TP8 deploy below.

Accuracy Tests

Test	Result
`test_k3_moe_sort_plan_e2e.py` (new, MI355X): routing, sort, MXFP8 quant, shared expert and the AITER MoE runner vs an fp32 reference, K3/K3.1 TP8 shapes, M=1..32	2 pass, 10 subtests
`test_moe_sorting_small.py`, `test_k3_sort_plan*.py`, K3 MoE kernel tests (MI355X)	23 pass, 211 subtests

Full K3 stack on MI355X TP8 (DFlash, prod config) vs B300, McNemar:

Benchmark	MI355X	B300	p
GPQA-D	89.4	87.9	0.55
AIME25	96.7	93.3	1.0
MATH500	92.2	92.2	1.0
MMLU-Pro	89.6	90.0	0.77
GSM8K	97.2	-	-

Kimi-K3.1 Pro on MI355X TP8 (DFlash, K3.1 bringup config), GSM8K (1319), this feature on vs off:

	On	Off	McNemar p
GSM8K	95.38	95.15	0.72 (17 vs 14 discordant)

Speed Tests and Profiling

Throughput change from turning this feature off, full K3 stack, MI355X TP8:

Workload	C=1	C=8	C=32
Agentic replay	-7.6%	-4.9%	-8.9%
90k prefix	-6.4%	-3.4%	-

Mean -6.2%.

Kimi-K3.1 Pro, MI355X TP8, output tok/min/GPU, one run per cell (0 failed requests):

Workload	On	Off	Off vs on
90k prefix, C=1	2184, 2162, 2170	2073, 2069	-4.6%
90k prefix, C=8	6013, 6071, 6012	5900, 5912	-2.1%
Agentic replay, C=1	1338, 1477	1301, 1240	-9.7%
Agentic replay, C=8	3619, 3437	3585, 3877	+5.8%
Agentic replay, C=12	3842, 3922	4012, 3994	+3.1%

The agentic replay does not pin output lengths, and two deploys of identical code differed by up to 23% at C=8 on K3, so the agentic rows are within noise; the fixed-length 90k prefix rows vary by under 1% between runs.

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.

CI States

Latest PR Test (Base): ❌ Run #42072711749
Latest PR Test (Extra): ❌ Run #42072711709
Latest PR Test (AMD ROCm 10): ❌ Run #42072711701

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

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> # Conflicts: # python/sglang/kernels/ops/moe/moe_sorting_small.py

rkarhila-amd

Please address the comments. Patching things gets perf quickly, but overriding the aiter behaviour is bad in a long run. Division of general and code-specific code is also not typical for this project.

In addition, please put the perf and accuracy numbers in tables so that human eyes have an easier time inspecting them quickly.

Lint seems to fail, but you need to update the code to latest anyway.

Also, Kimi-K3.1 is just out, so check if these changes are relevant for 3.1

rkarhila-amd · 2026-11-04T17:05:54Z

+
+
+@dataclass
+class K3SortPlan:


Please reconsider the ownership of the K3 plan and dispatch.

The k3_* Triton kernels belong here in kernels/ops/moe, and the K3 model should orchestrate them.

moe_sorting_small.py is shared, yet here it contains K3SortPlan, a K3-specific context variable, route-state handoff, and the shared-expert GEMM.

Could the generic sorter expose a model-agnostic sort-extension hook, with the K3 plan and implementation kept in a k3_* module? This preserves the generic sorter’s model independence while still allowing the Aiter-internal sort boundary to invoke the K3 fast path.

Done in b56a003. moe_sorting_small.py is model-agnostic again: it exposes a SortExtension hook, scoped with sort_extension(...), at the sort boundary AITER calls inside fused_moe. With no extension set, behavior is unchanged. K3SortPlan, the route-state handoff and the deferred shared-expert GEMM moved to kernels/ops/moe/k3_sort_plan.py. kimi_k3.py routes with k3_route_sort_mid.route_quant and runs the experts under k3_sort_plan(plan).

rkarhila-amd · 2026-11-04T17:17:36Z

+
+@functools.cache
+def apply_k3_moe_decode_cfg() -> None:
+    try:


Avoid try-excepts, they can make debugging more difficult (if anyone still makes that by hand).

The sglang images are built with specific releases of aiter, trust the correct version to be there. If the current version of aiter is not yet in sglang docker builds, wait for aiter release to be updated before pushing this PR to be merged.

Done in 97c4f52: the AITER imports are direct, with no try/except fallbacks. The PR relies on the AITER that ships in the image.

rkarhila-amd · 2026-11-04T17:23:42Z

@@ -0,0 +1,113 @@
+"""K3SortPlan: the model-owned inputs of the K3 sort kernels inside aiter fused_moe."""


My agent advices to reconsider the test plan:

The new behavior crosses K3 routing, the patched AITER sorter, MXFP8 quantization, shared-expert execution, and the MoE runner. The listed tests are predominantly unit-level.

Added test/registered/amd/test_k3_moe_sort_plan_e2e.py (c0e75fb). It runs K3 routing, the K3 sort plan, MXFP8 quant, the shared expert and AiterRunnerCore end to end, at the K3/K3.1 TP8 MoE shapes (896 experts, top-k 16, M=1..32), against an fp32 reference. It passes on MI355X: 2 tests, 10 subtests. With the rest of the K3 MoE GPU tests: 23 pass, 211 subtests.

With end-to-end test I kindly ask removing some of the overlapping tests:

Test Verdict Reason

test_k3_sort_plan_keeps_small_sort_at_513_experts Remove Same gate the M=8/12/16 cases already require: with a plan active, _small_sort_supported stays true past the expert-count cutoff. Those cases also require the kernel to take the GEMM. This one only reads the boolean, at E=513.

test_exit_runs_untaken_gemm in test_k3_sort_plan.py Remove Mocks shared_gemm_preact. test_untaken_shared_gemm_matches_standalone runs that same exit on a GPU and checks it against the direct kernel and against torch.

Removed both in 56c66e8: test_k3_sort_plan_keeps_small_sort_at_513_experts and test_exit_runs_untaken_gemm. The e2e test and the M=8/12/16 plan cases keep that coverage.

rkarhila-amd · 2026-11-04T17:28:26Z

+(M=16: ~139 vs 224), and at that shape the tuned M=16 stage1 (``_xcd4`` swizzle) is ~11 us
+slower than the plain variant. These rows were re-tuned on real K3 routing:
+same activation/weight quantization, same layouts, only the
+tile schedule changes, so outputs match the default kernels to fp32-reference noise.


I appreciate the effort and I can easily believe the statistics from dspark are different from the (probably) random sampling data data that is used in a lot of the tunings.

However patching the tuning table to me seems a very bad idea and I would recommend solving this in Aiter side, maybe having alternative speculative decoding tuning tables there.

Agreed. The runtime table patch is removed (k3_moe_cfg.py is deleted in d90bd7c). The retuned token 8/16 rows are in ROCm/aiter#7481, with a before/after table on DFlash-shaped and uniform routing (M=16 DFlash: 77.4 to 68.7 us; uniform: 100.9 to 98.6 us). Until that lands in the image, K3 uses AITER's current rows and stays correct.

…K3 plan to k3_sort_plan Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

… with try/except Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> # Conflicts: # python/sglang/kernels/ops/moe/__init__.py

…ve in AITER Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

rkarhila-amd

Looks much better, thank you. A few tests seem to be overlapping, please consider removing those.

#44987 seems to be on collision course with this one but is at draft state currently.

rkarhila-amd · 2026-11-05T15:37:47Z


 @functools.cache
 def apply_aiter_small_moe_sort_patch() -> bool:
    """Patch a stock aiter so decode-sized MoE sorting (+ stage1 mxfp8 quant) is one launch.


I have strong opinions against patching aiter functionality in sglang side, there are big risks in maintainability. But this damage is already done in #36559 and continued in #46691 and so it is not a blocker for this PR.

Understood, thanks. This PR no longer patches AITER: the runtime decode-table patch is gone (the rows are in ROCm/aiter#7481), and the AITER imports are direct. The remaining hook here is the model-agnostic SortExtension at the sort boundary, which does nothing unless a model sets it.

rkarhila-amd · 2026-11-05T15:50:11Z

@@ -0,0 +1,113 @@
+"""K3SortPlan: the model-owned inputs of the K3 sort kernels inside aiter fused_moe."""


With end-to-end test I kindly ask removing some of the overlapping tests:

Test Verdict Reason

test_k3_sort_plan_keeps_small_sort_at_513_experts Remove Same gate the M=8/12/16 cases already require: with a plan active, _small_sort_supported stays true past the expert-count cutoff. Those cases also require the kernel to take the GEMM. This one only reads the boolean, at E=513.

test_exit_runs_untaken_gemm in test_k3_sort_plan.py Remove Mocks shared_gemm_preact. test_untaken_shared_gemm_matches_standalone runs that same exit on a GPU and checks it against the direct kernel and against torch.

aryatsc · 2026-11-05T17:05:16Z

@rkarhila-amd on your Kimi-K3.1 question: I ran it on a Kimi-K3.1 Pro TP8 MI355X deploy (our K3.1 bringup config) with this feature on and off.

	On	Off
90k prefix C=1 (tok/min/GPU)	2184, 2162, 2170	2073, 2069
90k prefix C=8	6013, 6071, 6012	5900, 5912
GSM8K (1319)	95.38	95.15 (McNemar p = 0.72)

So it applies to K3.1 unchanged: about +4.6% at C=1 and +2.1% at C=8 on the fixed-length workload, with no accuracy change and 0 failed requests. The agentic replay numbers are in the description; they are within run-to-run noise. The shape gates are the same for both models, so there's no K3.1-specific code.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

aryatsc · 2026-11-05T17:15:52Z

@rkarhila-amd on #44987 (draft, @kevin-mii): it overlaps with this PR in the K3 decode routing in kimi_k3.py and in the route-sort handoff to the AITER MoE runner (moe_sort_handoff.py there, k3_sort_plan.py + SortExtension here). Both fuse the router top-k with the MoE sort for the same shapes, so only one should land. Happy to coordinate with @kevin-mii: if #44987 moves forward first, I'll rebase this PR onto it and keep only the parts it doesn't cover (the shared-expert GEMM overlap and the sort-plan hook).

rkarhila-amd

Thanks for the modifications. I don't see any blocking problems so I'll pass it on,

rkarhila-amd · 2026-11-06T14:02:34Z

+        # stable rank: pairs of earlier tokens on the same expert
+        rank = tl.zeros((TOPK,), tl.int32)
+        n_prev = t * TOPK
+        for s in range(0, n_prev, RANK_TILE):


My agent tried to convince me that a standard python range would crash the compiler and insisted on tl.arange.
But in practice this compiled well so no issue.

rkarhila-amd · 2026-11-06T14:04:50Z

+                ids, w = self._ref_topk(logits)
+                with k3_sort_plan():
+                    out = self._routed(x, ids, w)
+                self._assert_rel(out, self._routed(x, ids, w), 1e-2, "routed")


Agent raised a concern:

Both sides of the comparison call self._routed(...), which uses the same AITER/K3 implementation. Therefore, if the implementation produces an incorrect result consistently, both calls can produce the same incorrect result and the test still passes.

A stronger test would compare against an independent reference, such as:

a PyTorch fp32 MoE implementation;
AITER with the K3 optimization disabled;
independently computed expected sorting/permutation and quantization outputs.

I'm satisfied that you test the k3_sort_plan wrapper, up to you if you want to follow this.

[ROCm][Kimi-K3] Fused MoE routing, sort, quant and epilogue

d7b697b

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

aryatsc requested review from BBuf, DarkSharpness, Edwardf0t1, Fridge003, HaiShaw, HydraQYH, Ying1123, celve, ch-wan, ispobock, merrymercy and yuan-luo as code owners October 26, 2026 20:00

github-actions Bot added the jit-kernel label Oct 26, 2026

Merge branch 'main' into devin/pr-moe

32b3f91

aryatsc mentioned this pull request Oct 26, 2026

[Perf] [AMD] Kimi-K3: static FP8 dense layers and tuned decode GEMMs #46653

Open

3 of 5 tasks

aryatsc added 2 commits October 26, 2026 18:51

Merge branch 'main' into devin/pr-moe

c1c555a

Merge branch 'main' into devin/pr-moe

8247b9d

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> # Conflicts: # python/sglang/kernels/ops/moe/moe_sorting_small.py

This comment has been minimized.

Sign in to view

rkarhila-amd suggested changes Nov 4, 2026

View reviewed changes

aryatsc and others added 6 commits November 4, 2026 23:57

[AMD] MoE sort: expose a model-agnostic SortExtension hook; move the …

b56a003

…K3 plan to k3_sort_plan Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

[AMD] MoE sort: import the image's aiter directly instead of guarding…

97c4f52

… with try/except Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

[AMD] MoE runner: use the module-level zero_copy_context in run()

894600c

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

[AMD] Test the K3 MoE sort plan end to end through the AITER runner

c0e75fb

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

Merge remote-tracking branch 'origin/main' into w_pr-moe

4179511

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> # Conflicts: # python/sglang/kernels/ops/moe/__init__.py

[AMD] Kimi-K3: drop the runtime AITER decode-table patch; the rows li…

d90bd7c

…ve in AITER Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

aryatsc mentioned this pull request Nov 5, 2026

[Config] [MoE] Kimi-K3 A8W4: retune the token 8/16 decode rows for DFlash verify batches ROCm/aiter#7481

Open

3 of 10 tasks

rkarhila-amd suggested changes Nov 5, 2026

View reviewed changes

[AMD] Kimi-K3: drop two tests that overlap the e2e MoE coverage

56c66e8

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

rkarhila-amd approved these changes Nov 6, 2026

View reviewed changes

		@@ -0,0 +1,113 @@
		"""K3SortPlan: the model-owned inputs of the K3 sort kernels inside aiter fused_moe."""

Test	Verdict	Reason
test_k3_sort_plan_keeps_small_sort_at_513_experts	Remove	Same gate the M=8/12/16 cases already require: with a plan active, _small_sort_supported stays true past the expert-count cutoff. Those cases also require the kernel to take the GEMM. This one only reads the boolean, at E=513.
test_exit_runs_untaken_gemm in test_k3_sort_plan.py	Remove	Mocks shared_gemm_preact. test_untaken_shared_gemm_matches_standalone runs that same exit on a GPU and checks it against the direct kernel and against torch.



		@dataclass
		class K3SortPlan:

Conversation

aryatsc commented Oct 26, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

CI States

Uh oh!

This comment has been minimized.

rkarhila-amd left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rkarhila-amd left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aryatsc commented Nov 5, 2026

Uh oh!

aryatsc commented Nov 5, 2026

Uh oh!

rkarhila-amd left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

aryatsc commented Oct 26, 2026 •

edited by github-actions Bot

Loading