Repository navigation
[ROCm][Kimi-K3][Perf] Cut gfx942 MXFP4-to-int4 requant time under pure TP - #64004
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
simondanielsson
left a comment
simondanielsson
left a comment
There was a problem hiding this comment.
Thanks for the fix!
| # One expert per chunk gives NUM_EXPERTS chunks of about ten temporaries | ||
| # each. Later chunks must reuse the cached blocks, so device allocations | ||
| # stay a small constant instead of growing with the chunk count. | ||
| assert device_allocs <= 16 |
There was a problem hiding this comment.
Question: Whats the rationale behind 16 here?
There was a problem hiding this comment.
There wasn't a strong one: the fix made 5 device allocations and the old code 50, and 16 sat in between. It also ran without load_model's max_split_size_mb:20, so it missed the tail-chunk issue okorzh-amd found.
I replaced it with the exact property in 0758e29: on the Kimi-K3 TP8 w2 shape under max_split_size_mb:20, no device memory is allocated after the first chunk. The test records num_device_alloc at the start of each chunk by wrapping fp4_utils.mxfp4_to_f32, which still runs. It fails on the refactor-only parent (every chunk allocates) and on the previous revision (the last chunk allocates), and passes now. The outputs are in the PR description.
simondanielsson
added
the
ready
|
✅ @xaguilar-amd, CI is now available for this PR.
|
| w_all = weight.data.view(fp4_dtype) | ||
| s_all = scale.data.view(e8m0_dtype) | ||
| num_experts, rows, packed_cols = w_all.shape | ||
| chunk = max(1, _K3_INT4_CONVERT_CHUNK_VALUES // (rows * packed_cols * 2)) |
There was a problem hiding this comment.
The per-layer table and the Tradeoffs paragraph don't match the shapes this path gets. Kimi-K3's routed experts use routed_expert_hidden_size=3584, not hidden_size 7168. Could you rerun the microbenchmark at 3584, and say whether the before and after runs used the same base?
There was a problem hiding this comment.
You're right, thanks for spotting this! My microbenchmark hardcoded hidden_size 7168, so the per-layer table, the chunk counts and the Tradeoffs paragraph were wrong. The end-to-end numbers came from loading the real checkpoint though, so they were not affected.
I reran the microbenchmark at 3584 and updated the description. Before and after use the same base: the end-to-end "before" runs and the old microbenchmark used 74340c7 with only mxfp4.py swapped back, and mxfp4.py is identical between 74340c7 and the current base 3f45273.
At the real shapes, the previous revision's peak reserved memory was about 11.5 GiB per layer, above the old EP8 path's 5.7 GiB, so the "stays below the old EP8 peak" claim was false. Most of that came from the tail-chunk issue in your other comment. With that fixed and the budget halved to 2^27, the peak is 2.89 GiB allocated and 4.56 GiB reserved for both TP8 and EP8, below the old EP8 path (3.43 / 5.66 GiB). The conversion is also faster than at 2^28 (TP8 0.13 s vs 0.27 s per layer, EP8 0.30 s vs 0.55 s).
Thanks for the great feedback!
| assert out_scale is not None | ||
| out_packed[lo:hi].copy_(chunk_packed) | ||
| out_scale[lo * scale_stride : hi * scale_stride].copy_(chunk_scale) | ||
| # No empty_cache() here: the next chunk reuses these cached |
There was a problem hiding this comment.
Every Kimi-K3 weight ends with a shorter chunk: TP8 w13 is 896 = 9x97+23 and EP8 w13 is 112 = 9x12+4. load_model runs this conversion under max_split_size_mb:20. In that mode, PyTorch reuses a cached block for a request of 20 MiB or more only when the block is less than 20 MiB larger than the request. So most of the tail chunk's temporaries miss the cache.
There was a problem hiding this comment.
Yes, confirmed! At production shapes the first chunk made 20 device allocations, middle chunks 0, and the last, shorter chunk 5 to 21 more, held next to the cached full-size set.
Fixed now by shifting the last chunk back so it overlaps the previous one, so every chunk has the same shapes. The overlapping experts are recomputed to identical values, which costs less than one chunk per weight. Now no chunk after the first allocates device memory, and the new test checks that.
Alternatives I measured: calling empty_cache() before the short chunk fixes the peak but keeps an extra round of allocations and frees; balancing chunk sizes still leaves sizes one expert apart (about 88 MB of f32 under EP8), which still misses; preallocating the outputs before the loop added an allocation per chunk.
End to end with this revision: TP8 requant 25 s (was 59 s), startup 210 s (was 235 s); EP8 requant 24 s (was 63 s), startup 181 s (was 216 s). GSM8K 5-shot TP8 0.988 (250 questions), EP8 1.000 (100).
Thank you!
|
/ci run |
|
✅ Triggered Buildkite CI #97968 for commit |
|
/ci run |
|
✅ Triggered Buildkite CI #98073 for commit ℹ️ This PR is 16 commits behind upstream |
|
This pull request has merge conflicts that must be resolved before it can be |
…a helper Move the nested convert() out of Mxfp4MoEMethod._setup_kernel_k3_situ_gfx942 so it can be tested on its own. No behavior change. Signed-off-by: Xavier Aguilar <xavier.aguilarfruto@amd.com>
The gfx942 int4 requant converted 8 experts per chunk and called empty_cache() after every chunk. Pure TP keeps all 896 experts per rank at 1/tp width, so it ran 8x more chunks than EP for the same bytes, and each chunk released its temporaries with hipFree. On an 8x MI325X Kimi-K3 load the requant took about 19 minutes under TP8 and 4.6 minutes under EP8. Size chunks by unpacked values instead of experts, so TP and EP run the same number of chunks, and keep the temporaries cached across chunks. The cache is still emptied once per converted weight. The output is byte-identical for any chunk size. Signed-off-by: Xavier Aguilar <xavier.aguilarfruto@amd.com>
Under load_model's max_split_size_mb:20, a cached block of 20 MiB or more is only reused for a request less than 20 MiB smaller than it. The last chunk of every Kimi-K3 weight was shorter, so its temporaries missed the cache and were allocated next to the cached full-size set, raising the per-layer peak reserved memory to about 11.5 GiB. Shift the last chunk back so it overlaps the previous one and every chunk has the same shapes. Halve the chunk budget to 2^27 values. On the real per-rank shapes (routed expert hidden size 3584) this keeps the peak below the old EP8 path and is faster under both TP8 and EP8. Replace the device allocation threshold in the test with the exact property: no device allocations after the first chunk, checked on the Kimi-K3 TP8 w2 shape under load_model's allocator settings. Skip the test on a missing aiter package rather than one submodule. Co-authored-by: Claude Signed-off-by: Xavier Aguilar <xavier.aguilarfruto@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Replace the chunk-size self-comparison with a comparison against an unchunked reference that decodes MXFP4 as the OCP MX spec defines it and applies aiter's int4 quant and shuffles to all experts at once. The self-comparison passed for bugs that are the same in every chunk, such as a dropped block scale or a swapped nibble order; the reference test fails for both. Co-authored-by: Claude Signed-off-by: Xavier Aguilar <xavier.aguilarfruto@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
0eea9b7 to
53bb4f1
Compare
|
Rebased onto main. This picked up #62407, which keeps loaded params' storage in |
|
/ci run |
|
✅ Triggered Buildkite CI #98236 for commit |
fxmarty-amd
left a comment
fxmarty-amd
left a comment
There was a problem hiding this comment.
Nice work, a followup should migrate this as part of https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/online
|
fyi: vllm/vllm/model_executor/model_loader/base_loader.py Lines 158 to 167 in 255431b |
|
thanks @fxmarty-amd |
|
/ci run |
|
✅ Triggered Buildkite CI #98676 for commit |
Overview
Cut the gfx942 Kimi-K3 MXFP4-to-int4 MoE requant (
--quantization-config '{"moe":{"weight":"int4_per_group_32"}}', added in #51274) from ~19 min to 25 s under pure TP8. TP8 server startup drops from 22.5 min to 3.5 min. The converted weights are byte-identical.Validation
End to end: 8x MI325X (gfx942), ROCm 10.0, Kimi-K3 with the DSpark draft. The checkpoint was in the page cache for all startup numbers. Before and after use the same base; only
mxfp4.pydiffers./healthbefore/healthafterThe requant is the time between
Loading weights tookandOnline quantization as part of model loadingon rank 0.Accuracy (GSM8K, 5-shot,
lm_eval --model local-completions, with the fix):The weights are byte-identical to the current code (see below), so I didn't rerun accuracy on the baseline.
Per-layer microbenchmark: one full Kimi-K3 MoE layer (w13 + w2), per-rank shapes with
routed_expert_hidden_size3584, run undermax_split_size_mb:20likegpu_worker.load_model. Times are the median of 7 runs, alternating before and after. The old code's first run in a fresh process is faster (1.94 s TP8, 0.85 s EP8); itshipFreecost grows as the process maps more memory.empty_cache()calls beforeNew test:
tests/kernels/quantization/test_k3_gfx942_int4_requant.py. It runs on gfx942 only and is picked up by the existing "(MI300) Quantization Kernels" step. It checks that the output does not depend on chunking, and that no device memory is allocated after the first chunk on the Kimi-K3 TP8 w2 shape underload_model's allocator settings. It passes 4/4 with the fix and fails before it. Against the refactor-only parent (9acd4b5), every chunk allocates:Against the first revision of this PR (3bdfe80), the shorter last chunk allocates:
For those runs the test loaded
mxfp4.pyfrom each commit. On the parent I appended_K3_INT4_CONVERT_CHUNK_VALUESonly so the test'smonkeypatch.setattrresolves; the old code ignores it.tests/models/kimi_k3/test_gfx942_int4.pyandtests/quantization/test_quantization_config_args.pystill pass.Details
Root cause.
_setup_kernel_k3_situ_gfx942converted 8 experts per chunk and calledtorch.accelerator.empty_cache()after every chunk. Pure TP keeps all 896 experts per rank at 1/tp width, so TP8 ran 112 chunks per weight where EP8 runs 14, for the same bytes. That is 20,608 chunks over 92 MoE layers.Every
empty_cache()hands the chunk's temporaries back to the driver, and the next chunk allocates them again. In a live TP8 server, a native py-spy profile of the requant showed 92% of the time inemptyCache -> release_block -> hipFree(HSAmunmap/ioctl) and about 5% inhipMalloc. The conversion kernels themselves barely appeared. That worked out to about 56 ms per chunk.Fix.
_K3_INT4_CONVERT_CHUNK_VALUES = 1 << 27. TP8 and EP8 both run 19 chunks per w13 and 10 per w2.empty_cache(), so later chunks reuse the cached temporaries. The cache is still emptied once per converted weight.max_split_size_mb:20, a cached block of 20 MiB or more is only reused for a request less than 20 MiB smaller, so a shorter last chunk would allocate a second set of temporaries. The overlapping experts are recomputed to identical values, which costs less than one chunk per weight.per_1x32_i4_quant,shuffle_weightandshuffle_scale_for_int4all work per expert along the leading dimension. The new test checks this explicitly.Tradeoffs. The TP8 per-layer peak grows (allocated 1.56 to 2.89 GiB, reserved 1.85 to 4.56 GiB), because the old TP chunks were tiny. The EP8 peak shrinks (allocated 3.43 to 2.89 GiB, reserved 5.66 to 4.56 GiB). Both stay below the old EP8 peak, which already fit.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.