[ROCm][Kimi-K3] Add opt-in gfx942 MXFP4-to-int4 conversion by maeehart · Pull Request #51274 · vllm-project/vllm

maeehart · 2026-08-06T15:09:53Z

Summary

Add an explicit int4_per_group_32 MoE override for Kimi-K3 on gfx942. When selected, the MXFP4 expert weights are requantized once at load time and served through AITER's replacement 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 support the scaled FP4 conversion instructions used by the native MXFP4 path. This opt-in path pays the conversion cost once during model loading and then uses the replacement packed-int4 kernels from ROCm/aiter#4646.

Changes

Register int4_per_group_32 as a MoE weight quantization key.
Require the explicit --quantization-config.moe.weight int4_per_group_32 override on gfx942.
Require the replacement A16W4/int4 AITER entry point and fail closed unless it supports SiTUv2.
Convert eight experts at a time and release each source tensor before converting the next one, bounding transient memory during load.
Log a warning that the selected conversion is lossy.

Validation

Unit tests:

tests/quantization/test_quantization_config_args.py: 15 passed
tests/models/kimi_k3/test_gfx942_int4.py: 5 passed

Strict AITER comparison against the torch reference passed for the production EP shape:

model_dim=3584
intermediate_dim=3072
local_experts=112
topk=15
tokens=1,4,7,16,32
activation=SiTUv2
beta=4.0
linear_beta=25.0

Re-checked on 8x MI325X with TP8, expert parallelism, Triton RMSNorm, Triton MLA, AITER fused MoE, and full plus piecewise graph capture, using the replacement packed-int4 path from ROCm/aiter#4646 (AITER v0.1.21.post2, picked up in #55968).

Graph capture completed.
Chat completions returned coherent answers, including a Fibonacci Python snippet.
Five-shot GSM8K through lm_eval local-completions (num_fewshot=5, temperature=0.0, max_gen_toks=512) scored 0.9598 strict-match and 0.9598 flexible-extract on all 1,319 questions, stderr 0.0054. That is within noise of the 0.9560 A16W4 baseline (stderr 0.0056) on the same 1,319 questions.
vllm bench serve 1K input / 128 output completed 10/10 requests at concurrency 1 and 160/160 requests at concurrency 16, with no server errors. These serving runs used the untuned FlyDSL a16w4 fallback, so they are a functional check, not a tuned peak.

Long-context and DSpark validation against this AITER release remains pending.

Dependency and scope

ROCm/aiter#4646 is in AITER v0.1.21.post2. vLLM picked that tag up in #55968, so this branch is no longer blocked on the AITER release.

The conversion is an explicit opt-in for deployments that accept a lossy one-time weight conversion. It does not change the default checkpoint path.

Test plan

Quantization key and opt-in selection unit tests.
Strict production-shape comparison against the torch reference.
Full 1,319-question GSM8K after FlyDSL: port a16wi4 to new pipeline, clean old moe_gemm_2stage ROCm/aiter#4646.
Target-only graph capture and serving sanity.
Long-context and DSpark validation against the AITER release containing FlyDSL: port a16wi4 to new pipeline, clean old moe_gemm_2stage ROCm/aiter#4646.
ROCm CI on the rebased branch.

maeehart · 2026-08-06T15:11:06Z

I have read the DCO document and hereby sign off past commits made by me.

maeehart · 2026-08-12T20:42:42Z

I have read the DCO document and hereby sign off on the commits in this pull request.

mergify · 2026-08-14T21:24:18Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @maeehart.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

mergify · 2026-09-02T13:25:51Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @maeehart.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

xaguilar-amd · 2026-09-04T07:22:31Z

@maeehart AITER has been bumped to 0.1.21.post1 #52826 we can push forward this one now.

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

fxmarty-amd · 2026-09-09T16:59:32Z

Hi @maeehart, do you think this could use #55684 ?

This would involve:

Writing Mxfp4MoEMethod.dequantize_weight
Extending [Quantization] Layer re-quantization for linear layers through online quantization API (MXFP8 -> FP8 PTPC showcase) #55684 to MOE methods.
Adding int4 online MOE method to: https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/online (similar to Int8OnlineMoEMethod)

simondanielsson

Very nice work!

mergify · 2026-09-10T15:31:39Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @maeehart.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

maeehart · 2026-09-10T17:40:50Z

@fxmarty-amd I looked at routing this through #55684. I would keep this PR on the current Mxfp4MoEMethod load-time conversion.

