Skip to content

[AMD] Perf Kimi-K3 MoE optimization - #33838

Merged
HaiShaw merged 7 commits into
sgl-project:mainfrom
HaiShaw:perf/k3_moe-opt
Sep 3, 2026
Merged

HaiShaw merged 7 commits into
sgl-project:mainfrom
HaiShaw:perf/k3_moe-opt

Conversation

@kkHuang-amd

@kkHuang-amd kkHuang-amd commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator

Motivation

Kimi-K3 SiTUv2 MXFP4 MoE requires different preshuffled weight layouts for different activation quantization modes:

  • A16W4 and A8W4 use AITER's GU-interleaved shuffle_weight_a16w4 layout.
  • A4W4 uses the generic separated layout.
  • When both A8W4 and A4W4 flags are set, AITER gives A8W4 precedence.

AITER dependency

This PR requires ROCm/aiter#4534 (6dc26b7a817bba2ae92ff22a7a63c0c024f4c7e5) or a newer AITER revision. That PR provides the optimized Kimi-K3 Opus/FlyDSL A8W4 SiTUv2 MoE dispatch used by the benchmarks below and depends on FlyDSL 0.3.0.

The A4W4 SiTUv2 mode was introduced by ROCm/aiter#4463, which is already included in the required #4534 revision.

Modifications

  • Make Kimi-K3 SiTU weight preprocessing follow AITER's activation-mode precedence.
  • Keep A16W4/A8W4 on the GU-interleaved preshuffle path.
  • Route A4W4 through the generic separated shuffle path.
  • Preserve A8W4 precedence when both AITER mode flags are enabled.
  • Cache the AITER routing correction-bias dtype conversion and invalidate it when the loaded parameter changes.
  • Return the explicit published SiTU MoE output buffer instead of a distinct FFI wrapper/allocation.
  • Add registered AMD/CPU/CUDA tests for layout precedence, correction-bias caching, reload invalidation, and the zero-copy output contract.

Accuracy Tests

Environment: 8x MI35x/gfx950, TP8, public Kimi-K3 checkpoint.

  • Focused layout/copy/cache unit tests: 11 passed.
  • AITER strict K3 A4W4 oracle with SiTU beta 4.0 / linear beta 25.0: all 17 K3 cases passed.
  • A4W4 GSM8K 50: 1.000.
  • A4W4 GSM8K 200 after the fix: 0.980 (before fix: 0.005).
  • A8W4 GSM8K 200: 0.980.

Speed Tests and Profiling

Required AITER revision: 6dc26b7a8 from ROCm/aiter#4534.

Server command:

SGLANG_USE_AITER=1 \
SGLANG_AITER_K3_OPT=1 \
AITER_FLYDSL_FORCE=1 \
AITER_SITUV2_A8W4=1 \
AITER_SITUV2_A4W4=0 \
SGLANG_K3_FLYDSL_AR_NORM=1 \
sglang serve \
  --model-path /data/models/Kimi-K3 \
  --trust-remote-code \
  --tp 8 \
  --attention-backend triton \
  --dtype bfloat16 \
  --mem-fraction-static 0.85 \
  --cuda-graph-max-bs-decode 256 \
  --host 127.0.0.1 \
  --port 8000 \
  --disable-radix-cache \
  --reasoning-parser kimi_k3 \
  --tool-call-parser kimi_k3

Client workload:

  • Random dataset, random-range-ratio=1.0.
  • Fixed input/output lengths: 8192 / 1024.
  • num-prompts = 8 * concurrency.
  • warmup-requests = 2 * concurrency.

Final results with the MoE copy optimization, compared with the original MI35x reference:

Concurrency Total throughput (tok/s) Delta Per GPU (tok/s/GPU) Median E2EL (ms) Median TTFT (ms) Median TPOT (ms) Median ITL (ms)
2 882.44 +2.95% 110.31 20887.60 852.91 19.58 19.39
4 1594.32 +4.75% 199.29 23121.70 1609.19 21.04 20.67
8 2527.54 +1.18% 315.94 29152.66 2552.72 25.97 24.53
16 3893.59 +2.06% 486.70 37909.76 4785.77 32.39 28.88
32 5381.41 +2.68% 672.68 54853.02 8985.35 44.81 37.19

At concurrency 32 versus the original reference, total token throughput improves by 2.68%, median E2EL by 2.55%, TTFT by 3.57%, TPOT by 2.54%, and ITL by 2.29%.

The MoE copy change was also measured against the immediately preceding sweep with identical server/client settings:

