[ROCm][Perf] Kimi-K3 latent-MoE: overlap the shared all-reduce with the routed up-projection

[mpashkovskii]

Purpose

Kimi-K3's latent MoE combines a routed expert output — produced in a compressed
latent space, then up-projected back to hidden dim — with a shared expert output.
Under tensor parallelism both need an all-reduce. Today the ROCm runner
(ROCmLatentMoERunner) only implements the column-parallel tail
(_shard_up_proj_tail, "Tier 2"), which folds a sharded up-projection into the
shared partial before a single reduce. That tail is efficient for prefill-sized
batches but gives up all overlap: the reduce must follow the accumulate.

This PR adds a token-count-dispatched overlap tail so decode-sized batches stop
leaving the shared all-reduce on the critical path:

• New Tier 1 — _overlap_allreduce_tail. All-reduce + RMSNorm the latent, then
  run the full replicated up-projection GEMM on the default stream while the
  shared-expert all-reduce runs on the aux stream (maybe_execute_in_parallel +
  aux_stream()), and add. For batches under
  VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD the shared all-reduce is hidden behind
  the up-proj GEMM.
• Tier selection — _select_tail_tier. Small batches take the overlap tail;
  larger batches keep the column-parallel tail; when the up-projection rows do not
  divide evenly across ranks (or the routed scale is non-unit) the runner falls back
  to the overlap tail, which is correct at any shape.
• _allreduce_norm_latent_out. Fused AR+RMSNorm for the latent — collapses the
  latent all-reduce and its RMSNorm into a single AITER op, with a plain all_reduce +
  RMSNorm fallback. This helper is backed by a small addition to the base runner
  (moe_runner.py): a module-level _aiter_fused_ar_rmsnorm binding (via
  rocm_aiter_ops.get_fused_allreduce_rmsnorm_op()) plus a _get_zero_residual
  buffer helper. The underlying AITER op already exists in _aiter_ops.py; the base
  runner just did not expose it before.

The native (non-fused) path is unchanged: the fused path is still gated on
tp_size > 1, an un-reduced combine output, a shared expert, no sequence parallelism,
and unit routed scale — anything else defers to the base MoERunner.forward. No
existing base-runner behavior changes; the only base-runner edit is the additive
AR+RMSNorm helper above, covered by a new unit test.

Relates to #50682 (ROCm/AMD Kimi-K3 gap and roadmap tracking → Performance
Optimization). Tracked as PR #51437.

Files changed

• vllm/models/kimi_k3/amd/latent_moe_runner.py — the overlap tier, tier selection,
  and _allreduce_norm_latent_out.
• vllm/model_executor/layers/fused_moe/runner/moe_runner.py (+90/−7) — additive
  _aiter_fused_ar_rmsnorm binding + _get_zero_residual helper (no behavior change
  to existing paths).
• tests/models/kimi_k3/test_amd_latent_moe_runner.py — tier-selection unit tests +
  multi-GPU parity for both tails.
• tests/model_executor/test_moe_runner_fused_ar_rmsnorm.py — unit test for the new
  base-runner helper.

This feature has no kernel dependency

The overlap and column-parallel tails contain no activation — they are reductions,
an up-projection GEMM, and an add. The situ/situv2 activation lives entirely in the
stage-1 expert GEMM (_forward_entry), which runs before tier selection and is
identical across Tier 1, Tier 2, and the upstream baseline. So this PR does not depend
on any activation-kernel change, and the unit + multi-GPU parity tests — which feed
synthetic routed/shared tensors straight into the tail methods and never invoke the
expert GEMM — pass regardless of the AITER build.