#55684 is still open and linear-only. resolve_quant_method still rejects online requant of a pre-quantized MoE layer. A full Mxfp4MoEMethod.dequantize_weight to BF16 also does not fit. The MXFP4 unpack is about 8x the packed weight and peaks above 20 GiB per rank, which is why this PR converts eight experts at a time.

The decode path we measured is AITER FlyDSL a16w4 with SiTUv2 through AiterExperts. Int8OnlineMoEMethod selects Triton/Marlin/CPU. fused_flydsl_moe hardcodes SiLU. Either would drop the activation this model needs.

Five-shot GSM8K on this path is 0.9598 strict-match on 1,319 questions, stderr 0.0054.

MoE requant through the online API can follow once #55684 covers MoE. I do not want to block this opt-in gfx942 path on that.

maeehart · 2026-09-10T17:40:52Z

@dllehr-amd Can you review this for merge if you agree with keeping the current Mxfp4MoEMethod load-time path?

Five-shot GSM8K held at 0.9598 strict-match on 1,319 questions, stderr 0.0054. I replied to @fxmarty-amd on the online-requant thread and to @simondanielsson on the inline comments. The leftover duplicate SiTU entry is removed in 6a0a733.

Mergify currently reports merge conflicts with main, so a rebase is still needed before merge.

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> Signed-off-by: Hartikainen <mahartik@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> Signed-off-by: Hartikainen <mahartik@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> Signed-off-by: Hartikainen <mahartik@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> Signed-off-by: Hartikainen <mahartik@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> Signed-off-by: Hartikainen <mahartik@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> Signed-off-by: Hartikainen <mahartik@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> Signed-off-by: Hartikainen <mahartik@amd.com>

Switch the capability check to AITER's replacement A16W4/int4 entry point because the original packed-int4 implementation is being removed. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>

Drop the unused layer argument from the gfx942 int4 kernel factory call so the branch passes static type checking. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Hartikainen <mahartik@amd.com>

The rebase conflict resolution used backend_to_kernel_cls and os without importing them, which made mypy fail in pre-commit. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Signed-off-by: Hartikainen <mahartik@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

fxmarty-amd · 2026-09-10T18:07:45Z

MoE requant through the online API can follow once #55684 covers MoE. I do not want to block this opt-in gfx942 path on that.

yes, makes sense, we are still a bit far from it, I agree.

The decode path we measured is AITER FlyDSL a16w4 with SiTUv2 through AiterExperts. Int8OnlineMoEMethod selects Triton/Marlin/CPU. fused_flydsl_moe hardcodes SiLU. Either would drop the activation this model needs.

Well, this is the reason why oracles / MOE backends and abstractions as

vllm/vllm/model_executor/layers/fused_moe/modular_kernel.py

Lines 628 to 634 in 8359e15

    
               @staticmethod 
        
               @abstractmethod 
        
               def _supports_activation(activation: MoEActivation) -> bool: 
        
                   """ 
        
                   Whether the kernel supports a particular act function. 
        
                   """ 
        
                   raise NotImplementedError

were introduced.

IMO ideally, this should make use of https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py that is supposed to handle MOE backends for kInt4Static32 weights. I am not sure it makes sense to have this as part of MXFP4 oracle, that is normally meant to handle backends making use of MXFP4 weights?

The MOE oracle refactor was meant to separate the:

weight creation,
backend selection,
weight processing to be mapped to what backends except

and having the re-quantization step, custom backend selection & processing all under mxfp4.py here sounds like we are not adopting these oracle/backend abstractions.

I don't think this is blocking, just nit comment!

SiTU was already in _supports_activation. The second entry was a rebase leftover. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Hartikainen <mahartik@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

maeehart · 2026-09-10T18:37:48Z

I have read the DCO document and hereby sign off on the commits in this pull request.

maeehart · 2026-09-10T18:39:01Z

I have read the DCO document and hereby sign off on the commits in this pull request.

maeehart · 2026-09-10T18:39:02Z

@fxmarty-amd Agreed. I would keep this PR as the load-time special case.

int_wna16.py currently selects Marlin, Triton, FlashInfer TRTLLM, Humming, CPU, XPU, and emulation. It has no AITER backend. The measured path is AiterExperts FlyDSL a16w4 with SiTUv2. Routing the converted weights through that oracle would drop the activation.

The checkpoint is still MXFP4. The int4 tensors exist only after _setup_kernel_k3_situ_gfx942. That is why the convert and kernel pick still live in Mxfp4MoEMethod for this PR.

A follow-up can add AITER packed-int4 to the WNA16 oracle, or add Mxfp4MoeBackend.AITER_INT4_BF16 as in the inline comment. I do not want to do that rewrite in this PR.

