| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import math |
| 5 | from types import MethodType, SimpleNamespace |
| 6 | from unittest.mock import Mock |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | from torch import nn |
| 11 | |
| 12 | from vllm.config import ParallelConfig |
| 13 | from vllm.model_executor.layers.activation import SiluAndMul |
| 14 | from vllm.models.common.ops import sequence_parallel as sp_ops |
| 15 | from vllm.models.kimi_k3.nvidia import model as kimi_model |
| 16 | from vllm.models.kimi_k3.nvidia import mtp as kimi_mtp |
| 17 | from vllm.platforms import current_platform |
| 18 | |
| 19 | |
| 20 | class _IdentityNorm(nn.Module): |
| 21 | def __init__(self, hidden_size: int = 2) -> None: |
| 22 | super().__init__() |
| 23 | self.weight = nn.Parameter(torch.ones(hidden_size), requires_grad=False) |
| 24 | self.variance_epsilon = 1e-5 |
| 25 | |
| 26 | def forward( |
| 27 | self, |
| 28 | hidden_states: torch.Tensor, |
| 29 | residual: torch.Tensor | None = None, |
| 30 | ): |
| 31 | if residual is None: |
| 32 | return hidden_states |
| 33 | return hidden_states, residual |
| 34 | |
| 35 | |
| 36 | class _RecordingMoE(nn.Module): |
| 37 | def __init__(self) -> None: |
| 38 | super().__init__() |
| 39 | self.num_tokens = 0 |
| 40 | |
| 41 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 42 | self.num_tokens = hidden_states.shape[0] |
| 43 | return hidden_states |
| 44 | |
| 45 | |
| 46 | class _Projection(nn.Module): |
| 47 | def __init__(self, hidden_size: int = 2) -> None: |
| 48 | super().__init__() |
| 49 | self.weight = nn.Parameter( |
| 50 | torch.ones(1, hidden_size), |
| 51 | requires_grad=False, |
| 52 | ) |
| 53 | |
| 54 | |
| 55 | class _SequenceParallelMTPBlock: |
| 56 | use_sequence_parallel = True |
| 57 | |
| 58 | def __call__( |
| 59 | self, |
| 60 | *, |
| 61 | positions: torch.Tensor, |
| 62 | hidden_states: torch.Tensor, |
| 63 | residual: torch.Tensor | None, |
| 64 | ): |
| 65 | assert residual is None |
| 66 | return hidden_states * 2, None, hidden_states * 3 |
| 67 | |
| 68 | |
| 69 | def _mock_sequence_parallel_collectives(monkeypatch): |
| 70 | monkeypatch.setattr( |
| 71 | kimi_model, |
| 72 | "sp_reduce_scatter", |
| 73 | lambda tensor: tensor.chunk(2, dim=0)[0], |
| 74 | ) |
| 75 | monkeypatch.setattr( |
| 76 | kimi_model, |
| 77 | "sp_shard", |
| 78 | lambda tensor: torch.nn.functional.pad(tensor, (0, 0, 0, 1))[:2], |
| 79 | ) |
| 80 | monkeypatch.setattr( |
| 81 | kimi_model, |
| 82 | "sp_all_gather", |
| 83 | lambda tensor: torch.cat([tensor, tensor], dim=0), |
| 84 | ) |
| 85 | |
| 86 | |
| 87 | @pytest.mark.parametrize( |
| 88 | ("num_tokens", "is_padding", "tp_rank", "expected"), |
| 89 | [ |
| 90 | (1, None, 0, [False]), |
| 91 | (1, None, 1, [True]), |
| 92 | (5, None, 2, [False, True]), |
| 93 | (5, None, 3, [True, True]), |
| 94 | (5, [False, True, False, False, False], 0, [False, True]), |
| 95 | ], |
| 96 | ) |
| 97 | def test_sp_padding_mask_marks_added_rows( |
| 98 | monkeypatch, |
| 99 | num_tokens: int, |
| 100 | is_padding: list[bool] | None, |
| 101 | tp_rank: int, |
| 102 | expected: list[bool], |
| 103 | ): |
| 104 | monkeypatch.setattr(sp_ops, "get_tensor_model_parallel_world_size", lambda: 4) |
| 105 | monkeypatch.setattr(sp_ops, "get_tensor_model_parallel_rank", lambda: tp_rank) |
| 106 | |
| 107 | hidden_states = torch.empty(num_tokens, 2) |
| 108 | padding = torch.tensor(is_padding) if is_padding is not None else None |
| 109 | actual = sp_ops.sp_padding_mask(padding, hidden_states) |
| 110 | |
| 111 | torch.testing.assert_close(actual, torch.tensor(expected)) |
| 112 | |
| 113 | |
| 114 | @pytest.mark.parametrize( |
| 115 | ("data_parallel_size", "expected"), |
| 116 | [ |
| 117 | (1, False), |
| 118 | (2, True), |
| 119 | ], |
| 120 | ) |
| 121 | def test_moe_sequence_parallel_requires_data_parallel( |
| 122 | monkeypatch, |
| 123 | data_parallel_size: int, |
| 124 | expected: bool, |
| 125 | ): |
| 126 | monkeypatch.setattr(current_platform, "device_count", lambda: 2) |
| 127 | parallel_config = ParallelConfig( |
| 128 | tensor_parallel_size=2, |
| 129 | data_parallel_size=data_parallel_size, |
| 130 | enable_expert_parallel=True, |
| 131 | all2all_backend="allgather_reducescatter", |
| 132 | ) |
| 133 | |
| 134 | assert parallel_config.use_sequence_parallel_moe is expected |
| 135 | |
| 136 | |
| 137 | def test_kimi_decoder_layer_keeps_moe_states_sequence_sharded(monkeypatch): |
| 138 | layer = object.__new__(kimi_model.KimiDecoderLayer) |
| 139 | nn.Module.__init__(layer) |
| 140 | layer.use_attn_res = False |
| 141 | layer.use_sequence_parallel = True |
| 142 | layer.input_layernorm = _IdentityNorm() |
| 143 | layer.post_attention_layernorm = _IdentityNorm() |
| 144 | layer.mlp = _RecordingMoE() |
| 145 | layer._run_self_attn = MethodType( |
| 146 | lambda self, positions, hidden_states: hidden_states, |
| 147 | layer, |
| 148 | ) |
| 149 | |
| 150 | _mock_sequence_parallel_collectives(monkeypatch) |
| 151 | |
| 152 | positions = torch.arange(3) |
| 153 | full_hidden_states = torch.arange(6, dtype=torch.float32).view(3, 2) |
| 154 | hidden_states = kimi_model.sp_shard(full_hidden_states) |
| 155 | hidden_states, prefix_sum, residual = layer( |
| 156 | positions=positions, |
| 157 | hidden_states=hidden_states, |
| 158 | residual=None, |
| 159 | ) |
| 160 | |
| 161 | assert prefix_sum is None |
| 162 | assert hidden_states.shape == residual.shape == (2, 2) |
| 163 | assert layer.mlp.num_tokens == 2 |
| 164 | |
| 165 | hidden_states, prefix_sum, residual = layer( |
| 166 | positions=positions, |
| 167 | hidden_states=hidden_states, |
| 168 | residual=residual, |
| 169 | ) |
| 170 | |
| 171 | assert prefix_sum is None |
| 172 | assert hidden_states.shape == residual.shape == (2, 2) |
| 173 | assert layer.mlp.num_tokens == 2 |
| 174 | |
| 175 | |
| 176 | def test_kimi_attn_residual_states_stay_sequence_sharded(monkeypatch): |
| 177 | layer = object.__new__(kimi_model.KimiDecoderLayer) |
| 178 | nn.Module.__init__(layer) |
| 179 | layer.use_attn_res = True |
| 180 | layer.use_sequence_parallel = True |
| 181 | layer.prev_valid_blocks = 0 |
| 182 | layer.block_write_idx = 0 |
| 183 | layer.is_block_write_layer = False |
| 184 | layer.input_layernorm = _IdentityNorm() |
| 185 | layer.post_attention_layernorm = _IdentityNorm() |
| 186 | layer.self_attention_res_norm = _IdentityNorm() |
| 187 | layer.mlp_res_norm = _IdentityNorm() |
| 188 | layer.self_attention_res_proj = _Projection() |
| 189 | layer.mlp_res_proj = _Projection() |
| 190 | layer.mlp = _RecordingMoE() |
| 191 | layer._run_self_attn = MethodType( |
| 192 | lambda self, positions, hidden_states: hidden_states, |
| 193 | layer, |
| 194 | ) |
| 195 | |
| 196 | _mock_sequence_parallel_collectives(monkeypatch) |
| 197 | monkeypatch.setattr( |
| 198 | kimi_model, |
| 199 | "attn_res", |
| 200 | lambda prefix_sum, hidden_states, *args, **kwargs: ( |
| 201 | prefix_sum if hidden_states is None else prefix_sum + hidden_states |
| 202 | ), |
| 203 | ) |
| 204 | |
| 205 | prefix_sum = kimi_model.sp_shard(torch.arange(6, dtype=torch.float32).view(3, 2)) |
| 206 | block_residual = torch.zeros(2, 1, 2) |
| 207 | hidden_states, prefix_sum, block_residual = layer( |
| 208 | positions=torch.arange(3), |
| 209 | hidden_states=None, |
| 210 | prefix_sum=prefix_sum, |
| 211 | residual=block_residual, |
| 212 | ) |
| 213 | |
| 214 | assert hidden_states.shape == prefix_sum.shape == (2, 2) |
| 215 | assert block_residual.shape == (2, 1, 2) |
| 216 | assert layer.mlp.num_tokens == 2 |
| 217 | |
| 218 | |
| 219 | def test_kimi_mtp_restores_sequence_parallel_output(monkeypatch): |
| 220 | layer = object.__new__(kimi_mtp.KimiK3MultiTokenPredictorLayer) |
| 221 | nn.Module.__init__(layer) |
| 222 | layer.enorm = _IdentityNorm() |
| 223 | layer.hnorm = _IdentityNorm() |
| 224 | layer.eh_proj = nn.Identity() |
| 225 | object.__setattr__(layer, "mtp_block", _SequenceParallelMTPBlock()) |
| 226 | |
| 227 | final_norm = Mock(side_effect=lambda hidden_states: hidden_states + 1) |
| 228 | object.__setattr__( |
| 229 | layer, |
| 230 | "shared_head", |
| 231 | SimpleNamespace(norm=final_norm), |
| 232 | ) |
| 233 | |
| 234 | monkeypatch.setattr( |
| 235 | kimi_mtp, |
| 236 | "fused_mtp_input", |
| 237 | lambda positions, inputs_embeds, *args: inputs_embeds, |
| 238 | ) |
| 239 | monkeypatch.setattr( |
| 240 | kimi_mtp, |
| 241 | "sp_shard", |
| 242 | lambda tensor: torch.nn.functional.pad(tensor, (0, 0, 0, 1))[:2], |
| 243 | ) |
| 244 | monkeypatch.setattr( |
| 245 | kimi_mtp, |
| 246 | "sp_all_gather", |
| 247 | lambda tensor: torch.cat([tensor, tensor], dim=0), |
| 248 | ) |
| 249 | |
| 250 | inputs_embeds = torch.arange(6, dtype=torch.float32).view(3, 2) |
| 251 | logits_hidden_states, hidden_states = layer( |
| 252 | input_ids=torch.zeros(3, dtype=torch.long), |
| 253 | positions=torch.arange(3), |
| 254 | previous_hidden_states=torch.zeros_like(inputs_embeds), |
| 255 | inputs_embeds=inputs_embeds, |
| 256 | ) |
| 257 | |
| 258 | sharded_states = torch.nn.functional.pad(inputs_embeds, (0, 0, 0, 1))[:2] |
| 259 | expected_hidden_states = torch.cat( |
| 260 | [sharded_states * 5, sharded_states * 5], |
| 261 | dim=0, |
| 262 | )[:3] |
| 263 | torch.testing.assert_close(hidden_states, expected_hidden_states) |
| 264 | torch.testing.assert_close(logits_hidden_states, expected_hidden_states + 1) |
| 265 | final_norm.assert_called_once() |
| 266 | torch.testing.assert_close(final_norm.call_args.args[0], expected_hidden_states) |
| 267 | |
| 268 | |
| 269 | @pytest.mark.parametrize( |
| 270 | ("enabled", "use_sequence_parallel", "eligible", "tp_size", "expected"), |
| 271 | [ |
| 272 | (True, True, True, 8, True), |
| 273 | (False, True, True, 8, False), # opt-in only |
| 274 | (True, False, True, 8, False), # replication only exists under SP |
| 275 | (True, True, False, 8, False), # FusedMoE path owns the reduction |
| 276 | (True, True, True, 1, False), # nothing to shard |
| 277 | (True, True, True, 5, False), # 6144 % 5 -- would fail divide() |
| 278 | ], |
| 279 | ) |
| 280 | def test_shard_sequence_parallel_mlp_gating( |
| 281 | monkeypatch, |
| 282 | enabled: bool, |
| 283 | use_sequence_parallel: bool, |
| 284 | eligible: bool, |
| 285 | tp_size: int, |
| 286 | expected: bool, |
| 287 | ): |
| 288 | monkeypatch.setattr(kimi_model.envs, "VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT", enabled) |
| 289 | monkeypatch.setattr( |
| 290 | kimi_model, "get_tensor_model_parallel_world_size", lambda: tp_size |
| 291 | ) |
| 292 | |
| 293 | assert ( |
| 294 | kimi_model.shard_sequence_parallel_mlp( |
| 295 | hidden_size=7168, |
| 296 | intermediate_size=6144, |
| 297 | use_sequence_parallel=use_sequence_parallel, |
| 298 | eligible=eligible, |
| 299 | ) |
| 300 | is expected |
| 301 | ) |
| 302 | |
| 303 | |
| 304 | def test_sharded_sequence_parallel_mlp_matches_replicated(default_vllm_config): |
| 305 | """Sharded SP MLP must reproduce the replicated result for every token. |
| 306 | |
| 307 | Each rank owns a *disjoint* token shard, so a weight shard alone cannot |
| 308 | finish a rank's own tokens: the ranks must gather the full token set, |
| 309 | compute partial sums over their intermediate shard, and reduce-scatter. |
| 310 | Splicing per-rank feature slices together instead silently mixes different |
| 311 | tokens and produces plausible-looking garbage. |
| 312 | """ |
| 313 | tp_size, hidden, intermediate, tokens_per_rank = 4, 16, 12, 3 |
| 314 | torch.manual_seed(0) |
| 315 | num_tokens = tp_size * tokens_per_rank |
| 316 | x = torch.randn(num_tokens, hidden) |
| 317 | gate_weight = torch.randn(intermediate, hidden) |
| 318 | up_weight = torch.randn(intermediate, hidden) |
| 319 | down_weight = torch.randn(hidden, intermediate) |
| 320 | act_fn = SiluAndMul() |
| 321 | |
| 322 | replicated = act_fn(x @ torch.cat([gate_weight, up_weight]).T) @ down_weight.T |
| 323 | |
| 324 | shard = intermediate // tp_size |
| 325 | # Every rank all-gathers the full token set, then computes its partial. |
| 326 | partials = [ |
| 327 | act_fn( |
| 328 | x |
| 329 | @ torch.cat( |
| 330 | [ |
| 331 | gate_weight[r * shard : (r + 1) * shard], |
| 332 | up_weight[r * shard : (r + 1) * shard], |
| 333 | ] |
| 334 | ).T |
| 335 | ) |
| 336 | @ down_weight[:, r * shard : (r + 1) * shard].T |
| 337 | for r in range(tp_size) |
| 338 | ] |
| 339 | reduced = torch.stack(partials).sum(0) |
| 340 | # Reduce-scatter: rank r keeps only its own token shard. |
| 341 | for r in range(tp_size): |
| 342 | mine = reduced[r * tokens_per_rank : (r + 1) * tokens_per_rank] |
| 343 | expected = replicated[r * tokens_per_rank : (r + 1) * tokens_per_rank] |
| 344 | torch.testing.assert_close(mine, expected, atol=1e-5, rtol=1e-5) |
| 345 | |
| 346 | |
| 347 | def test_sp_all_gather_uses_custom_kernel(monkeypatch): |
| 348 | hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2) |
| 349 | expected = torch.cat([hidden_states, hidden_states]) |
| 350 | custom_all_gather = Mock(return_value=expected) |
| 351 | device_communicator = SimpleNamespace( |
| 352 | custom_all_gather=custom_all_gather, |
| 353 | ) |
| 354 | monkeypatch.setattr( |
| 355 | sp_ops, |
| 356 | "get_tp_group", |
| 357 | lambda: SimpleNamespace(device_communicator=device_communicator), |
| 358 | ) |
| 359 | fallback = Mock(side_effect=AssertionError("unexpected fallback")) |
| 360 | monkeypatch.setattr(sp_ops, "tensor_model_parallel_all_gather", fallback) |
| 361 | |
| 362 | output = sp_ops.sp_all_gather(hidden_states) |
| 363 | |
| 364 | torch.testing.assert_close(output, expected) |
| 365 | custom_all_gather.assert_called_once_with(hidden_states) |
| 366 | fallback.assert_not_called() |
| 367 | |
| 368 | |
| 369 | def test_sp_reduce_scatter_uses_custom_kernel_after_padding(monkeypatch): |
| 370 | hidden_states = torch.arange(6, dtype=torch.float32).view(3, 2) |
| 371 | expected = torch.arange(4, dtype=torch.float32).view(2, 2) |
| 372 | custom_reduce_scatter = Mock(return_value=expected) |
| 373 | device_communicator = SimpleNamespace( |
| 374 | custom_reduce_scatter=custom_reduce_scatter, |
| 375 | ) |
| 376 | monkeypatch.setattr( |
| 377 | sp_ops, |
| 378 | "get_tp_group", |
| 379 | lambda: SimpleNamespace(device_communicator=device_communicator), |
| 380 | ) |
| 381 | monkeypatch.setattr( |
| 382 | sp_ops, |
| 383 | "get_tensor_model_parallel_world_size", |
| 384 | lambda: 2, |
| 385 | ) |
| 386 | fallback = Mock(side_effect=AssertionError("unexpected fallback")) |
| 387 | monkeypatch.setattr(sp_ops, "tensor_model_parallel_reduce_scatter", fallback) |
| 388 | |
| 389 | output = sp_ops.sp_reduce_scatter(hidden_states) |
| 390 | |
| 391 | torch.testing.assert_close(output, expected) |
| 392 | padded = custom_reduce_scatter.call_args.args[0] |
| 393 | assert padded.shape == (4, 2) |
| 394 | torch.testing.assert_close(padded[:3], hidden_states) |
| 395 | torch.testing.assert_close(padded[3], torch.zeros(2)) |
| 396 | fallback.assert_not_called() |
| 397 | |
| 398 | |
| 399 | @pytest.mark.parametrize("shape", [(3,), (3, 2, 2)]) |
| 400 | def test_sp_shard_pads_only_the_token_axis(monkeypatch, shape): |
| 401 | hidden_states = torch.arange(math.prod(shape), dtype=torch.float32).view(shape) |
| 402 | monkeypatch.setattr( |
| 403 | sp_ops, |
| 404 | "get_tensor_model_parallel_world_size", |
| 405 | lambda: 2, |
| 406 | ) |
| 407 | monkeypatch.setattr(sp_ops, "get_tensor_model_parallel_rank", lambda: 1) |
| 408 | |
| 409 | output = sp_ops.sp_shard(hidden_states) |
| 410 | |
| 411 | padding = hidden_states.new_zeros((1, *shape[1:])) |
| 412 | expected = torch.cat([hidden_states, padding])[2:] |
| 413 | torch.testing.assert_close(output, expected) |
| 414 | |
| 415 | |
| 416 | def test_sp_collectives_fall_back_without_custom_kernel(monkeypatch): |
| 417 | hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2) |
| 418 | monkeypatch.setattr( |
| 419 | sp_ops, |
| 420 | "get_tp_group", |
| 421 | lambda: SimpleNamespace(device_communicator=None), |
| 422 | ) |
| 423 | monkeypatch.setattr( |
| 424 | sp_ops, |
| 425 | "get_tensor_model_parallel_world_size", |
| 426 | lambda: 2, |
| 427 | ) |
| 428 | all_gather = Mock(return_value=hidden_states) |
| 429 | reduce_scatter = Mock(return_value=hidden_states) |
| 430 | monkeypatch.setattr(sp_ops, "tensor_model_parallel_all_gather", all_gather) |
| 431 | monkeypatch.setattr( |
| 432 | sp_ops, |
| 433 | "tensor_model_parallel_reduce_scatter", |
| 434 | reduce_scatter, |
| 435 | ) |
| 436 | |
| 437 | torch.testing.assert_close(sp_ops.sp_all_gather(hidden_states), hidden_states) |
| 438 | torch.testing.assert_close(sp_ops.sp_reduce_scatter(hidden_states), hidden_states) |
| 439 | all_gather.assert_called_once_with(hidden_states, 0) |
| 440 | reduce_scatter.assert_called_once_with(hidden_states, 0) |