main
py 78 lines 2.71 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from unittest.mock import Mock, patch
5
6
7 def test_qwen3_5_lm_head_receives_quant_config():
8 from vllm.model_executor.models.qwen3_5 import Qwen3_5ForCausalLMBase
9
10 mock_quant_config = Mock()
11
12 mock_hf_config = Mock()
13 mock_hf_config.tie_word_embeddings = False
14 mock_hf_config.vocab_size = 128
15 mock_hf_config.hidden_size = 64
16
17 mock_vllm_config = Mock()
18 mock_vllm_config.model_config.hf_text_config = mock_hf_config
19 mock_vllm_config.cache_config.mamba_cache_mode = "align"
20 mock_vllm_config.scheduler_config = Mock()
21 mock_vllm_config.quant_config = mock_quant_config
22 mock_vllm_config.lora_config = None
23
24 mock_pp_group = Mock()
25 mock_pp_group.is_last_rank = True
26
27 with (
28 patch("vllm.model_executor.models.qwen3_5.Qwen3_5Model") as MockModel,
29 patch("vllm.model_executor.models.qwen3_5.ParallelLMHead") as MockLMHead,
30 patch("vllm.model_executor.models.qwen3_5.LogitsProcessor"),
31 patch(
32 "vllm.model_executor.models.qwen3_5.get_pp_group",
33 return_value=mock_pp_group,
34 ),
35 ):
36 MockModel.return_value.make_empty_intermediate_tensors = Mock()
37
38 Qwen3_5ForCausalLMBase(vllm_config=mock_vllm_config)
39
40 MockLMHead.assert_called_once()
41 call_kwargs = MockLMHead.call_args.kwargs
42 assert call_kwargs["quant_config"] is mock_quant_config
43
44
45 def test_qwen3_5_mtp_lm_head_receives_quant_config():
46 from vllm.config import CompilationMode
47 from vllm.model_executor.models.qwen3_5_mtp import Qwen3_5MTP
48
49 mock_quant_config = Mock()
50
51 mock_hf_config = Mock()
52 mock_hf_config.tie_word_embeddings = False
53 mock_hf_config.vocab_size = 128
54 mock_hf_config.hidden_size = 64
55
56 mock_vllm_config = Mock()
57 mock_vllm_config.model_config.hf_text_config = mock_hf_config
58 mock_vllm_config.cache_config.mamba_cache_mode = "align"
59 mock_vllm_config.compilation_config.mode = CompilationMode.NONE
60 mock_vllm_config.quant_config = mock_quant_config
61
62 mock_pp_group = Mock()
63 mock_pp_group.is_last_rank = True
64
65 with (
66 patch("vllm.model_executor.models.qwen3_5_mtp.Qwen3_5MultiTokenPredictor"),
67 patch("vllm.model_executor.models.qwen3_5_mtp.ParallelLMHead") as MockLMHead,
68 patch("vllm.model_executor.models.qwen3_5_mtp.LogitsProcessor"),
69 patch(
70 "vllm.model_executor.models.qwen3_5_mtp.get_pp_group",
71 return_value=mock_pp_group,
72 ),
73 ):
74 Qwen3_5MTP(vllm_config=mock_vllm_config)
75
76 MockLMHead.assert_called_once()
77 call_kwargs = MockLMHead.call_args.kwargs
78 assert call_kwargs["quant_config"] is mock_quant_config