| 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 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 |
| 30 | from vllm.utils.network_utils import get_open_port |
| 31 | |
| 32 | pytestmark = pytest.mark.skipif( |
| 33 | not current_platform.is_rocm(), |
| 34 | reason="ROCmLatentMoERunner is only wired up on ROCm", |
| 35 | ) |
| 36 | |
| 37 | HIDDEN_SIZE = 7168 |
| 38 | LATENT_SIZE = 3584 |
| 39 | EPS = 1e-5 |
| 40 | DTYPE = torch.bfloat16 |
| 41 | |
| 42 | |
| 43 | def _build_transform(device: torch.device) -> KimiRoutedOutputTransform: |
| 44 | norm = RMSNorm(LATENT_SIZE, eps=EPS).to(device=device, dtype=DTYPE) |
| 45 | up_proj = ReplicatedLinear( |
| 46 | LATENT_SIZE, |
| 47 | HIDDEN_SIZE, |
| 48 | bias=False, |
| 49 | params_dtype=DTYPE, |
| 50 | prefix="routed_expert_up_proj", |
| 51 | ).to(device=device) |
| 52 | |
| 53 | torch.manual_seed(0) |
| 54 | norm.weight.data.copy_(1 + 0.1 * torch.randn_like(norm.weight)) |
| 55 | up_proj.weight.data.copy_(torch.randn_like(up_proj.weight) / LATENT_SIZE**0.5) |
| 56 | return KimiRoutedOutputTransform(norm, up_proj) |
| 57 | |
| 58 | |
| 59 | def _tail_runner( |
| 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, |
| 75 | "_up_proj_shard_size": HIDDEN_SIZE // tp_world, |
| 76 | "_logged_sharded_tail": False, |
| 77 | "moe_config": SimpleNamespace( |
| 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 | ), |
| 83 | } |
| 84 | for name, value in attrs.items(): |
| 85 | object.__setattr__(runner, name, value) |
| 86 | return runner |
| 87 | |
| 88 | |
| 89 | def _all_reduced(tensor: torch.Tensor, group) -> torch.Tensor: |
| 90 | reduced = tensor.clone() |
| 91 | dist.all_reduce(reduced, group=group) |
| 92 | return reduced |
| 93 | |
| 94 | |
| 95 | def _rank_partials( |
| 96 | num_tokens: int, device: torch.device |
| 97 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 98 | routed = torch.randn(num_tokens, LATENT_SIZE, device=device, dtype=DTYPE) |
| 99 | shared = torch.randn(num_tokens, HIDDEN_SIZE, device=device, dtype=DTYPE) |
| 100 | return routed.mul_(0.01), shared |
| 101 | |
| 102 | |
| 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) |
| 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)): |
| 111 | torch.manual_seed(100 * iteration + rank + 1) |
| 112 | routed_output, shared_output = _rank_partials(num_tokens, device) |
| 113 | |
| 114 | expected = F.linear( |
| 115 | F.rms_norm( |
| 116 | _all_reduced(routed_output, group), |
| 117 | (LATENT_SIZE,), |
| 118 | transform.norm.weight, |
| 119 | EPS, |
| 120 | ), |
| 121 | transform.up_proj.weight, |
| 122 | ) |
| 123 | expected.add_(_all_reduced(shared_output, group)) |
| 124 | |
| 125 | actual = runner._shard_up_proj_tail(routed_output, shared_output, None) |
| 126 | |
| 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( |
| 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) |
| 145 | runner = _tail_runner(transform, tp_world) |
| 146 | group = get_tp_group().device_group |
| 147 | |
| 148 | torch.manual_seed(rank + 1) |
| 149 | routed_output, shared_output = _rank_partials(8, device) |
| 150 | before = shared_output.clone() |
| 151 | latent = transform.norm(_all_reduced(routed_output, group)) |
| 152 | |
| 153 | captured: dict = {} |
| 154 | real_reduce = runner._maybe_reduce_final_output |
| 155 | |
| 156 | def _capture(states, trunc_size, output_is_reduced=None): |
| 157 | captured["states"] = states.clone() |
| 158 | captured["output_is_reduced"] = output_is_reduced |
| 159 | return real_reduce(states, trunc_size, output_is_reduced) |
| 160 | |
| 161 | object.__setattr__(runner, "_maybe_reduce_final_output", _capture) |
| 162 | runner._shard_up_proj_tail(routed_output, shared_output, None) |
| 163 | # Restoring the method breaks the runner -> _capture -> runner cycle, so the |
| 164 | # captured device tensor is freed before the process group is torn down. |
| 165 | object.__setattr__(runner, "_maybe_reduce_final_output", real_reduce) |
| 166 | |
| 167 | assert captured["output_is_reduced"] is False |
| 168 | |
| 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]) |
| 173 | |
| 174 | torch.testing.assert_close( |
| 175 | local[:, start:end], before[:, start:end] + projected, atol=8e-2, rtol=3e-2 |
| 176 | ) |
| 177 | torch.testing.assert_close(local[:, :start], before[:, :start], atol=0, rtol=0) |
| 178 | torch.testing.assert_close(local[:, end:], before[:, end:], atol=0, rtol=0) |
| 179 | |
| 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 | |
| 187 | |
| 188 | def _worker(local_rank: int, world_size: int, port: str, check: str) -> None: |
| 189 | device = torch.device(f"cuda:{local_rank}") |
| 190 | torch.accelerator.set_device_index(device) |
| 191 | with ensure_current_vllm_config(): |
| 192 | init_test_distributed_environment( |
| 193 | world_size, 1, local_rank, port, local_rank=local_rank |
| 194 | ) |
| 195 | _CHECKS[check](device, world_size, local_rank) |
| 196 | |
| 197 | |
| 198 | def _run_ranks(check: str, tp_size: int) -> None: |
| 199 | spawn( |
| 200 | _worker, |
| 201 | args=(tp_size, str(get_open_port()), check), |
| 202 | nprocs=tp_size, |
| 203 | join=True, |
| 204 | ) |
| 205 | |
| 206 | |
| 207 | @multi_gpu_test(num_gpus=4) |
| 208 | def test_sharded_tail_tp4_matches_replicated_projection() -> None: |
| 209 | _run_ranks("matches_replicated", 4) |
| 210 | |
| 211 | |
| 212 | @multi_gpu_test(num_gpus=8) |
| 213 | 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) |
| 225 | |
| 226 | |
| 227 | def _runner(**attrs) -> ROCmLatentMoERunner: |
| 228 | runner = object.__new__(ROCmLatentMoERunner) |
| 229 | for name, value in attrs.items(): |
| 230 | object.__setattr__(runner, name, value) |
| 231 | return runner |
| 232 | |
| 233 | |
| 234 | @pytest.mark.parametrize( |
| 235 | "shardable,pre_reduced,expect_sharded", |
| 236 | [ |
| 237 | (True, False, True), |
| 238 | (True, True, False), |
| 239 | (False, False, False), |
| 240 | (False, True, False), |
| 241 | ], |
| 242 | ) |
| 243 | def test_forward_shards_only_when_the_tail_is_valid( |
| 244 | shardable: bool, |
| 245 | pre_reduced: bool, |
| 246 | expect_sharded: bool, |
| 247 | monkeypatch: pytest.MonkeyPatch, |
| 248 | ) -> None: |
| 249 | """The shard needs an un-reduced shared partial to accumulate into.""" |
| 250 | base_output, sharded_output = torch.zeros(1), torch.ones(1) |
| 251 | monkeypatch.setattr(moe_runner.MoERunner, "forward", lambda *a, **k: base_output) |
| 252 | monkeypatch.setattr( |
| 253 | ROCmLatentMoERunner, "_fused_forward", lambda *a, **k: sharded_output |
| 254 | ) |
| 255 | monkeypatch.setattr( |
| 256 | ROCmLatentMoERunner, "_fused_output_is_reduced", property(lambda _: pre_reduced) |
| 257 | ) |
| 258 | runner = _runner(_tail_shardable=shardable) |
| 259 | |
| 260 | result = runner.forward(torch.zeros(1), torch.zeros(1)) |
| 261 | |
| 262 | assert result is (sharded_output if expect_sharded else base_output) |
| 263 | |
| 264 | |
| 265 | @pytest.fixture |
| 266 | def build_runner(monkeypatch: pytest.MonkeyPatch): |
| 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, |
| 283 | has_shared_experts: bool = True, |
| 284 | routed_scaling_factor: float = 1.0, |
| 285 | ) -> None: |
| 286 | self.moe_config = SimpleNamespace( |
| 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, |
| 293 | up_proj=SimpleNamespace(weight=torch.zeros(hidden, LATENT_SIZE)) |
| 294 | if has_up_proj |
| 295 | else None, |
| 296 | ) |
| 297 | self.routed_scaling_factor = routed_scaling_factor |
| 298 | self._shared_experts = object() if has_shared_experts else None |
| 299 | |
| 300 | monkeypatch.setattr(moe_runner.MoERunner, "__init__", _base_init) |
| 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: |
| 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 | [ |
| 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"), |
| 361 | pytest.param({"is_sequence_parallel": True}, id="sequence-parallel"), |
| 362 | pytest.param({"routed_scaling_factor": 2.0}, id="routed-scaling-factor"), |
| 363 | ], |
| 364 | ) |
| 365 | def test_falls_back_when_the_config_breaks_the_shard( |
| 366 | build_runner, override: dict |
| 367 | ) -> None: |
| 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 |
| 377 | assert runner._up_proj_shard_size == 0 |