maeehart · 2026-09-11T14:50:51Z

Can you get this reviewed @dllehr-amd. We have discussed this PR also separately and the design concerns are now resolved.

maeehart mentioned this pull request Aug 6, 2026

[ROCm][AMD] Kimi-K3 Gap and Roadmap Tracking #50682

Open

13 tasks

mergify Bot added quantization kimi k3 rocm Related to AMD ROCm labels Aug 6, 2026

github-project-automation Bot added this to AMD Aug 6, 2026

github-project-automation Bot moved this to Todo in AMD Aug 6, 2026

maeehart force-pushed the k3-gfx942-int4-lossy branch from 410a582 to be8095e Compare August 6, 2026 15:26

maeehart mentioned this pull request Aug 12, 2026

[FlyDSL] Support SiTUv2 in the packed-int4 MoE stage1 epilogue ROCm/aiter#4471

Closed

1 task

omirosh mentioned this pull request Aug 12, 2026

[ROCm][Bugfix] Kimi-K3: opt-in gfx942 MXFP4-to-int4 MoE conversion sgl-project/sglang#34566

Draft

5 tasks

fxmarty-amd mentioned this pull request Aug 13, 2026

[RFC]: Extended online quantization roadmap #52167

Open

1 task

mergify Bot added the needs-rebase label Aug 14, 2026

maeehart mentioned this pull request Aug 18, 2026

[ROCm][AMD] Kimi-K3 gfx942 / MI325X Gap and Roadmap #52803

Open

maeehart force-pushed the k3-gfx942-int4-lossy branch from a694c54 to 0b152a8 Compare August 24, 2026 15:36

mergify Bot removed the needs-rebase label Aug 24, 2026

maeehart force-pushed the k3-gfx942-int4-lossy branch from 1575b38 to 10e3a00 Compare August 24, 2026 20:32

mergify Bot added the needs-rebase label Sep 2, 2026

fxmarty-amd mentioned this pull request Sep 7, 2026

[Quantization] Layer re-quantization for linear layers through online quantization API (MXFP8 -> FP8 PTPC showcase) #55684

Open

maeehart force-pushed the k3-gfx942-int4-lossy branch from 10e3a00 to 2110215 Compare September 9, 2026 16:19

maeehart marked this pull request as ready for review September 9, 2026 16:23

maeehart requested review from AndreasKaratzas, mgoin, robertgshaw2-redhat, tjtanaa and yewentao256 as code owners September 9, 2026 16:23

maeehart requested review from ProExpertProg, houseroad and tlrmchlsmth as code owners September 9, 2026 16:23

claude Bot reviewed Sep 9, 2026

View reviewed changes

mergify Bot removed the needs-rebase label Sep 9, 2026

simondanielsson reviewed Sep 9, 2026

View reviewed changes

Comment thread vllm/model_executor/layers/quantization/mxfp4.py

Comment thread vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py Outdated

Comment thread vllm/model_executor/layers/quantization/mxfp4.py

Comment thread vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

mergify Bot added the needs-rebase label Sep 10, 2026

maeehart and others added 10 commits September 10, 2026 10:57

maeehart force-pushed the k3-gfx942-int4-lossy branch from 6a0a733 to 0fc3721 Compare September 10, 2026 18:10

maeehart force-pushed the k3-gfx942-int4-lossy branch from 0fc3721 to 97eb956 Compare September 10, 2026 19:12

mergify Bot removed the needs-rebase label Sep 10, 2026

Uh oh!

Uh oh!

Conversation

maeehart commented Aug 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Why

Changes

Validation

Dependency and scope

Test plan

Uh oh!

maeehart commented Aug 6, 2026

Uh oh!

maeehart commented Aug 12, 2026

Uh oh!

mergify Bot commented Aug 14, 2026

Uh oh!

mergify Bot commented Sep 2, 2026

Uh oh!

xaguilar-amd commented Sep 4, 2026

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

fxmarty-amd commented Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

simondanielsson left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mergify Bot commented Sep 10, 2026

Uh oh!

maeehart commented Sep 10, 2026

Uh oh!

maeehart commented Sep 10, 2026

Uh oh!

fxmarty-amd commented Sep 10, 2026

Uh oh!

maeehart commented Sep 10, 2026

Uh oh!

maeehart commented Sep 10, 2026

Uh oh!

maeehart commented Sep 10, 2026

Uh oh!

maeehart commented Sep 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

maeehart commented Aug 6, 2026 •

edited

Loading

fxmarty-amd commented Sep 9, 2026 •

edited

Loading