[AMD] Perf Kimi-K3 MoE optimization - #33838
Conversation
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.
Combine the sgl-project#33838 layout and cache fixes with sgl-project#33916 zero-copy support so the validated Kimi-K3 AITER path avoids redundant per-layer materialization.
Keep the Kimi-K3 SiTU zero-copy output contract while integrating the latest main changes. Co-authored-by: Cursor <cursoragent@cursor.com>
|
@amd-bot ci-status |
CI Status for PR #33838Merge 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 Caution PR CI is incomplete AND the changed code is untested. The registration validator fails during base-a/stage-a on all vendors, so Changed files: 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
Other Executed Failures
Details / what to do before mergeRoot 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:
( Coverage — must do before merge:
Generated by amd-bot using Claude Code CLI |
|
@amd-bot ci-status |
CI Status for PR #33838Merge verdict: ⛔ Not safe to merge. This PR modifies MoE routing ( Caution PR CI is incomplete AND the changed code is failing where it was exercised.
Changed files: 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
Other Executed Failures
(Cascade/gate jobs collapsed and excluded from counts: Details / what to do before merge
Generated by amd-bot using Claude Code CLI |
|
@amd-bot ci-status |
CI Status for PR #33838Merge 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 Caution This PR's core AMD code is not exercised by any completed PR-CI test. Warning PR CI is also incomplete on other vendors: the NVIDIA main run 32319934749 was cancelled at ~05:17 mid Changed files: 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
Coverage detail (what ran vs. what didn't)
Details / what to do before merge
Generated by amd-bot using Claude Code CLI |
|
@amd-bot ci-status |
CI Status for PR #33838 — [AMD] Perf Kimi-K3 MoE optimizationMerge verdict: Not ready to merge. Warning Two coverage gaps, both matter here:
Changed files: 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/ AMD Executed Failures
Other Executed Failures
(Cascade-only, not counted: Details / what to do before merge
Generated by amd-bot using Claude Code CLI |
|
@amd-bot ci-status |
CI Status for PR #33838Merge verdict: Warning This PR's added tests ( Caution PR CI on Changed files: Executed CI failure attribution: AMD: 6 failures (0 related) · Others: 4 failures (0 related). AMD Executed Failures
Other Executed Failures
Details / what to do before merge
Generated by amd-bot using Claude Code CLI |
Combine the sgl-project#33838 layout and cache fixes with sgl-project#33916 zero-copy support so the validated Kimi-K3 AITER path avoids redundant per-layer materialization.
HaiShaw
left a comment
There was a problem hiding this comment.
LGTM. Please update cookbook w.r.t. updated ENVAR @kkHuang-amd
Co-authored-by: wunhuang <wunhuang@amd.com>
Motivation
Kimi-K3 SiTUv2 MXFP4 MoE requires different preshuffled weight layouts for different activation quantization modes:
shuffle_weight_a16w4layout.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
Accuracy Tests
Environment: 8x MI35x/gfx950, TP8, public Kimi-K3 checkpoint.
11 passed.4.0/ linear beta25.0: all 17 K3 cases passed.1.000.0.980(before fix:0.005).0.980.Speed Tests and Profiling
Required AITER revision:
6dc26b7a8from ROCm/aiter#4534.Server command:
Client workload:
random-range-ratio=1.0.8192 / 1024.num-prompts = 8 * concurrency.warmup-requests = 2 * concurrency.Final results with the MoE copy optimization, compared with the original MI35x reference:
At concurrency 32 versus the original reference, total token throughput improves by
2.68%, median E2EL by2.55%, TTFT by3.57%, TPOT by2.54%, and ITL by2.29%.The MoE copy change was also measured against the immediately preceding sweep with identical server/client settings:
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, versus537.33 tok/sfor post-fix A8W4.Checklist
pre-commit run --all-filespasses).Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ci.CI States
Latest PR Test (Base): ❌ Run #32319934749
Latest PR Test (Extra): ❌ Run #32319934659
Latest PR Test (AMD ROCm 7.2): ❌ Run #32319934688