@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 43d691ec6b

[ROCm][Perf] Kimi-K3 Shard Latent MoE up-projection for ROCm path (#51253)

Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com>

kliuae committed Aug 7, 2026 at 14:11 UTC 43d691ec6b1d26d3ef3d8725a7c7e4d8556eb984
4 files changed +453
tests/models/kimi_k3/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/models/kimi_k3/test_amd_latent_moe_runner.py new
+297
@@ -0,0 +1,297 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""The ROCm latent-MoE tail must equal the replicated up-projection.
4 +
5 +A wrong shard offset or a dropped accumulation still runs and still produces
6 +plausible text, so these pin the arithmetic rather than the behaviour.
7 +"""
8 +
9 +from types import SimpleNamespace
10 +
11 +import pytest
12 +import torch
13 +import torch.distributed as dist
14 +import torch.nn.functional as F
15 +from torch.multiprocessing import spawn
16 +
17 +from tests.utils import (
18 + ensure_current_vllm_config,
19 + init_test_distributed_environment,
20 + multi_gpu_test,
21 +)
22 +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.latent_moe_runner import ROCmLatentMoERunner
27 +from vllm.models.kimi_k3.amd.linear import KimiRoutedOutputTransform
28 +from vllm.platforms import current_platform
29 +from vllm.utils.network_utils import get_open_port
30 +
31 +pytestmark = pytest.mark.skipif(
32 + not current_platform.is_rocm(),
33 + reason="ROCmLatentMoERunner is only wired up on ROCm",
34 +)
35 +
36 +HIDDEN_SIZE = 7168
37 +LATENT_SIZE = 3584
38 +EPS = 1e-5
39 +DTYPE = torch.bfloat16
40 +
41 +
42 +def _build_transform(device: torch.device) -> KimiRoutedOutputTransform:
43 + norm = RMSNorm(LATENT_SIZE, eps=EPS).to(device=device, dtype=DTYPE)
44 + up_proj = ReplicatedLinear(
45 + LATENT_SIZE,
46 + HIDDEN_SIZE,
47 + bias=False,
48 + params_dtype=DTYPE,
49 + prefix="routed_expert_up_proj",
50 + ).to(device=device)
51 +
52 + torch.manual_seed(0)
53 + norm.weight.data.copy_(1 + 0.1 * torch.randn_like(norm.weight))
54 + up_proj.weight.data.copy_(torch.randn_like(up_proj.weight) / LATENT_SIZE**0.5)
55 + return KimiRoutedOutputTransform(norm, up_proj)
56 +
57 +
58 +def _tail_runner(
59 + transform: KimiRoutedOutputTransform, tp_size: int
60 +) -> ROCmLatentMoERunner:
61 + """A runner carrying only what the tail reads, so no engine is needed.
62 +
63 + ``_maybe_reduce_final_output`` is left as the real base-class method, so the
64 + all-reduce that stitches the shards is the real collective.
65 + """
66 + runner = object.__new__(ROCmLatentMoERunner)
67 + attrs = {
68 + "routed_output_transform": transform,
69 + "_up_proj_shard_size": HIDDEN_SIZE // tp_size,
70 + "_logged_sharded_tail": False,
71 + "moe_config": SimpleNamespace(
72 + tp_size=tp_size,
73 + ep_size=1,
74 + is_sequence_parallel=False,
75 + skip_final_all_reduce=False,
76 + ),
77 + }
78 + for name, value in attrs.items():
79 + object.__setattr__(runner, name, value)
80 + return runner
81 +
82 +
83 +def _all_reduced(tensor: torch.Tensor, group) -> torch.Tensor:
84 + reduced = tensor.clone()
85 + dist.all_reduce(reduced, group=group)
86 + return reduced
87 +
88 +
89 +def _rank_partials(
90 + num_tokens: int, device: torch.device
91 +) -> tuple[torch.Tensor, torch.Tensor]:
92 + routed = torch.randn(num_tokens, LATENT_SIZE, device=device, dtype=DTYPE)
93 + shared = torch.randn(num_tokens, HIDDEN_SIZE, device=device, dtype=DTYPE)
94 + return routed.mul_(0.01), shared
95 +
96 +
97 +def _check_matches_replicated(device: torch.device, tp_size: int, rank: int) -> None:
98 + transform = _build_transform(device)
99 + runner = _tail_runner(transform, tp_size)
100 + group = get_tp_group().device_group
101 +
102 + for iteration, num_tokens in enumerate((1, 5, 8, 16, 5)):
103 + torch.manual_seed(100 * iteration + rank + 1)
104 + routed_output, shared_output = _rank_partials(num_tokens, device)
105 +
106 + expected = F.linear(
107 + F.rms_norm(
108 + _all_reduced(routed_output, group),
109 + (LATENT_SIZE,),
110 + transform.norm.weight,
111 + EPS,
112 + ),
113 + transform.up_proj.weight,
114 + )
115 + expected.add_(_all_reduced(shared_output, group))
116 +
117 + actual = runner._shard_up_proj_tail(routed_output, shared_output, None)
118 +
119 + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
120 +
121 +
122 +def _check_writes_only_its_own_shard(
123 + device: torch.device, tp_size: int, rank: int
124 +) -> None:
125 + """Two ranks that swapped both weight rows and write offsets still sum to
126 + the right total, so inspect each slice before the final collective."""
127 + transform = _build_transform(device)
128 + runner = _tail_runner(transform, tp_size)
129 + group = get_tp_group().device_group
130 +
131 + torch.manual_seed(rank + 1)
132 + routed_output, shared_output = _rank_partials(8, device)
133 + before = shared_output.clone()
134 + latent = transform.norm(_all_reduced(routed_output, group))
135 +
136 + captured: dict = {}
137 + real_reduce = runner._maybe_reduce_final_output
138 +
139 + def _capture(states, trunc_size, output_is_reduced=None):
140 + captured["states"] = states.clone()
141 + captured["output_is_reduced"] = output_is_reduced
142 + return real_reduce(states, trunc_size, output_is_reduced)
143 +
144 + object.__setattr__(runner, "_maybe_reduce_final_output", _capture)
145 + runner._shard_up_proj_tail(routed_output, shared_output, None)
146 + # Restoring the method breaks the runner -> _capture -> runner cycle, so the
147 + # captured device tensor is freed before the process group is torn down.
148 + object.__setattr__(runner, "_maybe_reduce_final_output", real_reduce)
149 +
150 + assert captured["output_is_reduced"] is False
151 +
152 + shard = HIDDEN_SIZE // tp_size
153 + start, end = rank * shard, (rank + 1) * shard
154 + local = captured["states"]
155 + projected = F.linear(latent, transform.up_proj.weight[start:end])
156 +
157 + torch.testing.assert_close(
158 + local[:, start:end], before[:, start:end] + projected, atol=8e-2, rtol=3e-2
159 + )
160 + torch.testing.assert_close(local[:, :start], before[:, :start], atol=0, rtol=0)
161 + torch.testing.assert_close(local[:, end:], before[:, end:], atol=0, rtol=0)
162 +
163 +
164 +_CHECKS = {
165 + "matches_replicated": _check_matches_replicated,
166 + "own_shard_only": _check_writes_only_its_own_shard,
167 +}
168 +
169 +
170 +def _worker(local_rank: int, world_size: int, port: str, check: str) -> None:
171 + device = torch.device(f"cuda:{local_rank}")
172 + torch.accelerator.set_device_index(device)
173 + with ensure_current_vllm_config():
174 + init_test_distributed_environment(
175 + world_size, 1, local_rank, port, local_rank=local_rank
176 + )
177 + _CHECKS[check](device, world_size, local_rank)
178 +
179 +
180 +def _run_ranks(check: str, tp_size: int) -> None:
181 + spawn(
182 + _worker,
183 + args=(tp_size, str(get_open_port()), check),
184 + nprocs=tp_size,
185 + join=True,
186 + )
187 +
188 +
189 +@multi_gpu_test(num_gpus=4)
190 +def test_sharded_tail_tp4_matches_replicated_projection() -> None:
191 + _run_ranks("matches_replicated", 4)
192 +
193 +
194 +@multi_gpu_test(num_gpus=8)
195 +def test_sharded_tail_tp8_matches_replicated_projection() -> None:
196 + _run_ranks("matches_replicated", 8)
197 +
198 +
199 +@multi_gpu_test(num_gpus=4)
200 +def test_sharded_tail_tp4_writes_only_its_own_shard() -> None:
201 + _run_ranks("own_shard_only", 4)
202 +
203 +
204 +def _runner(**attrs) -> ROCmLatentMoERunner:
205 + runner = object.__new__(ROCmLatentMoERunner)
206 + for name, value in attrs.items():
207 + object.__setattr__(runner, name, value)
208 + return runner
209 +
210 +
211 +@pytest.mark.parametrize(
212 + "shardable,pre_reduced,expect_sharded",
213 + [
214 + (True, False, True),
215 + (True, True, False),
216 + (False, False, False),
217 + (False, True, False),
218 + ],
219 +)
220 +def test_forward_shards_only_when_the_tail_is_valid(
221 + shardable: bool,
222 + pre_reduced: bool,
223 + expect_sharded: bool,
224 + monkeypatch: pytest.MonkeyPatch,
225 +) -> None:
226 + """The shard needs an un-reduced shared partial to accumulate into."""
227 + base_output, sharded_output = torch.zeros(1), torch.ones(1)
228 + monkeypatch.setattr(moe_runner.MoERunner, "forward", lambda *a, **k: base_output)
229 + monkeypatch.setattr(
230 + ROCmLatentMoERunner, "_fused_forward", lambda *a, **k: sharded_output
231 + )
232 + monkeypatch.setattr(
233 + ROCmLatentMoERunner, "_fused_output_is_reduced", property(lambda _: pre_reduced)
234 + )
235 + runner = _runner(_tail_shardable=shardable)
236 +
237 + result = runner.forward(torch.zeros(1), torch.zeros(1))
238 +
239 + assert result is (sharded_output if expect_sharded else base_output)
240 +
241 +
242 +@pytest.fixture
243 +def build_runner(monkeypatch: pytest.MonkeyPatch):
244 + """Construct through the real subclass ``__init__``, stubbing only the base."""
245 +
246 + def _base_init(
247 + self,
248 + *,
249 + tp_size: int = 8,
250 + hidden: int = HIDDEN_SIZE,
251 + is_sequence_parallel: bool = False,
252 + has_up_proj: bool = True,
253 + has_shared_experts: bool = True,
254 + routed_scaling_factor: float = 1.0,
255 + ) -> None:
256 + self.moe_config = SimpleNamespace(
257 + tp_size=tp_size, is_sequence_parallel=is_sequence_parallel
258 + )
259 + self.routed_output_transform = SimpleNamespace(
260 + norm=None,
261 + up_proj=SimpleNamespace(weight=torch.zeros(hidden, LATENT_SIZE))
262 + if has_up_proj
263 + else None,
264 + )
265 + self.routed_scaling_factor = routed_scaling_factor
266 + self._shared_experts = object() if has_shared_experts else None
267 +
268 + monkeypatch.setattr(moe_runner.MoERunner, "__init__", _base_init)
269 + return ROCmLatentMoERunner
270 +
271 +
272 +def test_shards_under_the_kimi_k3_serving_config(build_runner) -> None:
273 + runner = build_runner()
274 +
275 + assert runner._tail_shardable
276 + assert runner._up_proj_shard_size == HIDDEN_SIZE // 8
277 +
278 +
279 +@pytest.mark.parametrize(
280 + "override",
281 + [
282 + pytest.param({"tp_size": 1}, id="no-tp"),
283 + pytest.param({"hidden": HIDDEN_SIZE + 2}, id="hidden-not-divisible-by-tp"),
284 + pytest.param({"has_up_proj": False}, id="no-up-proj"),
285 + pytest.param({"has_shared_experts": False}, id="no-shared-partial"),
286 + pytest.param({"is_sequence_parallel": True}, id="sequence-parallel"),
287 + pytest.param({"routed_scaling_factor": 2.0}, id="routed-scaling-factor"),
288 + ],
289 +)
290 +def test_falls_back_when_the_config_breaks_the_shard(
291 + build_runner, override: dict
292 +) -> None:
293 + """Each of these makes the sharded tail wrong, not merely slower."""
294 + runner = build_runner(**override)
295 +
296 + assert not runner._tail_shardable
297 + assert runner._up_proj_shard_size == 0
vllm/models/kimi_k3/amd/latent_moe_runner.py new
+152
@@ -0,0 +1,152 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +
5 +from vllm.distributed import (
6 + get_tensor_model_parallel_rank,
7 + tensor_model_parallel_all_reduce,
8 +)
9 +from vllm.logger import init_logger
10 +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner, _unpack
11 +
12 +logger = init_logger(__name__)
13 +
14 +
15 +class ROCmLatentMoERunner(MoERunner):
16 + """MoE runner for latent MoE with a replicated routed up-projection.
17 +
18 + Mirrors CUDA's LatentMoERunner, but currently only the up projection
19 + -sharded path is implemented. (Tier 2)
20 +
21 + Native path: the replicated up-proj produces the full hidden dim on every
22 + rank, so the base runner combines routed + shared correctly at any TP size.
23 + """
24 +
25 + def __init__(
26 + self,
27 + *args,
28 + **kwargs,
29 + ) -> None:
30 + super().__init__(*args, **kwargs)
31 +
32 + transform = self.routed_output_transform
33 + up_proj = getattr(transform, "up_proj", None)
34 + tp_size = self.moe_config.tp_size
35 +
36 + self._up_proj_shard_size = 0
37 + self._tail_shardable = (
38 + up_proj is not None
39 + and tp_size > 1
40 + and up_proj.weight.shape[0] % tp_size == 0
41 + and self._shared_experts is not None
42 + and not self.moe_config.is_sequence_parallel
43 + and self.routed_scaling_factor == 1.0
44 + )
45 + if self._tail_shardable:
46 + assert up_proj is not None
47 + self._up_proj_shard_size = up_proj.weight.shape[0] // tp_size
48 + else:
49 + logger.warning_once(
50 + "K3 latent-MoE tail is not shardable under this config, "
51 + "falling back to the replicated up-projection.",
52 + scope="global",
53 + )
54 + self._logged_sharded_tail = False
55 +
56 + def _shard_up_proj_tail(
57 + self,
58 + fused_output: torch.Tensor,
59 + shared_output: torch.Tensor,
60 + trunc_size: int | None,
61 + ) -> torch.Tensor:
62 + """
63 + Tier 2: column-parallel up-projection folded into the final reduce.
64 + """
65 + if not self._logged_sharded_tail:
66 + self._logged_sharded_tail = True
67 + logger.info_once(
68 + "Kimi-K3 latent-MoE tail: up-projecting only this rank's "
69 + "hidden shard into the shared output.",
70 + scope="global",
71 + )
72 +
73 + transform = self.routed_output_transform
74 + assert transform is not None
75 +
76 + latent = tensor_model_parallel_all_reduce(fused_output)
77 + if transform.norm is not None:
78 + latent = transform.norm(latent)
79 +
80 + shard_size = self._up_proj_shard_size
81 + shard_start = get_tensor_model_parallel_rank() * shard_size
82 + up_proj_shard = transform.up_proj.weight.narrow(0, shard_start, shard_size)
83 + hidden_shard = shared_output.narrow(-1, shard_start, shard_size)
84 +
85 + # hidden_shard += latent @ up_proj_shard.T, accumulated in the GEMM's
86 + # beta-add epilogue so folding in the shared partial costs no kernel.
87 + hidden_shard.addmm_(latent, up_proj_shard.t())
88 +
89 + return self._maybe_reduce_final_output(
90 + shared_output, trunc_size, output_is_reduced=False
91 + )
92 +
93 + def forward(
94 + self,
95 + hidden_states: torch.Tensor,
96 + router_logits: torch.Tensor,
97 + input_ids: torch.Tensor | None = None,
98 + shared_experts_input: torch.Tensor | None = None,
99 + ) -> torch.Tensor:
100 + if self._tail_shardable and not self._fused_output_is_reduced:
101 + return self._fused_forward(
102 + hidden_states, router_logits, input_ids, shared_experts_input
103 + )
104 + return super().forward(
105 + hidden_states, router_logits, input_ids, shared_experts_input
106 + )
107 +
108 + def _fused_forward(
109 + self,
110 + hidden_states: torch.Tensor,
111 + router_logits: torch.Tensor,
112 + input_ids: torch.Tensor | None,
113 + shared_experts_input: torch.Tensor | None,
114 + ) -> torch.Tensor:
115 + # When the caller pre-applies the routed input transform outside the
116 + # runner (e.g. to overlap it on a separate stream), it passes the
117 + # already-transformed routed input as ``hidden_states`` and the original
118 + # hidden states as ``shared_experts_input``; skip the transform then.
119 + if shared_experts_input is None:
120 + hidden_states, shared_experts_input = self.apply_routed_input_transform(
121 + hidden_states
122 + )
123 +
124 + hidden_states, og_hidden_dim_pre_xform, og_hidden_dim_post_xform = (
125 + self._maybe_pad_hidden_states(
126 + shared_experts_input,
127 + hidden_states,
128 + )
129 + )
130 +
131 + result = self._forward_entry(
132 + hidden_states,
133 + router_logits,
134 + shared_experts_input,
135 + input_ids,
136 + self._encode_layer_name(),
137 + self.moe_config.hidden_dim_unpadded
138 + if self._quant_method.has_unpadded_output
139 + else 0,
140 + )
141 +
142 + shared_output, fused_output = _unpack(result)
143 + assert shared_output is not None
144 +
145 + if og_hidden_dim_pre_xform is not None:
146 + fused_output = fused_output[..., :og_hidden_dim_pre_xform]
147 +
148 + result = self._shard_up_proj_tail(
149 + fused_output, shared_output, og_hidden_dim_post_xform
150 + )
151 +
152 + return self._maybe_add_zero_expert_output(result)
vllm/models/kimi_k3/amd/linear.py
+2
@@ -62,6 +62,7 @@ from vllm.model_executor.models.utils import (
62 maybe_prefix,
63 )
64 from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention
65 +from vllm.models.kimi_k3.amd.latent_moe_runner import ROCmLatentMoERunner
66 from vllm.models.kimi_k3.amd.ops.attn_res import attn_res
67 from vllm.sequence import IntermediateTensors
68 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
@@ -289,6 +290,7 @@ class KimiMoE(nn.Module):
290 routed_scaling_factor=self.routed_scaling_factor,
291 routed_input_transform=self.routed_expert_down_proj,
292 routed_output_transform=self.routed_output_transform,
293 + runner_cls=ROCmLatentMoERunner if self.use_latent_moe else None,
294 )
295 if self.padded_moe_intermediate_size != moe_intermediate_size:
296 w13_weight = getattr(self.experts, "w13_weight", None)