[ROCm][Kimi-K3] Add opt-in gfx942 MXFP4-to-int4 conversion - #51274
Draft
maeehart wants to merge 7 commits into
Draft
[ROCm][Kimi-K3] Add opt-in gfx942 MXFP4-to-int4 conversion#51274maeehart wants to merge 7 commits into
maeehart wants to merge 7 commits into
Conversation
12 tasks
Contributor
Author
|
I have read the DCO document and hereby sign off past commits made by me. |
gfx942 (MI325X, MI300X) has no native MXFP4 matmul, so the day-0 Kimi-K3 expert path falls through to code that dies in LLVM codegen. Convert the MXFP4 expert weights to int4 with groupwise bf16 scales at load time and let the existing FlyDSL SiTU stage1 kernel consume them. The gfx950 native MXFP4 path is left untouched. Two supporting changes are needed for the same reason. The AITER MoE expert backend refuses kMxfp4Static outside gfx950 and does not list the SiTU activation as supported. The AITER top-k/top-p sampler is a gfx950-only prebuilt that segfaults on gfx942, so sampling falls back to the native torch implementation there. Co-authored-by: Cursor <cursoragent@cursor.com> (cherry picked from commit 2253680) Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>
The gfx942 path requantized Kimi-K3's MXFP4 experts to groupwise int4 whenever the hardware matched, without the user asking. The conversion is lossy, so it is now opt-in through --quantization-config.moe.weight int4 and gfx942 keeps the native MXFP4 path otherwise. Also refuse to load when the installed AITER predates ROCm/aiter#4471. Before that fix the packed-int4 stage1 dropped the requested activation and hardcoded SiLU, so Kimi-K3 served fluent text while computing SiLU instead of the SiTUv2 its config asks for. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>
The gfx942 Kimi-K3 path asks for groupwise int4 MoE weights through --quantization-config.moe.weight, but QUANT_KEY_NAMES had no name for that scheme, so the flag was rejected before the path could be selected. Register the existing kInt4Static32 key under int4_per_group_32 and match on the parsed QuantKey rather than a raw string. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>
mxfp4_to_f32 splits packed nibbles with repeat_interleave and then gathers through an f32 lookup table, so the working tensor is 8x the packed weight before per_1x32_i4_quant shrinks it again. Materializing that for a whole expert tensor peaks above 20 GiB per rank and fails once the weights are resident, which is what pure tensor parallel hits since it holds all experts per rank. Convert 8 experts at a time and free each slice, bounding the transient without changing the result. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>
The ROCm branch of mxfp4_round_up_hidden_size_and_intermediate_size() rounds the per-partition intermediate size up to 256. Kimi-K3 has moe_intermediate_size 3072, so a TP8 shard is 384 and gets rounded to 512. That inflates every w13 tensor from (896, 768, 1792) to (896, 1024, 1792), a 33 percent increase on all 92 MoE layers, which is about 38 GiB per rank. Pure TP8 then sits at 248.69 GiB resident with 2.12 GiB free and dies during the int4 conversion. The round-up is not needed here. AITER's resolve_flydsl_stage1_tile_n() already downgrades tile_n from 256 to 128 for a non-256-aligned inter_dim, and 128 divides 384 exactly. Mxfp4MoEMethod already skipped the round-up for the gfx950 SiTU path, so extend the same condition to the gfx942 int4 path. Measured on 8 MI325X at TP8 without expert parallel, max-model-len 16384, gpu-memory-utilization 0.97: resident drops from 248.69 GiB to 192.51 GiB, which matches the expert-parallel figure exactly, free rises from 2.12 GiB to 58.61 GiB, all 96 shards load with zero out-of-memory workers, and the KV cache is 948,305 tokens. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>
Keep the draft focused by dropping the unrelated sampler fallback. Add unit coverage for the groupwise-int4 key, explicit gfx942 selection, and native Kimi-K3 intermediate shape. Make the lossy conversion visible in startup logs and use the accelerator cache API while bounding conversion memory. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>
Current main already preserves Kimi-K3's native AITER intermediate shape. Drop the duplicate early return and its redundant test. Also remove a stale unreachable setup tail left by porting the original broader branch. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>
maeehart
force-pushed
the
k3-gfx942-int4-lossy
branch
from
August 6, 2026 15:26
410a582 to
be8095e
Compare
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
1 participant
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.
Summary
Add an explicit
int4_per_group_32MoE override for Kimi-K3 on gfx942. When selected, the MXFP4 expert weights are requantized once at load time and served through AITER's BF16 x packed-int4 FlyDSL path.The conversion is lossy and never happens by architecture alone. The default MXFP4 path is unchanged.
Why
gfx942 does not have the scaled MXFP4 MFMA used by the native CDNA4 path. The fallback can serve the model, but it dequantizes the expert weights during every forward. This opt-in path pays the conversion once during model loading and then uses the existing packed-int4 kernels.
The original implementation was part of the broader closed #50319. This draft keeps only the conversion path, its memory bounds, and the configuration needed to select it.
Changes
int4_per_group_32as an online MoE weight quantization key.--quantization-config.moe.weight int4_per_group_32override on gfx942.Validation
Unit tests:
End-to-end integration on 8x MI325X, Kimi-K3 TP8 with expert parallelism, DSpark with seven speculative tokens, prefix caching, and a 262,144-token model length:
The end-to-end image also included vLLM #50578 and #51065 plus the merged AITER large-KV offset fix. These numbers are integration evidence for the complete serving stack, not an isolated performance delta for this PR.
Dependency and scope
This remains a draft until ROCm/aiter#4471 lands in an AITER release used by vLLM.
The preferred long-term path is a numerically correct native A16W4/MXFP4 kernel such as ROCm/aiter#4502. This PR is an explicit stopgap for deployments that accept a lossy one-time conversion. It does not replace that work.
Test plan