Repository navigation
[Perf][ROCm][Kimi-K3] Fuse all-reduce with RMSNorm in latent-MoE - #57920
amd-mghanimi wants to merge 11 commits into
Conversation
Call AITER's 1-stage custom AR+RMSNorm from the latent MoE runner instead of a plain all-reduce plus a separate RMSNorm, matching the NVIDIA FlashInfer site. Prefill-sized tensors still fall back because they miss the 1-stage gate. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Mehdi Ghanimifard <mehdi.ghanimifard@amd.com>
Kimi-K3 runs under breakable CUDA graphs, not a torch.compile eager/compile split, so the comments and logs should name that execution path. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Mehdi Ghanimifard <mehdi.ghanimifard@amd.com>
ROCm uplift PRs must not edit vllm/models/kimi_k3/nvidia/. The common fused_allreduce_rms_norm_out helper remains for the AMD path. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Mehdi Ghanimifard <mehdi.ghanimifard@amd.com>
T1 fusion targets CompilationMode.NONE, not torch.compile: no compile decorator, no eager_break_during_capture, zeros_like instead of a process-global residual cache so capture can record the launch. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Mehdi Ghanimifard <mehdi.ghanimifard@amd.com>
Revert vllm/models/common/ops/fused_allreduce_rms_norm.py: it is the flashinfer helper consumed by the NVIDIA K3/DSV32 models and has no ROCm path to add. The AMD tail now calls torch.ops.vllm.rocm_aiter_fused_allreduce_rmsnorm directly, matching how amd/mla.py fires fused_mla_dual_rms_norm. Gate on AITER custom AR being the active backend, on aiter exposing a runtime hidden_dim (K3's latent is 3584, outside the old template set), and on the one-stage kernel bound; everything else keeps the unfused all-reduce. Zero residuals are cached per shape and never resized, so a graph captured against one size cannot replay against a reallocated buffer. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Mehdi Ghanimifard <mehdi.ghanimifard@amd.com>
AITER's gfx1250 custom-AR backend has no fused_allreduce_rmsnorm, so the K3 tail falls back to plain AR+RMS there. The test is one TP2 spawn. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Mehdi Ghanimifard <mehdi.ghanimifard@amd.com>
amd-mghanimi
changed the title
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
@claude review |
| return buf | ||
|
|
||
|
|
||
| def can_fuse_allreduce_rms_norm(hidden_states: torch.Tensor) -> bool: |
There was a problem hiding this comment.
Suggestion: Can we re-use _can_use_aiter_fused_ar_rms instead (and add the on_gfx1250 gate there👍)
There was a problem hiding this comment.
this is also a valid point, if fused_ar_rms needs to live in vllm/model_executor/layers it makes sense to reuse as much as possible between different variations of ar+norm such as vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py
If there is a clear consensus on this I will update the code accordingly.
There was a problem hiding this comment.
I believe we are talking about this:
fused_allreduce_gemma_rms_norm.py
If I understand correctly, the issue is that in the current design we want to keep the code cleaner by keeping flashinfer/trtllm and aiter separate and independent from each other. but that implementation still keeps them closely related. Perhaps we need a separate PR for this.
| from vllm.model_executor.layers.layernorm import RMSNorm | ||
|
|
||
| # Zero residual, cached per shape so CUDA-graph captures keep a stable pointer. | ||
| _zero_residuals: dict[tuple[torch.Size, torch.dtype, torch.device], torch.Tensor] = {} |
There was a problem hiding this comment.
Question: Would the WorkspaceManager be safer here instead of manual allocation and caching?
There was a problem hiding this comment.
I think the term zero_residual is related to difference between RMSNorm and AddRMSNorm, so when there is nothing to add, we still need a dummy zero tensor. it takes a few kilo bytes of memory.
The old implementation has this pattern:
vllm/vllm/compilation/passes/fusion/allreduce_rms_fusion.py
Lines 1184 to 1191 in 606d124
My understanding is that this way we make sure that the dummy cache is allocated once for all other calls since they are stored like key/value dict. I guess we could also use get_persistent/get_persistent_resource then we have to make dependency on v1 workers? since we don't want to change this tensor maybe it is an overkill?
The helper is not Kimi-K3-specific, so keep it next to other reusable AMD ops. Signed-off-by: Mehdi Ghanimifard <mehdi.ghanimifard@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This branch has not been deployed
2 participants
Purpose
When running KimiK3 with aiter backend, there is an opportunity to fuse all-reduce and RMSnorm using the existing fusion kernels:
torch.ops.vllm.rocm_aiter_fused_allreduce_rmsnormThere is an explicit usage of cross-device all reduce followed by rmsnorm in Latent-MoE tail here:
vllm/vllm/models/kimi_k3/amd/latent_moe_runner.py
Lines 76 to 78 in b8cf275
To be able to enable this, we needed to check some guards and implement the call interface for
rocm_aiter_fused_allreduce_rmsnorm. So, without updating aiter this PR simply defines:vllm/models/common/amd/ops/fused_ar_rms.py, with a unittest to make sure the guards are in place for other configs.Fusion is gated with
can_fuse_allreduce_rms_normchecks:VLLM_ROCM_USE_AITER_CUSTOM_AR) (default is on)fused_allreduce_rmsnorm;looking at latest AITER v0.1.22 gfx1250 header vs fused op);
fused_allreduce_rmsnormdoesn't exist on gfx1250 yet.supports_dynamic_hidden_dim(K3 latent is 3584; pre-0.1.12 aitertemplate-specializes on
{512, 1024, 2048, 4096}and silently no-ops)should_custom_arand one-stage fused AR+RMS (two-stage is slower thanplain AR + norm)
Additional notes
On TP8 bf16 latent
[tokens, 3584], one-stage is< 128 KB(~≤18 tokens).CudaCommunicator.all_reducetries QuickReduce before AITER custom AR;QuickReduce INT4 needs ≳292 tokens. So, this fusion practically only hits Decode and only on concurrency ~≤18 tokens.
Larger capture sizes keep unfused all-reduce.
Confirm against the
Using [...] all-reduce backendsstartup log.Zero residuals are cached per
(shape, dtype, device)and never resized, so agraph captured at one size cannot replay against a reallocated buffer.
AI-assisted: Cursor was used to draft and iterate on this change. I reviewed
the code and own the submission.
Test Plan
tests/models/kimi_k3/test_amd_fused_ar_rms.pyon ROCm (2 GPUs): fused(4 tokens) and fallback (256 tokens) vs explicit AR + RMSNorm
PR at MC4 for ISL/OSL 8k/1k and 60k/600. Expect the fused
1-stage kernel to replace one
cross_device_reduce_1stageand oneadd_rmsnorm_quantper launch.concurrency. Expect ITL movement only where one-stage fusion can fire
(decode ≲18 tokens, e.g. MC4/MC8); higher MC is a control.
Serve config
We only ran the tests on MI355X. This was tested by patching this nightly image onto
vllm/vllm-openai-rocm:nightly-rocm100-af1c01499b289be555c475669ba50a88e96d846eas
nightly-af1c014-k3-fused-ar-rms:The profile traces was captured with these additional configs:
(so the cropped steps was used for the call statisitic.)
Test Result
Unit test:
tests/models/kimi_k3/test_amd_fused_ar_rms.pycovers fused (4tokens) and unfused fallback (256 tokens).
Profiler, rank-0 call counts, same decode capture. Each fused launch consumed
one all-reduce and one norm (Δs agree to the unit). The fused kernel is absent
from baseline.
8k/1k MC4
aiter::cross_device_reduce_1stageaiter::allreduce_fusion_kernel_1stageaiter::add_rmsnorm_quant_kernel60k/600 MC4
aiter::cross_device_reduce_1stageaiter::allreduce_fusion_kernel_1stageaiter::add_rmsnorm_quant_kernelThe perfetto screenshot shows how mc4 looks like with/without (from 30us down to 33us):
no fusion:
all reduce fusion:
When serving, median ITL (ms) and output tok/s recieve the uplift on affected MC.
One-stage fusion only covers decode batches ≲18 tokens (MC4/MC8 here). Higher MC does not hit by this PR.
Accuracy Test Results
I ran
lm_eval's needle in a haystacknia_single2with two different batch sizesnum_concurrent{8 and 16}:Two numbers per cell for C=8/C=16
Between running the experiments and submitting this PR, there was an aiter bump. So, for accuracy tests I switched to the latest image:
vllm-openai-rocm/nightly-rocm100-3df4ae153eb385e27b52f26c81f8edb9e20b9984However, there was an OOM errors in the base image, so I patched both runs with #57923 as a temporary solution to examile accuracy here.
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.