Concurrency Throughput delta E2EL delta TTFT delta TPOT delta ITL delta
2 +0.86% -0.84% -10.61% -0.15% -1.12%
4 +1.57% -1.17% -0.46% -1.17% -1.15%
8 +0.93% -0.87% -0.34% -0.99% -1.01%
16 +0.64% -0.65% -0.10% -0.74% -0.79%
32 +0.42% -0.39% -0.57% -0.33% -0.53%

Throughput improves by 0.42%–1.57% and TPOT/ITL improve at every tested concurrency after adding the copy optimization. The C2 TTFT delta is based on a small 16-request sample and should be treated as higher variance.

A4W4 is now numerically correct but is not the performance default for this workload: its median output throughput was 530.84 tok/s, versus 537.33 tok/s for post-fix A8W4.

Checklist

  • Format code with pre-commit (pre-commit run --all-files passes).
  • Add registered unit tests.
  • Documentation is not required; this does not add a user-facing interface.
  • Provide accuracy and speed benchmark results.
  • 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 #32319934749
Latest PR Test (Extra): ❌ Run #32319934659
Latest PR Test (AMD ROCm 7.2): ❌ Run #32319934688

Your Name and others added 2 commits August 6, 2026 07:03
Use the separated MXFP4 shuffle for A4W4 while preserving A8W4 precedence so real-checkpoint MoE outputs remain numerically valid.
Cache the AITER routing-bias cast and make the SiTU runner honor its published output buffer so graph replay avoids redundant per-layer materialization.
wunhuang and others added 2 commits August 19, 2026 03:45
Keep the Kimi-K3 SiTU zero-copy output contract while integrating the latest main changes.

Co-authored-by: Cursor <cursoragent@cursor.com>
@kkHuang-amd

Copy link
Copy Markdown
Collaborator Author

@amd-bot ci-status

@amd-bot

amd-bot commented Aug 19, 2026

Copy link
Copy Markdown

@kkHuang-amd

CI Status for PR #33838

Merge verdict: 🔴 Do not merge. This PR breaks CI on every platform (CUDA, CPU, XPU, NPU, AMD ROCm 7.2) with a single self-inflicted root cause: two of the three new test files are missing the required if __name__ == "__main__": entry, and the test-registration validator rejects them at collection time. Because this fails in the earliest stage (base-a / stage-a), all downstream stages were fast-fail skipped — so PR CI is incomplete and the PR's actual MoE code changes were never exercised.

Caution