Environment note (not a feature dependency): serving Kimi-K3 int4 on gfx942 at all
requires ROCm/aiter#4471 ("[FlyDSL] Support SiTUv2 in the packed-int4 MoE stage1
epilogue"), because Kimi-K3 uses hidden_act = situ and the current packed-int4 FlyDSL
stage1 otherwise raises
NotImplementedError: split-K stage1 activation supports only 'silu', got 'situv2'
(or silently computes SiLU). This affects the upstream Tier-2 baseline (#51253)
exactly as much as this PR — it is a prerequisite for reproducing the end-to-end
serving numbers below, not something this feature introduces or relies on. On an AITER
without split-K SiTUv2 support, force the non-split-K stage1 with AITER_KSPLIT=1.

Relationship to #50657

#50657 ("[ROCm][Perf] Fuse and FP8-pack Kimi-K3 latent MoE output tail") targets the
same boundary but with a different mechanism: a dedicated AITER FlyDSL fused-tail
kernel plus optional FP8 packing of the up-projection weight, which depends on
unmerged AITER PRs (ROCm/aiter#4496, #4503). This PR is complementary — it schedules
existing in-tree collectives/GEMMs across two streams and adds an AITER fused AR+RMSNorm
helper to the base runner (backed by an op already present in _aiter_ops.py, so no new
kernel). The two overlap only on the runner_cls= wiring line in amd/linear.py
(already present on main); if both land, that single line needs reconciling to a
runner that can select between the two strategies.

Test Plan

Unit (no GPU) — tier-selection and fused-path gates:

.venv/bin/python -m pytest -q \
  tests/models/kimi_k3/test_amd_latent_moe_runner.py \
  -k "not tp4 and not tp8"

Multi-GPU arithmetic parity — each tail must equal the replicated up-projection:

# 4- and 8-GPU ROCm host
.venv/bin/python -m pytest -q tests/models/kimi_k3/test_amd_latent_moe_runner.py

The multi-GPU suite covers, at TP4 and TP8:

• test_overlap_tail_tp{4,8}_matches_replicated_projection — new Tier 1 output equals
  shared + up_proj(rms_norm(all_reduce(routed))).
• test_shard_tail_tp{4,8}_matches_replicated_projection — Tier 2 unchanged.
• test_shard_tail_tp4_writes_only_its_own_shard — the column-parallel tail writes
  only this rank's hidden shard before the final collective.

End-to-end serving parity (8× MI355X, gfx950, moonshotai/Kimi-K3):

vllm serve /model --served-model-name moonshotai/Kimi-K3 \
  --tensor-parallel-size 8 --trust-remote-code --moe-backend auto \
  --gpu-memory-utilization 0.9 --max-num-seqs 32 \
  --max-num-batched-tokens 4096 --reasoning-parser kimi_k3

Compare the candidate against main (baseline = Tier 2 only). Exercise both tiers by
sweeping batch size around VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD (small decode →
overlap tier; large prefill → column-parallel tier) and force the column-parallel path
with VLLM_DISABLE_SHARED_EXPERTS_STREAM=1. Note: EP must be off for the fused
tail to engage — with expert parallel enabled the MoE combine already reduces the
fused output, so the runner takes the native path and neither tail runs.

Test Result

All numbers below are placeholders — fill from the runs described in the Test Plan.
Environment: 8× MI325X (gfx942), moonshotai/Kimi-K3, int4_per_group_32, TP=8,
EP off. Serving requires an AITER that handles SiTUv2 (ROCm/aiter#4471), or
AITER_KSPLIT=1 to force the non-split-K stage1 — an environment prerequisite for K3
on gfx942, not specific to this feature.

Correctness

• Unit (-k "not tp4 and not tp8"): TODO PASS/FAIL (N passed).
• Multi-GPU parity (TP4 + TP8, overlap & shard tails): TODO PASS/FAIL.
• Accuracy parity vs main (fused tail must not change outputs within tolerance):
  TODO — e.g. gsm8k / small lm_eval, baseline vs candidate.

Decode latency / throughput — batch below
VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD so the overlap tail (Tier 1) is active.
Baseline = main (Tier 2 only); Candidate = this PR.

Batch (tokens)	Tail tier	Baseline TPOT (ms)	Candidate TPOT (ms)	Δ %	Baseline tok/s	Candidate tok/s	Δ %
1	overlap	TODO	TODO	TODO	TODO	TODO	TODO
8	overlap	TODO	TODO	TODO	TODO	TODO	TODO
32	overlap	TODO	TODO	TODO	TODO	TODO	TODO
128	overlap	TODO	TODO	TODO	TODO	TODO	TODO
256	overlap	TODO	TODO	TODO	TODO	TODO	TODO
4096 (prefill)	column-parallel	TODO	TODO	TODO	TODO	TODO	TODO

Overlap evidence — trace excerpt showing the shared-expert all_reduce running on
the aux stream concurrently with the routed up-projection mm on the default stream
(confirms the reduce is hidden, not merely reordered): TODO — paste kernel timeline.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

AI assistance was used for this change.

[mergify]

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

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