main
py 480 lines 17.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from types import SimpleNamespace
5
6 import pytest
7 import torch
8 from torch import nn
9
10 from vllm.config import ParallelConfig
11 from vllm.model_executor.models.utils import WeightsMapper
12 from vllm.platforms import current_platform
13
14 pytestmark = pytest.mark.skipif(
15 not current_platform.is_rocm(), reason="ROCm-specific tests"
16 )
17
18
19 @pytest.mark.parametrize("hash_routing", [False, True], ids=["regular", "hash"])
20 @pytest.mark.parametrize("vision", [False, True], ids=["text", "vision"])
21 @torch.inference_mode()
22 def test_rocm_moe_routing_and_shared_experts_match_reference(
23 hash_routing, vision, monkeypatch, default_vllm_config, dist_init
24 ) -> None:
25 """AMD V4 routing, clamped experts and shared output agree with PyTorch."""
26 from vllm.forward_context import set_forward_context
27 from vllm.models.deepseek_v4.amd.model import DeepseekV4MoE
28 from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID
29 from vllm.utils.torch_utils import set_default_torch_dtype
30 from vllm.v1.worker.workspace import (
31 init_workspace_manager,
32 is_workspace_manager_initialized,
33 )
34
35 monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "0")
36 monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "0")
37 default_vllm_config.kernel_config.moe_backend = "triton"
38 config = SimpleNamespace(
39 hidden_size=512,
40 n_routed_experts=8,
41 num_experts_per_tok=2,
42 moe_intermediate_size=512,
43 swiglu_limit=0.05,
44 norm_topk_prob=True,
45 scoring_func="sqrtsoftplus",
46 routed_scaling_factor=1.5,
47 num_hash_layers=int(hash_routing),
48 vocab_size=IMAGE_SENTINEL_BASE_ID + 8,
49 topk_method="noaux_tc",
50 vision_n_layers=int(vision),
51 n_shared_experts=1,
52 hidden_act="silu",
53 )
54 torch.manual_seed(42)
55 with torch.device("cuda"), set_default_torch_dtype(torch.bfloat16):
56 moe = DeepseekV4MoE(
57 SimpleNamespace(
58 model_config=SimpleNamespace(hf_config=config),
59 quant_config=None,
60 parallel_config=ParallelConfig(),
61 ),
62 prefix="model.layers.0.ffn",
63 )
64 for parameter in moe.parameters():
65 if parameter.is_floating_point():
66 parameter.normal_(std=0.04)
67 hidden = torch.randn(17, config.hidden_size) * 0.1
68 input_ids = torch.arange(17, dtype=torch.int32)
69 if vision:
70 # Include both ends of the five-token sentinel interval and the
71 # first ordinary token above it.
72 input_ids[1:4] = torch.tensor(
73 [
74 IMAGE_SENTINEL_BASE_ID,
75 IMAGE_SENTINEL_BASE_ID + 4,
76 IMAGE_SENTINEL_BASE_ID + 5,
77 ],
78 dtype=torch.int32,
79 )
80 if moe.gate.tid2eid is not None:
81 rows = torch.arange(config.vocab_size)
82 moe.gate.tid2eid.copy_(torch.stack((rows % 8, (rows + 3) % 8), dim=1))
83 if moe.gate.e_score_correction_bias is not None:
84 moe.gate.e_score_correction_bias.copy_(torch.arange(8) * 0.1)
85 if moe.gate.bias_vl is not None:
86 moe.gate.bias_vl.copy_(torch.arange(8) * -0.1)
87
88 scores = torch.nn.functional.softplus(
89 hidden.float() @ moe.gate.weight.float().T
90 ).sqrt()
91 bias = moe.gate.e_score_correction_bias
92 selection_scores = scores if bias is None else scores + bias
93 expected_ids = selection_scores.topk(2, dim=-1).indices
94 if hash_routing:
95 expected_ids = moe.gate.tid2eid[input_ids.long()].long()
96 if vision:
97 image_mask = (input_ids >= IMAGE_SENTINEL_BASE_ID) & (
98 input_ids < IMAGE_SENTINEL_BASE_ID + 5
99 )
100 expected_ids[image_mask] = (
101 (scores + moe.gate.bias_vl).topk(2, dim=-1).indices[image_mask]
102 )
103 expected_weights = scores.gather(1, expected_ids)
104 expected_weights *= config.routed_scaling_factor / expected_weights.sum(
105 dim=-1, keepdim=True
106 )
107 router_logits, _ = moe.gate(hidden)
108 weights, ids = moe.experts.router.select_experts(
109 hidden, router_logits, topk_indices_dtype=torch.int32, input_ids=input_ids
110 )
111 torch.testing.assert_close(ids.long(), expected_ids)
112 torch.testing.assert_close(weights, expected_weights, rtol=1e-5, atol=1e-6)
113
114 def expert_reference(x, gate_up_weight, down_weight):
115 gate, up = (
116 torch.nn.functional.linear(x, gate_up_weight).float().chunk(2, dim=-1)
117 )
118 gate = gate.clamp(max=config.swiglu_limit)
119 up = up.clamp(-config.swiglu_limit, config.swiglu_limit)
120 activated = (torch.nn.functional.silu(gate) * up).to(x.dtype)
121 return torch.nn.functional.linear(activated, down_weight).float()
122
123 routed = moe.experts.routed_experts
124 expected = torch.zeros_like(hidden, dtype=torch.float32)
125 for expert_id in range(config.n_routed_experts):
126 token, slot = torch.where(expected_ids == expert_id)
127 expected[token] += (
128 expert_reference(
129 hidden[token], routed.w13_weight[expert_id], routed.w2_weight[expert_id]
130 )
131 * expected_weights[token, slot, None]
132 )
133 shared = moe.shared_experts
134 expected += expert_reference(
135 hidden, shared.gate_up_proj.weight, shared.down_proj.weight
136 )
137 moe.experts._quant_method.process_weights_after_loading(routed)
138 if not is_workspace_manager_initialized():
139 init_workspace_manager(torch.device("cuda:0"))
140 with set_forward_context(None, default_vllm_config, num_tokens=hidden.shape[0]):
141 actual = moe(hidden, input_ids)
142 torch.testing.assert_close(actual.float(), expected, rtol=0.02, atol=5e-5)
143
144
145 def test_rocm_packed_kv_cache_auto_uses_ds_mla_layout() -> None:
146 from vllm.config import CacheConfig
147 from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype
148
149 cache_config = CacheConfig()
150
151 resolved_dtype, torch_dtype = _resolve_dsv4_kv_cache_dtype(
152 use_fp8_ds_mla_layout=True,
153 kv_cache_dtype=cache_config.cache_dtype,
154 cache_config=cache_config,
155 )
156
157 assert resolved_dtype == "fp8_ds_mla"
158 assert torch_dtype is torch.uint8
159 assert cache_config.cache_dtype == "fp8_ds_mla"
160
161
162 def test_rocm_packed_kv_cache_rejects_unquantized_dtype() -> None:
163 from vllm.config import CacheConfig
164 from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype
165
166 cache_config = CacheConfig(cache_dtype="bfloat16")
167
168 with pytest.raises(ValueError, match="only supports fp8 kv-cache"):
169 _resolve_dsv4_kv_cache_dtype(
170 use_fp8_ds_mla_layout=True,
171 kv_cache_dtype=cache_config.cache_dtype,
172 cache_config=cache_config,
173 )
174
175
176 def test_vl_mapper_preserves_rocm_weight_mapping() -> None:
177 from vllm.models.deepseek_v4.amd.model import _make_deepseek_v4_weights_mapper
178 from vllm.models.deepseek_v4.common.vl_model import (
179 _make_deepseek_v4_vl_weights_mapper,
180 )
181
182 text_mapper = _make_deepseek_v4_weights_mapper("fp4", fuse_shared_experts=True)
183 mapper = _make_deepseek_v4_vl_weights_mapper(text_mapper, image_enabled=True)
184
185 assert mapper._map_name("layers.3.attn.wq_a.input_scale") == (
186 "language_model.model.layers.3.attn.wq_a.input_scale_2"
187 )
188 assert mapper._map_name("layers.3.ffn.shared_experts.w2.weight") == (
189 "language_model.model.layers.3.ffn.shared_experts.w2.weight"
190 )
191 assert mapper._map_name("head.weight") == "language_model.lm_head.weight"
192
193
194 def test_rocm_moe_wires_vision_routing_on_hash_and_regular_layers(
195 monkeypatch: pytest.MonkeyPatch,
196 ) -> None:
197 from vllm.models.deepseek_v4.amd import model as rocm_model
198 from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID
199
200 captured: list[dict] = []
201
202 class FakeGate(nn.Module):
203 def __init__(self, **kwargs) -> None:
204 super().__init__()
205
206 def fake_factory(**kwargs):
207 captured.append(kwargs)
208 return nn.Identity()
209
210 monkeypatch.setattr(rocm_model, "GateLinear", FakeGate)
211 monkeypatch.setattr(rocm_model, "FusedMoEFactory", fake_factory)
212 monkeypatch.setattr(rocm_model, "get_tensor_model_parallel_world_size", lambda: 1)
213 monkeypatch.setattr(rocm_model, "get_tensor_model_parallel_rank", lambda: 0)
214
215 config = SimpleNamespace(
216 hidden_size=16,
217 n_routed_experts=4,
218 num_experts_per_tok=2,
219 moe_intermediate_size=8,
220 swiglu_limit=None,
221 norm_topk_prob=True,
222 scoring_func="sqrtsoftplus",
223 num_hash_layers=1,
224 vocab_size=32,
225 topk_method="noaux_tc",
226 vision_n_layers=1,
227 n_shared_experts=None,
228 )
229 vllm_config = SimpleNamespace(
230 model_config=SimpleNamespace(hf_config=config),
231 quant_config=None,
232 parallel_config=ParallelConfig(),
233 )
234
235 hash_moe = rocm_model.DeepseekV4MoE(vllm_config, prefix="model.layers.0.ffn")
236 regular_moe = rocm_model.DeepseekV4MoE(vllm_config, prefix="model.layers.1.ffn")
237
238 assert hash_moe.gate.tid2eid is not None
239 assert regular_moe.gate.tid2eid is None
240 for moe, factory_kwargs in zip((hash_moe, regular_moe), captured, strict=True):
241 assert moe.gate.e_score_correction_bias is not None
242 assert moe.gate.bias_vl is not None
243 assert factory_kwargs["bias_vl"] is moe.gate.bias_vl
244 assert factory_kwargs["image_sentinel_lo"] == IMAGE_SENTINEL_BASE_ID
245
246
247 def test_rocm_fse_rejects_data_parallel_deployment(
248 monkeypatch: pytest.MonkeyPatch,
249 ) -> None:
250 from vllm.models.deepseek_v4.amd.model import _fuse_shared_experts_enabled
251
252 config = SimpleNamespace(n_shared_experts=1)
253
254 monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "1")
255 with pytest.raises(ValueError, match="data parallelism"):
256 _fuse_shared_experts_enabled(config, ParallelConfig(data_parallel_size=2))
257
258 # Expert parallelism already disables FSE: no raise, FSE stays off.
259 assert not _fuse_shared_experts_enabled(
260 config, ParallelConfig(data_parallel_size=2, enable_expert_parallel=True)
261 )
262 assert _fuse_shared_experts_enabled(config, ParallelConfig())
263
264
265 def test_rocm_mtp_forwards_input_ids_for_vision_routing(
266 monkeypatch: pytest.MonkeyPatch,
267 ) -> None:
268 from vllm.models.deepseek_v4.amd import mtp as rocm_mtp
269
270 hidden_size = 4
271 hc_mult = 2
272
273 class FakeNorm(nn.Module):
274 def __init__(self) -> None:
275 super().__init__()
276 self.weight = nn.Parameter(torch.ones(hidden_size))
277 self.variance_epsilon = 1e-6
278
279 class FakeMTPBlock(nn.Module):
280 use_fused_mhc = False
281
282 def __init__(self) -> None:
283 super().__init__()
284 self.seen_input_ids: torch.Tensor | None = None
285
286 def forward(
287 self,
288 *,
289 positions: torch.Tensor,
290 x: torch.Tensor,
291 input_ids: torch.Tensor | None,
292 ):
293 self.seen_input_ids = input_ids
294 return x, None, None, None
295
296 def passthrough_mtp_input(
297 inputs_embeds: torch.Tensor,
298 positions: torch.Tensor,
299 previous_hidden_states: torch.Tensor,
300 *args,
301 ) -> tuple[torch.Tensor, torch.Tensor]:
302 return inputs_embeds, previous_hidden_states
303
304 monkeypatch.setattr(
305 rocm_mtp, "_FUSED_MTP_INPUT_RMSNORM_KERNEL", passthrough_mtp_input
306 )
307
308 layer = object.__new__(rocm_mtp.DeepSeekV4MultiTokenPredictorLayer)
309 nn.Module.__init__(layer)
310 layer.config = SimpleNamespace(hidden_size=hidden_size)
311 layer.hc_mult = hc_mult
312 layer.enorm = FakeNorm()
313 layer.hnorm = FakeNorm()
314 layer.e_proj = nn.Identity()
315 layer.h_proj = nn.Identity()
316 layer.mtp_block = FakeMTPBlock()
317
318 input_ids = torch.tensor([11, 12])
319 positions = torch.tensor([3, 4])
320 inputs_embeds = torch.arange(8, dtype=torch.float32).view(2, hidden_size)
321 previous_hidden_states = torch.arange(16, dtype=torch.float32).view(2, -1)
322
323 output = layer(
324 input_ids,
325 positions,
326 previous_hidden_states,
327 inputs_embeds,
328 )
329
330 assert layer.mtp_block.seen_input_ids is input_ids
331 expected = previous_hidden_states.view(2, hc_mult, hidden_size)
332 expected = expected + inputs_embeds.unsqueeze(-2)
333 torch.testing.assert_close(output, expected.flatten(1))
334
335
336 def test_rocm_compute_logits_local_skips_gather() -> None:
337 from vllm.models.deepseek_v4.amd.model import DeepseekV4ForCausalLM
338
339 calls: list[tuple[nn.Module, torch.Tensor, bool]] = []
340
341 def logits_processor(
342 lm_head: nn.Module, hidden_states: torch.Tensor, *, skip_gather: bool = False
343 ) -> torch.Tensor:
344 calls.append((lm_head, hidden_states, skip_gather))
345 return hidden_states + 1
346
347 model = object.__new__(DeepseekV4ForCausalLM)
348 nn.Module.__init__(model)
349 model.lm_head = nn.Identity()
350 model.logits_processor = logits_processor
351 hidden_states = torch.tensor([4.0])
352
353 result = model.compute_logits_local(hidden_states)
354
355 assert torch.equal(result, torch.tensor([5.0]))
356 assert calls == [(model.lm_head, hidden_states, True)]
357
358
359 class _FakeLanguageModel(nn.Module):
360 finalizes_weights_during_load = False
361
362 def __init__(self) -> None:
363 super().__init__()
364 self.tensor_a = nn.Parameter(torch.zeros(1))
365 self.tensor_c = nn.Parameter(torch.zeros(1))
366 self.finalized_values: list[tuple[float, float]] = []
367
368 def process_weights_after_loading(self) -> None:
369 self.finalized_values.append((self.tensor_a.item(), self.tensor_c.item()))
370
371 def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
372 return hidden_states + 1
373
374
375 def test_vl_wrapper_streams_then_delegates_finalization() -> None:
376 from vllm.models.deepseek_v4.common.vl_model import (
377 DeepseekV4ForConditionalGeneration,
378 )
379
380 model = object.__new__(DeepseekV4ForConditionalGeneration)
381 nn.Module.__init__(model)
382 model.language_model = _FakeLanguageModel()
383 model.vision = nn.Module()
384 model.vision.tensor_b = nn.Parameter(torch.zeros(1))
385 model.hf_to_vllm_mapper = WeightsMapper()
386
387 def interleaved_weights():
388 yield "language_model.tensor_a", torch.tensor([1.0])
389 assert model.language_model.tensor_a.item() == 1.0
390 yield "vision.tensor_b", torch.tensor([2.0])
391 assert model.vision.tensor_b.item() == 2.0
392 yield "language_model.tensor_c", torch.tensor([3.0])
393
394 loaded = model.load_weights(interleaved_weights())
395
396 assert loaded == {
397 "language_model.tensor_a",
398 "vision.tensor_b",
399 "language_model.tensor_c",
400 }
401 assert model.language_model.finalized_values == []
402
403 model.process_weights_after_loading()
404
405 assert model.language_model.finalized_values == [(1.0, 3.0)]
406 assert torch.equal(
407 model.compute_logits_local(torch.tensor([4.0])), torch.tensor([5.0])
408 )
409 model.process_weights_after_loading()
410 assert model.language_model.finalized_values == [(1.0, 3.0)]
411
412
413 class _FakeFinalizingLanguageModel(_FakeLanguageModel):
414 finalizes_weights_during_load = True
415
416 def __init__(self) -> None:
417 super().__init__()
418 self.load_calls = 0
419
420 def load_weights(self, weights) -> set[str]:
421 self.load_calls += 1
422 loaded = set()
423 for name, value in weights:
424 getattr(self, name).data.copy_(value)
425 loaded.add(name)
426 self.process_weights_after_loading()
427 return loaded
428
429
430 def test_vl_wrapper_groups_child_that_finalizes_during_load() -> None:
431 from vllm.models.deepseek_v4.common.vl_model import (
432 DeepseekV4ForConditionalGeneration,
433 )
434
435 model = object.__new__(DeepseekV4ForConditionalGeneration)
436 nn.Module.__init__(model)
437 model.language_model = _FakeFinalizingLanguageModel()
438 model.vision = nn.Module()
439 model.vision.tensor_b = nn.Parameter(torch.zeros(1))
440 model.hf_to_vllm_mapper = WeightsMapper()
441
442 loaded = model.load_weights(
443 iter(
444 (
445 ("language_model.tensor_a", torch.tensor([1.0])),
446 ("vision.tensor_b", torch.tensor([2.0])),
447 ("language_model.tensor_c", torch.tensor([3.0])),
448 )
449 )
450 )
451
452 assert loaded == {
453 "language_model.tensor_a",
454 "vision.tensor_b",
455 "language_model.tensor_c",
456 }
457 assert model.language_model.load_calls == 1
458 assert model.language_model.finalized_values == [(1.0, 3.0)]
459
460 # The framework's later model-level hook must not double-finalize a child
461 # which already completed this work in load_weights.
462 model.process_weights_after_loading()
463 assert model.language_model.finalized_values == [(1.0, 3.0)]
464
465
466 def test_vl_wrapper_dummy_load_delegates_finalization() -> None:
467 from vllm.models.deepseek_v4.common.vl_model import (
468 DeepseekV4ForConditionalGeneration,
469 )
470
471 model = object.__new__(DeepseekV4ForConditionalGeneration)
472 nn.Module.__init__(model)
473 model.language_model = _FakeFinalizingLanguageModel()
474
475 # DummyModelLoader bypasses model.load_weights(), so no finalized marker
476 # exists and the framework-level hook must still delegate to the child.
477 model.process_weights_after_loading()
478 assert model.language_model.finalized_values == [(0.0, 0.0)]
479 model.process_weights_after_loading()
480 assert model.language_model.finalized_values == [(0.0, 0.0)]