PR CI is incomplete AND the changed code is untested. The registration validator fails during base-a/stage-a on all vendors, so base-b/base-c (and the AMD GPU MoE path) were skipped by fast-fail — those are not tested. Separately, the whole MXFP4 SiTU change in mxfp4.py is gated behind AITER_SITUV2_A8W4/AITER_SITUV2_A4W4 and needs a new AITER revision (ROCm/aiter#4534); the PR even carries the amd-aiter-not-ready label. Even after fixing the __main__ entries, no PR-CI job on this SHA verifies the end-to-end Kimi-K3 MoE path. Fix the test files, then re-run and confirm the covering suite actually executes.

Changed files: moe/topk.py (+30/-1), quantization/mxfp4.py (+26/-9), test_topk_correction_bias_cache.py (+72), test_mxfp4_situ_output.py (+97), test_mxfp4_situ_weight_layout.py (+31)

Executed CI failure attribution: AMD: 1 executed failure (1 related) · Others: ~12 executed failures (10 related, 1 infra, plus gate/cascade). All "related" failures share one root cause.

AMD Executed Failures

Job Test File Test Function Error Related? Why
stage-a-test-1-gpu-small-amd-rocm720 test/registered/unit/layers/quantization/test_mxfp4_situ_weight_layout.py N/A (collection) ValueError: missing 'if __name__ == "__main__":' entry 🔴 New test file added by this PR lacks the required entrypoint; validator rejects it
call-gate / pr-gate (Extra AMD) N/A N/A Missing required label 'run-ci-extra' → exit 1 🟢 By-design opt-in gate, not a code failure
wait-for-stage-a-amd-rocm720, pr-test-amd-rocm720-finish N/A N/A fast-fail cascade 🟢 Cascade from the stage-a failure above

Other Executed Failures

Job Test File Test Function Error Related? Why
base-a-test-cpu (0–8) test/registered/unit/layers/quantization/test_mxfp4_situ_output.py N/A (collection) ValueError: missing 'if __name__ == "__main__":' entry 🔴 Same root cause; new file lacks entrypoint
stage-a-test-1-gpu-xpu test_mxfp4_situ_output.py N/A same ValueError 🔴 Same root cause
base-a-test-1-npu-a2 / run (0) (+ base-b NPU shards) test_mxfp4_situ_output.py N/A same ValueError 🔴 Same root cause
multimodal-gen-test-1-npu-a3 N/A N/A modelscope 404: openmoss/MOVA-360p model_index.json not found 🟢 Infra/model-download failure, unrelated to this PR
base-a-test-1-gpu-small (0) N/A N/A Fast-fail: skipping — root cause base-a-test-cpu 🟢 Cascade skip
call-gate / pr-gate (Extra) N/A N/A Missing required label 'run-ci-extra' 🟢 By-design opt-in gate

Details / what to do before merge

Root cause (🔴, single fix): The test-registration validator requires every registered pytest-style file to end with an explicit runner. Two new files are missing it:

  • test_mxfp4_situ_output.py — add sys.exit(pytest.main([__file__, "-v"]))
  • test_mxfp4_situ_weight_layout.py — add sys.exit(pytest.main([__file__, "-v"]))

(test_topk_correction_bias_cache.py already has if __name__ == "__main__": unittest.main() — verified 37286f7.) Verified in log: ROCm720 line 4459, CPU line 1197.

Coverage — must do before merge:

  1. Fix the two __main__ entries and push, so the suite actually collects and downstream (base-b/base-c) runs instead of being skipped.
  2. The MXFP4 SiTU code path (mxfp4.py) is AITER/env-gated — the new unit test only checks the pure _aiter_situ_uses_gu_interleaved_weights() helper; the real Kimi-K3 A4W4/A8W4 layout correctness fix is not validated by PR CI. Verify it manually against the required AITER revision (ROCm/aiter#4534) before merge (amd-aiter-not-ready label suggests this hasn't happened yet).
  3. Unrelated NPU multimodal-gen 404 (infra) and the Extra pr-gate failures (missing run-ci-extra label) are not blockers; add run-ci-extra only if you want the Extra pipeline signal.

Generated by amd-bot using Claude Code CLI

@kkHuang-amd

Copy link
Copy Markdown
Collaborator Author

@amd-bot ci-status

@amd-bot

amd-bot commented Aug 19, 2026

Copy link
Copy Markdown

@kkHuang-amd

CI Status for PR #33838

Merge verdict:Not safe to merge. This PR modifies MoE routing (topk.py) and MXFP4 quantization (mxfp4.py), and AMD CI — which actually does exercise those paths — shows two failures inside the changed code: a NaN in fused-MoE output (test_fused_moe.py) and an MXFP4 MoE accuracy regression (test_quark_mxfp4.py, 0.914/0.916 < 0.92 threshold). Separately, NVIDIA PR Test never completed — all base-b shards were fast-fail-skipped after an unrelated infra failure, so NVIDIA gives no real signal, and PR Test (NPU) is still running.

Caution

PR CI is incomplete AND the changed code is failing where it was exercised.

  • AMD ROCm 7.2 ran the MoE/MXFP4 tests and they failed (🔴 below) — this is the code this PR changes.
  • NVIDIA PR Test base-b cascade-skipped 13 shards after one infra failure (test_hicache_variants.py: GPU still not idle, 69.5 GiB leaked by pid 916002), so the MXFP4 combine-output change is unverified on NVIDIA / gpt-oss. Update the branch & rerun (or add bypass-fastfail sparingly) to get full NVIDIA signal.
  • PR Test (NPU) base-c-test-perf-2-npu-a3 is still in progress — do not treat NPU as passed.

Changed files: topk.py (+30/-1), mxfp4.py (+26/-9), 3 new unit tests under test/registered/unit/layers/ (+221)

Executed CI failure attribution: AMD: 4 failures (2 🔴 related · 1 🟡 · 1 🟢) · Others: 4 executed failures (all 🟢 infra/env), NVIDIA/XPU/Xeon/NPU · NPU still pending

AMD Executed Failures

Job Test File Test Function Error Related? Why
stage-b-1gpu-small (mi300, 8) test/registered/moe/test_fused_moe.py test_single_expert_routing (m=1, rsf=1.0) AssertionError: Tensor-likes are not close! ... nan (100%) 🔴 Fused-MoE numerics produce NaN; PR changes MoE routing/correction_bias handling in topk.py (aiter grouped-topk path)
stage-b-1gpu-small (mi35x-1) test/registered/quant/test_quark_mxfp4.py test_gsm8k (TestFP8ToMXFP4MOETP1) AssertionError: 0.916 not greater than 0.92 🔴 MXFP4 MoE accuracy just under threshold; PR rewrites mxfp4.py weight-layout selection + combine-output — the exact numerics this test checks
stage-b-large-8gpu-mi35x-disagg test/registered/amd/disaggregation/test_disaggregation_basic.py test_gsm8k POST /generate → 500; ConnectionRefused 127.0.0.1:11200; KeyError: 'answer' (0/4) 🟡 Decode/bootstrap server unreachable → could be a PR-induced MoE forward crash or a disagg setup/infra issue; stack not conclusive
stage-b-2gpu-large (mi300-2gpu, 1) test/registered/perf/test_bench_one_batch_2gpu.py test_torch_compile_tp2_bs1 ImportError: cannot import name 'specialize_impl' from 'triton.runtime.jit' → HIP graph-capture fail 🟢 torch.compile/Triton version mismatch in the image; unrelated to the diff

Other Executed Failures

Job Test File Test Function Error Related? Why
base-b-1gpu-large (2) (NVIDIA root cause) test/registered/hicache/test_hicache_variants.py TestHiCachePage.setUpClass RuntimeError: GPU still not idle after 30s: GPU0 uses 70.02 GiB (pid 916002) 🟢 Leaked GPU memory from a prior job; runner-state infra flake in hicache (not MoE/quant). Fast-failed 13 sibling base-b shards
stage-a-1gpu-xpu test/registered/xpu/test_xpu_basic.py benchmark parse OSError: libsycl.so.8: cannot open shared object file 🟢 Missing SYCL runtime lib in XPU image; env issue
build-test (xeon-gnr, base-b-test-cpu) test/registered/cpu/test_subblock_sparse_attention.py ModuleNotFoundError: No module named 'imageio' + HF 401 GatedRepoError 🟢 Missing dep + gated-model auth; env/infra
multimodal-gen-test-1-npu-a3 sglang/multimodal_gen/.../test_server_1_npu.py test_diffusion_generation[*] openai.APITimeoutError + custom container implementation failed 🟢 NPU diffusion server timeout / runner-container failure; unrelated to diff

(Cascade/gate jobs collapsed and excluded from counts: wait-for-base-b, pr-test-finish, wait-for-stage-b-amd-rocm720, pr-test-amd-rocm720-finish, finish, call-gate / pr-gate, pr-test-extra-finish, pr-test-amd-extra-finish.)

Details / what to do before merge

  • 🔴 test_fused_moe.py::test_single_expert_routing (NaN): highest priority. The aiter grouped-topk correction_bias_for_dtype change and/or the MXFP4 combine-output change (return StandardCombineInput(hidden_states=symm_output)) plausibly feed NaN into the fused-MoE output. Run this test locally on MI300 with and without the diff to confirm it's a regression vs. pre-existing.
  • 🔴 test_quark_mxfp4.py::test_gsm8k (0.914/0.916 < 0.92): the PR's whole point is fixing MXFP4 SiTU weight-layout selection; an accuracy dip on the MXFP4 MoE gsm8k gate directly implicates it. Confirm the same test passes on main on mi35x — if it does, this is a real numerical regression from the layout logic.
  • 🟡 disagg test_gsm8k: get the prefill/decode server-side stack (the 500 origin) to decide infra vs. MoE-forward crash; the KeyError: 'answer' is just the downstream symptom of empty generations.
  • Coverage gap (NVIDIA): the mxfp4.py combine-output change can affect gpt-oss on NVIDIA, but base-b never ran to completion. Rebase/rerun PR Test so the NVIDIA MXFP4/MoE suites actually execute before merge.
  • 🟢 XPU / Xeon / NPU / hicache failures are infra/env and not blockers for this diff, but they do mean those pipelines aren't giving a clean pass either.

Generated by amd-bot using Claude Code CLI

@kkHuang-amd

Copy link
Copy Markdown
Collaborator Author

@amd-bot ci-status

@amd-bot

amd-bot commented Aug 20, 2026

Copy link
Copy Markdown

@kkHuang-amd

CI Status for PR #33838

Merge verdict: Do not merge yet — PR CI is incomplete and the PR's core code path is unverified. The whole point of this PR is the AMD Kimi-K3 mxfp4 SiTU MoE optimization, but AMD stage-b never ran (the two new AMD-only mxfp4 tests were skipped), the NVIDIA main run was cancelled mid-flight before base-c, and NPU/Xeon are still pending. The only two executed failures (XPU, NPU) are unrelated to this PR. So "0 related failures" here does not mean safe.

Caution

This PR's core AMD code is not exercised by any completed PR-CI test. mxfp4.py changes are entirely AMD/aiter SiTU (Kimi-K3) paths, and the two new tests test/registered/unit/layers/quantization/test_mxfp4_situ_output.py + test_mxfp4_situ_weight_layout.py register only on AMD stage-b (register_amd_ci(stage="stage-b")). AMD stage-b was skipped because wait-for-stage-a-amd-rocm720 timed out (~4h), and no regular PR Test (AMD) run exists for this SHA. The AMD-only _use_aiter-guarded correction_bias_for_dtype call in topk.py:2163 is likewise unexercised. Green does NOT verify the mxfp4 optimization — an AMD maintainer must re-run AMD CI (stage-b must complete) before merge.

Warning

PR CI is also incomplete on other vendors: the NVIDIA main run 32319934749 was cancelled at ~05:17 mid base-b, so base-c (multi-GPU) was never reached; PR Test (NPU) 32319934529 and PR Test (Xeon) 32319934826 are still queued; and both Extra workflows fast-failed the label gate (missing run-ci-extra) so they never tested anything.

Changed files: topk.py (+30/-1), mxfp4.py (+26/-9), test_topk_correction_bias_cache.py (+72), test_mxfp4_situ_output.py (+112), test_mxfp4_situ_weight_layout.py (+37)

Executed CI failure attribution: AMD: 0 executed test failures (stage-b skipped — see caution) · Others: 2 failures (0 related) · NVIDIA base-a/b passed but run cancelled before base-c · NPU/Xeon pending

Other Executed Failures

Job Test File Test Function Error Related? Why
PR Test (XPU) / stage-b-test-1-gpu-xpu test/registered/attention/test_chunk_gated_delta_rule.py TestChunkGatedDeltaRule.test_padded_state_index_is_skipped [all_padded] AssertionError: padded rows wrote into the state pool (2/2 attempts) 🟢 Gated-delta-rule linear-attention state-pool test; PR only touches MoE topk + mxfp4 quant — no shared code path
PR Test (NPU) / multimodal-gen-test-1-npu-a3 sglang/multimodal_gen/test/server/ascend/test_server_1_npu.py TestDiffusionServerOneNpu::test_diffusion_generation[flux/ernie/z_image...] Perf-threshold asserts (assert 3026 <= 443) + custom container implementation failed 🟢 NPU diffusion image-gen latency thresholds + runner/container infra failure; unrelated to MoE/quant changes

Coverage detail (what ran vs. what didn't)

Changed code Covering test Ran? Result
topk.py correction_bias_for_dtype cache logic test_topk_correction_bias_cache.py (CPU/CUDA/AMD stage-a) ✅ ran on NVIDIA base-a 1-gpu-small passed ("passed": true)
topk.py AMD _use_aiter-guarded call (select_experts) integration only, AMD grouped-topk path ❌ AMD stage-b skipped unverified
mxfp4.py _aiter_situ_uses_gu_interleaved_weights / GU-interleave layout test_mxfp4_situ_weight_layout.py (AMD stage-b) ❌ skipped unverified
mxfp4.py finalized-kernel output-pointer contract test_mxfp4_situ_output.py (AMD stage-b) ❌ skipped unverified

Details / what to do before merge

  • Get AMD CI to complete (mandatory). PR Test ROCm 7.2 (AMD) stage-a passed but wait-for-stage-a-amd-rocm720 hit Timeout waiting for stage-a-amd-rocm720 jobs (~4h — a matrix partition never got a runner), which skipped all of stage-b. Re-run the AMD workflow so stage-b-test-1-gpu-small-amd actually executes test_mxfp4_situ_output.py + test_mxfp4_situ_weight_layout.py. This is the only signal that verifies the PR.
  • Re-run / wait out the NVIDIA main run32319934749 was cancelled before base-c; base-a and completed base-b shards passed, but multi-GPU (base-c) never ran.
  • Extra workflows are unlabeledPR Test Extra and Extra (AMD) failed the gate with Missing required label 'run-ci-extra'. Add the label (e.g. /tag-and-rerun-ci extra) if extra coverage is wanted; otherwise these are expected non-failures.
  • XPU/NPU reds need no action from this author — both are unrelated (XPU gated-delta-rule test; NPU diffusion perf/infra). Worth flagging to the XPU/NPU owners separately, but they don't block this PR's correctness.

Generated by amd-bot using Claude Code CLI

@kkHuang-amd

Copy link
Copy Markdown
Collaborator Author

@amd-bot ci-status

@amd-bot

amd-bot commented Aug 20, 2026

Copy link
Copy Markdown

@kkHuang-amd

CI Status for PR #33838 — [AMD] Perf Kimi-K3 MoE optimization

Merge verdict: Not ready to merge. PR Test ROCm 7.2 (AMD) is still in progress (3 stage-b-1-gpu-small shards running), so AMD PR CI is incomplete — do not treat it as green. NVIDIA (base-a/base-b), AMD stage-a, Arm64/MLX/MUSA/sgl-router all passed, and the PR's new unit tests passed where they ran. The executed red jobs are dominated by infra/unrelated failures (triton version mismatch, HF gated-repo 401, NPU perf thresholds, XPU attention), but two AMD MoE/quant tests need a look. Separately, the PR's actual value — real Kimi-K3 A4W4/A8W4 correctness — is not exercised e2e by PR CI.

Warning

Two coverage gaps, both matter here:

  1. AMD is not done. PR Test ROCm 7.2 (AMD) is still running; a later shard could still fail. Wait for completion before concluding.
  2. The core fix is not e2e-tested. The bug this PR fixes (SiTU MXFP4 A4W4/A8W4 weight-layout corruption, cosine-sim 0.0057/GSM8K 0.005) only reproduces with AITER #4534 + a real Kimi-K3 checkpoint, which no PR test runs. The 3 new tests are unit-level only: test_mxfp4_situ_weight_layout.py checks env-var precedence logic, test_mxfp4_situ_output.py checks a mocked FFI output-buffer identity, and test_topk_correction_bias_cache.py checks the dtype cache. Green here does not verify the real-checkpoint fix — the author should run the Kimi-K3 A4W4/A8W4 GSM8K comparison manually before merge.

Changed files: layers/moe/topk.py (+30/-1), layers/quantization/mxfp4.py (+26/-9), + 3 new unit tests (+221)

Executed CI failure attribution: AMD: 4 executed failures (0 clearly PR-caused, 2 to check) · Others: 4 executed failures (0 related) · AMD ROCm720 pending (3 shards). Gate/*-finish jobs are fast-fail cascades, collapsed below.

AMD Executed Failures

Job Test File Test Function Error Related? Why
stage-b-1gpu-small mi35x test/registered/quant/test_quark_mxfp4.py TestNVFP4ToMXFP4MOETP1.test_gsm8k AssertionError: 0.832 not greater than 0.85 🟡 mxfp4.py is modified, but the changed lines are all k3_situ-gated; for this non-SiTU NVFP4→MXFP4 path the logic is equivalent to before. Borderline threshold (first retry 0.85 not > 0.85) → likely flake, but verify since it's the same file family.
stage-b-1gpu-small mi300 (0) test/registered/moe/test_torch_compile_moe.py TestTorchCompileMoe.setUpClass Triton HIP Code 900: op not permitted when stream is capturingHIPBLAS_STATUS_INTERNAL_ERROR → graph capture failed, server exit -9 🟡 MoE test; PR touches topk.py, but the new code runs only for _use_aiter and use_grouped_topk and correction_bias is not None. Error signature is HIP graph-capture / hipblas infra, not a dtype-cache logic error. Confirm this test uses grouped-topk+correction-bias before blaming the PR.
stage-b-2gpu-large (1) test/registered/perf/test_bench_one_batch_2gpu.py N/A (scheduler init) ImportError: cannot import name 'specialize_impl' from triton.runtime.jit; hipIpcOpenMemHandle invalid argument; scheduler died exit -6 🟢 Triton version mismatch in the CI image + IPC handle failure — pure environment, cannot be caused by a Python MoE-layout change.
stage-b-8gpu-mi35x-disagg test/registered/amd/disaggregation/test_disaggregation_basic.py test_gsm8k RuntimeError: Internal Server Error 🟡 Disagg server error; could touch the grouped-topk correction-bias path. Generic 500, no stack pointing at changed code — needs the server-side log to attribute.
wait-for-stage-b-amd-rocm720 N/A N/A fast-fail gate tripped Cascade of the four rows above; not an independent failure.

Other Executed Failures

Job Test File Test Function Error Related? Why
PR Test (NPU) multimodal-gen .../ascend/test_server_1_npu.py test_diffusion_generation[ernie/flux] perf 3026 <= 443 + JIT build code=127 + "contact self-hosted runner administrator" 🟢 NPU diffusion + runner infra; no overlap with MoE/quant.
PR Test (NPU) base-c-perf-2 .../qwen3_6_27b/test_npu_qwen3_6_27b_...50ms.py test_npu_qwen3_6_27b_2p_in64k_out1k_50ms AssertionError: 71.44 not <= 51.0 (latency) 🟢 NPU Qwen3 latency threshold; unrelated backend.
PR Test (XPU) stage-b-1gpu-xpu test/registered/attention/test_chunk_gated_delta_rule.py test_padded_state_index_is_skipped AssertionError: padded rows wrote into the state pool (2 attempts) 🟢 XPU attention (gated-delta) test; no MoE code path touched.
PR Test (Xeon) build-test test/registered/cpu/test_subblock_sparse_attention.py N/A ModuleNotFoundError: imageio + HF GatedRepoError 401 (Llama-3.1-8B) 🟢 Missing package + HF auth; pure infra.

(Cascade-only, not counted: PR Test Extra pr-gate/finish, PR Test Extra (AMD) pr-gate/finish, NPU/XPU finish jobs — all aggregators, no executed test failures of their own.)

Details / what to do before merge

  • Wait for PR Test ROCm 7.2 (AMD) to finish — 3 stage-b-1-gpu-small shards are still running; do not merge on partial AMD signal.
  • Verify the fix e2e (most important): run the Kimi-K3 SiTU A4W4 vs A8W4 GSM8K / cosine-sim comparison against a real checkpoint with AITER ≥ #4534. PR CI's unit tests do not cover this.
  • Confirm the two 🟡 AMD tests are not PR-caused: test_quark_mxfp4 is a non-SiTU path (changed lines are situ-gated) so almost certainly a borderline-threshold flake — re-run to confirm; test_torch_compile_moe and test_disaggregation_basic show infra-style graph-capture / 500 errors — check whether they reproduce on main (a quick unrelated-PR ROCm720 run had its stage-b jobs pass, so these are not universally red).
  • NPU/XPU/Xeon reds are unrelated infra (triton/imageio/HF-401, NPU perf thresholds) — not blockers for this AMD PR, but the author should confirm required vendor pipelines finish.

Generated by amd-bot using Claude Code CLI

@kkHuang-amd

Copy link
Copy Markdown
Collaborator Author

@amd-bot ci-status

@amd-bot

amd-bot commented Aug 24, 2026

Copy link
Copy Markdown

@kkHuang-amd

CI Status for PR #33838

Merge verdict: ⚠️ Not clean to merge yet, but no PR-caused failures. mergeable_state is blocked (required AMD rocm720 checks red + gate failing). All 10 executed failures across AMD/XPU/NPU/CPU are in code paths unrelated to this PR (MoE topk correction-bias cache + mxfp4 SiTU weight-layout) — they are pre-existing accuracy-threshold flakes, infra/port/connection errors, and unrelated attention/perf tests. This PR's own added unit tests ran and passed on CPU, CUDA, and AMD. The catch: PR CI does not exercise the actual runtime value of this PR (real Kimi-K3 A4W4/A8W4 MoE numerics), which needs a real K3 checkpoint + AITER #4534 the author validated manually.

Warning

This PR's added tests (test_mxfp4_situ_weight_layout.py, test_mxfp4_situ_output.py, test_topk_correction_bias_cache.py) all passed on PR CI, but they only cover the pure helper _aiter_situ_uses_gu_interleaved_weights() (env-var precedence), the zero-copy buffer-identity contract (mocked), and the correction-bias cache. The real behavioral change — routing A4W4 SiTU MoE through the generic separated layout so a real Kimi-K3 checkpoint produces correct output — is gated behind SGLANG_USE_AITER + AITER_SITUV2_A4W4/A8W4 env vars and requires the AITER #4534 revision + a real K3 checkpoint on 8×MI35x. No PR-CI job runs that path. Green CI does not verify K3 MoE correctness; rely on the author's manual GSM8K results (A4W4 200: 0.980, A8W4 200: 0.980) for that.

Caution

PR CI on pr-test-amd-rocm720 is incomplete: wait-for-stage-b-amd-rocm720 and stage-c-test-4-gpu-amd-rocm720 (mi300, 0) were cancelled (fast-fail cascade), so part of the rocm720 stage-c matrix and base-c-test-8-gpu-b300 were not tested. Author should re-run AMD CI (or use bypass-fastfail sparingly) so required pipelines finish before merge.

Changed files: python/sglang/srt/layers/moe/topk.py (+30/-1), python/sglang/srt/layers/quantization/mxfp4.py (+26/-9), 3 new registered unit tests (+221).

Executed CI failure attribution: AMD: 6 failures (0 related) · Others: 4 failures (0 related). *-finish / call-gate / pr-gate jobs are fast-fail aggregators collapsed into their root causes below.

AMD Executed Failures

Job Test File Test Function Error Related? Why
stage-b-1gpu-small-mi35x-rocm720 test/registered/quant/test_quark_mxfp4.py test_gsm8k AssertionError: 0.918 not greater than 0.92 🟢 Razor-thin GSM8K threshold flake; PR's mxfp4 change is behaviorally identical for the non-SiTU path this Quark test uses (new branch is situ-gated).
stage-c-large-8gpu-mi35x-rocm720 test/registered/amd/perf/mi35x/test_qwen35_fp8_ar_fusion_mi35x.py N/A Server exit -9; Missing index file: model.safetensors...index.json 🟢 Server crash / model-fetch infra; AR-fusion perf path, untouched by PR.
stage-c-large-8gpu-rocm720 (mi300,3) test/registered/ops/test_aiter_allreduce_fusion_amd.py N/A 'types.SimpleNamespace' object has no attribute 'moe_ep_size' 🟢 Test-harness attribute error in allreduce fusion; no topk/mxfp4 involvement.
stage-c-large-8gpu-rocm720 (mi300,0) test/registered/amd/test_moriep_small.py N/A EADDRINUSE ... port: 36573 ... address already in use 🟢 Port-collision infra flake at server bring-up.
stage-b-2gpu-large-rocm720 (mi300,1) test/registered/perf/test_bench_one_batch_2gpu.py N/A HIP error ... Capture cuda graph failed: wait timeout ... aiter_ipc 🟢 AITER IPC / CUDA-graph capture timeout; runtime infra, not MoE weight/topk.
stage-b-8gpu-mi35x-disagg-rocm720 test/registered/amd/disaggregation/test_disaggregation_basic.py N/A ConnectionRefusedError: Connect call failed 127.0.0.1:11200 🟢 Disagg server never came up; connection/infra flake.

Other Executed Failures

Job Test File Test Function Error Related? Why
stage-b-test-1-gpu-xpu test/registered/attention/test_chunk_gated_delta_rule.py test_...all_padded AssertionError: padded rows wrote into the state pool 🟢 Gated-delta attention state-pool test; unrelated subsystem.
base-c-test-perf-2-npu-a3 test/registered/npu/performance/.../test_npu_qwen3_6_27b_...50ms.py N/A AssertionError: 71.44 not less than or equal to 51.0 🟢 NPU latency-threshold perf test; unrelated to MoE code.
multimodal-gen-test-1-npu-a3 sglang/multimodal_gen/test/server/ascend/test_server_1_npu.py test_diffusion_generation latency 3649 <= 1646 fail + FAILED: [code=127] cuda_0.o 🟢 NPU diffusion perf/compile failures; unrelated to LLM MoE.
build-test (xeon-gnr, base-b-test-cpu) test/registered/cpu/test_subblock_sparse_attention.py N/A ModuleNotFoundError: No module named 'imageio' + HF 401 GatedRepoError 🟢 Missing dep / gated-repo auth in CPU env; environment issue.

Details / what to do before merge

  • Coverage (most important): the K3 MoE numerical correctness this PR fixes is not run by any PR-CI job. Before merge, confirm the author's manual validation (8×MI35x, TP8, real K3 checkpoint, AITER 6dc26b7a8 from ROCm/aiter#4534): SGLANG_USE_AITER=1 AITER_SITUV2_A4W4=1 GSM8K and the A8W4 path. Also confirm the required AITER revision is what production/CI images will ship, since the layout precedence depends on it.
  • CI completeness: re-run pr-test-amd-rocm720 so the cancelled stage-c-4-gpu matrix and base-c-test-8-gpu-b300 actually execute; don't merge on the partial rocm720 signal.
  • test_quark_mxfp4.py (🟢): 0.918 vs 0.92 is a borderline flake and the changed mxfp4 branch is situ-gated (non-SiTU path is provably unchanged), but a quick re-run will confirm it's not a real regression since it's the only failing test touching mxfp4.py.
  • Remaining failures are infra/perf/threshold/env flakes in unrelated subsystems — safe to ignore for this PR, though they should be re-run to clear the gate.

Generated by amd-bot using Claude Code CLI

@HaiShaw

HaiShaw commented Sep 3, 2026

Copy link
Copy Markdown
Collaborator

@HaiShaw HaiShaw left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Please update cookbook w.r.t. updated ENVAR @kkHuang-amd

@HaiShaw
HaiShaw merged commit a600147 into sgl-project:main Sep 3, 2026
413 of 467 checks passed
StevenChenSE pushed a commit to StevenChenSE/sglang that referenced this pull request Sep 6, 2026
Co-authored-by: wunhuang <wunhuang@amd.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants