[ROCm][Perf] Kimi-K3 Enable sharded latent MoE up-projection under EP (#54956)
Signed-off-by: Xavier Aguilar <xavier.aguilarfruto@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
xaguilar-amd committed
Sep 28, 2026 at 16:31 UTC
a7ff44355c1cadd1d85b82f4adc9694ce7864b56
2 files changed
+97
-16
tests/models/kimi_k3/test_amd_latent_moe_runner.py
+95
-15
@@ -23,6 +23,7 @@ from vllm.distributed import get_tp_group
23
from vllm.model_executor.layers.fused_moe.runner import moe_runner
24
from vllm.model_executor.layers.layernorm import RMSNorm
25
from vllm.model_executor.layers.linear import ReplicatedLinear
26
+from vllm.models.kimi_k3.amd import latent_moe_runner
27
from vllm.models.kimi_k3.amd.latent_moe_runner import ROCmLatentMoERunner
28
from vllm.models.kimi_k3.amd.linear import KimiRoutedOutputTransform
29
from vllm.platforms import current_platform
@@ -56,21 +57,26 @@ def _build_transform(device: torch.device) -> KimiRoutedOutputTransform:
57
58
59
def _tail_runner(
59
- transform: KimiRoutedOutputTransform, tp_size: int
60
+ transform: KimiRoutedOutputTransform, tp_world: int, use_ep: bool = False
61
) -> ROCmLatentMoERunner:
62
"""A runner carrying only what the tail reads, so no engine is needed.
63
64
``_maybe_reduce_final_output`` is left as the real base-class method, so the
65
all-reduce that stitches the shards is the real collective.
66
+
67
+ ``tp_world`` is the size of the TP process group, which is what the hidden
68
+ dim is split by. Under expert parallelism the MoE config reports
69
+ ``tp_size=1``/``ep_size=tp_world`` for that same group, so ``use_ep`` pins
70
+ that the final reduce still fires.
71
"""
72
runner = object.__new__(ROCmLatentMoERunner)
73
attrs = {
74
"routed_output_transform": transform,
69
- "_up_proj_shard_size": HIDDEN_SIZE // tp_size,
75
+ "_up_proj_shard_size": HIDDEN_SIZE // tp_world,
76
"_logged_sharded_tail": False,
77
"moe_config": SimpleNamespace(
72
- tp_size=tp_size,
73
- ep_size=1,
78
+ tp_size=1 if use_ep else tp_world,
79
+ ep_size=tp_world if use_ep else 1,
80
is_sequence_parallel=False,
81
skip_final_all_reduce=False,
82
),
@@ -94,9 +100,11 @@ def _rank_partials(
100
return routed.mul_(0.01), shared
101
102
97
-def _check_matches_replicated(device: torch.device, tp_size: int, rank: int) -> None:
103
+def _check_matches_replicated(
104
+ device: torch.device, tp_world: int, rank: int, use_ep: bool = False
105
+) -> None:
106
transform = _build_transform(device)
99
- runner = _tail_runner(transform, tp_size)
107
+ runner = _tail_runner(transform, tp_world, use_ep)
108
group = get_tp_group().device_group
109
110
for iteration, num_tokens in enumerate((1, 5, 8, 16, 5)):
@@ -119,13 +127,22 @@ def _check_matches_replicated(device: torch.device, tp_size: int, rank: int) ->
127
torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
128
129
130
+def _check_matches_replicated_under_ep(
131
+ device: torch.device, tp_world: int, rank: int
132
+) -> None:
133
+ """The served config: the routed experts are expert-parallel, so the MoE
134
+ config reports no tensor parallelism at all for a group that still splits
135
+ the hidden dim ``tp_world`` ways."""
136
+ _check_matches_replicated(device, tp_world, rank, use_ep=True)
137
+
138
+
139
def _check_writes_only_its_own_shard(
123
- device: torch.device, tp_size: int, rank: int
140
+ device: torch.device, tp_world: int, rank: int
141
) -> None:
142
"""Two ranks that swapped both weight rows and write offsets still sum to
143
the right total, so inspect each slice before the final collective."""
144
transform = _build_transform(device)
128
- runner = _tail_runner(transform, tp_size)
145
+ runner = _tail_runner(transform, tp_world)
146
group = get_tp_group().device_group
147
148
torch.manual_seed(rank + 1)
@@ -149,7 +166,7 @@ def _check_writes_only_its_own_shard(
166
167
assert captured["output_is_reduced"] is False
168
152
- shard = HIDDEN_SIZE // tp_size
169
+ shard = HIDDEN_SIZE // tp_world
170
start, end = rank * shard, (rank + 1) * shard
171
local = captured["states"]
172
projected = F.linear(latent, transform.up_proj.weight[start:end])
@@ -163,6 +180,7 @@ def _check_writes_only_its_own_shard(
180
181
_CHECKS = {
182
"matches_replicated": _check_matches_replicated,
183
+ "matches_replicated_ep": _check_matches_replicated_under_ep,
184
"own_shard_only": _check_writes_only_its_own_shard,
185
}
186
@@ -196,6 +214,11 @@ def test_sharded_tail_tp8_matches_replicated_projection() -> None:
214
_run_ranks("matches_replicated", 8)
215
216
217
+@multi_gpu_test(num_gpus=8)
218
+def test_sharded_tail_tp8_ep8_matches_replicated_projection() -> None:
219
+ _run_ranks("matches_replicated_ep", 8)
220
+
221
+
222
@multi_gpu_test(num_gpus=4)
223
def test_sharded_tail_tp4_writes_only_its_own_shard() -> None:
224
_run_ranks("own_shard_only", 4)
@@ -241,12 +264,19 @@ def test_forward_shards_only_when_the_tail_is_valid(
264
265
@pytest.fixture
266
def build_runner(monkeypatch: pytest.MonkeyPatch):
244
- """Construct through the real subclass ``__init__``, stubbing only the base."""
267
+ """Construct through the real subclass ``__init__``, stubbing only the base.
268
+
269
+ ``tp_world`` is the size of the TP process group and is varied independently
270
+ of ``moe_config.tp_size``, because under expert parallelism those two are
271
+ different numbers for the same eight ranks.
272
+ """
273
+ world = {"size": 8}
274
275
def _base_init(
276
self,
277
*,
278
tp_size: int = 8,
279
+ ep_size: int = 1,
280
hidden: int = HIDDEN_SIZE,
281
is_sequence_parallel: bool = False,
282
has_up_proj: bool = True,
@@ -254,7 +284,9 @@ def build_runner(monkeypatch: pytest.MonkeyPatch):
284
routed_scaling_factor: float = 1.0,
285
) -> None:
286
self.moe_config = SimpleNamespace(
257
- tp_size=tp_size, is_sequence_parallel=is_sequence_parallel
287
+ tp_size=tp_size,
288
+ ep_size=ep_size,
289
+ is_sequence_parallel=is_sequence_parallel,
290
)
291
self.routed_output_transform = SimpleNamespace(
292
norm=None,
@@ -266,20 +298,63 @@ def build_runner(monkeypatch: pytest.MonkeyPatch):
298
self._shared_experts = object() if has_shared_experts else None
299
300
monkeypatch.setattr(moe_runner.MoERunner, "__init__", _base_init)
269
- return ROCmLatentMoERunner
301
+ monkeypatch.setattr(
302
+ latent_moe_runner,
303
+ "get_tensor_model_parallel_world_size",
304
+ lambda: world["size"],
305
+ )
306
+
307
+ def _build(*, tp_world: int = 8, **kwargs) -> ROCmLatentMoERunner:
308
+ world["size"] = tp_world
309
+ return ROCmLatentMoERunner(**kwargs)
310
+
311
+ return _build
312
313
314
def test_shards_under_the_kimi_k3_serving_config(build_runner) -> None:
273
- runner = build_runner()
315
+ """``--tensor-parallel-size 8``, MoE experts tensor-parallel."""
316
+ runner = build_runner(tp_world=8, tp_size=8, ep_size=1)
317
318
assert runner._tail_shardable
319
assert runner._up_proj_shard_size == HIDDEN_SIZE // 8
320
321
322
+def test_shards_under_the_kimi_k3_serving_config_with_expert_parallel(
323
+ build_runner,
324
+) -> None:
325
+ """``--tensor-parallel-size 8 --enable-expert-parallel``, what is served.
326
+
327
+ EP reports ``moe_config.tp_size == 1`` because each rank owns whole experts,
328
+ but the hidden dim is still split across all eight TP ranks, so the tail
329
+ must shard exactly as it does without EP.
330
+ """
331
+ runner = build_runner(tp_world=8, tp_size=1, ep_size=8)
332
+
333
+ assert runner._tail_shardable
334
+ assert runner._up_proj_shard_size == HIDDEN_SIZE // 8
335
+
336
+
337
+def test_shards_by_the_process_group_not_the_dp_flattened_tp_size(
338
+ build_runner,
339
+) -> None:
340
+ """Without EP, ``moe_config.tp_size`` is flattened across DP: TP=2 x DP=2
341
+ reports 4 for a TP process group of 2.
342
+
343
+ Only those 2 ranks share a hidden dim and take part in the final
344
+ all-reduce, so sharding by 4 would leave half of it with no routed
345
+ contribution on either rank.
346
+ """
347
+ runner = build_runner(tp_world=2, tp_size=4, ep_size=1)
348
+
349
+ assert runner._tail_shardable
350
+ assert runner._up_proj_shard_size == HIDDEN_SIZE // 2
351
+
352
+
353
@pytest.mark.parametrize(
354
"override",
355
[
282
- pytest.param({"tp_size": 1}, id="no-tp"),
356
+ pytest.param({"tp_world": 1, "tp_size": 1}, id="single-rank"),
357
+ pytest.param({"tp_world": 1, "tp_size": 4}, id="dp-without-tp"),
358
pytest.param({"hidden": HIDDEN_SIZE + 2}, id="hidden-not-divisible-by-tp"),
359
pytest.param({"has_up_proj": False}, id="no-up-proj"),
360
pytest.param({"has_shared_experts": False}, id="no-shared-partial"),
@@ -290,7 +365,12 @@ def test_shards_under_the_kimi_k3_serving_config(build_runner) -> None:
365
def test_falls_back_when_the_config_breaks_the_shard(
366
build_runner, override: dict
367
) -> None:
293
- """Each of these makes the sharded tail wrong, not merely slower."""
368
+ """Each of these makes the sharded tail wrong, not merely slower.
369
+
370
+ ``dp-without-tp`` is DP=4 with no TP: the hidden dim is not split at all,
371
+ so there is nothing to shard and no all-reduce to stitch shards back
372
+ together, however large ``moe_config.tp_size`` claims to be.
373
+ """
374
runner = build_runner(**override)
375
376
assert not runner._tail_shardable
vllm/models/kimi_k3/amd/latent_moe_runner.py
+2
-1
@@ -6,6 +6,7 @@ import torch
6
7
from vllm.distributed import (
8
get_tensor_model_parallel_rank,
9
+ get_tensor_model_parallel_world_size,
10
tensor_model_parallel_all_reduce,
11
)
12
from vllm.logger import init_logger
@@ -33,7 +34,7 @@ class ROCmLatentMoERunner(MoERunner):
34
35
transform = self.routed_output_transform
36
up_proj = getattr(transform, "up_proj", None)
36
- tp_size = self.moe_config.tp_size
37
+ tp_size = get_tensor_model_parallel_world_size()
38
39
self._up_proj_shard_size = 0
40
self._tail_shardable = (