A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 09:45 UTC
[Bugfix] Fix test_eagle test (#18223)
Signed-off-by: Lucia Fang <fanglu@fb.com>
Lucia Fang committed
May 15, 2025 at 15:59 UTC
8795eb9975561e19fe642b39b42e0c18280ac796
1 file changed
+8
-2
tests/v1/spec_decode/test_eagle.py
+8
-2
@@ -115,14 +115,15 @@ def test_prepare_inputs():
115
("eagle3", lambda k: _create_proposer("eagle3", k), eagle3_dir,
116
('model', 'embed_tokens')),
117
])
118
+@mock.patch('vllm.v1.spec_decode.eagle.get_pp_group')
119
@mock.patch('vllm.v1.spec_decode.eagle.get_layers_from_vllm_config')
120
@mock.patch('vllm.v1.spec_decode.eagle.ModelRegistry')
121
@mock.patch('vllm.v1.spec_decode.eagle.get_model_loader')
122
@mock.patch('vllm.v1.spec_decode.eagle.set_default_torch_dtype')
123
@mock.patch('vllm.v1.spec_decode.eagle.set_current_vllm_config')
124
def test_load_model(mock_set_config, mock_set_dtype, mock_get_loader,
124
- mock_registry, mock_get_layers, method, proposer_helper,
125
- draft_model_dir, target_attribute_path):
125
+ mock_registry, mock_get_layers, mock_get_pp_group, method,
126
+ proposer_helper, draft_model_dir, target_attribute_path):
127
128
# Setup mock for model class
129
mock_model_cls = mock.MagicMock()
@@ -158,6 +159,11 @@ def test_load_model(mock_set_config, mock_set_dtype, mock_get_loader,
159
# Make mock_get_layers return different values for each call
160
mock_get_layers.side_effect = [target_attn_layers, all_attn_layers]
161
162
+ # Setup mock for pp group to return the appropriate value for world size
163
+ mock_pp_group = mock.MagicMock()
164
+ mock_pp_group.world_size = 2 if method == "eagle" else 1
165
+ mock_get_pp_group.return_value = mock_pp_group
166
+
167
# Setup model loader mock
168
mock_loader = mock.MagicMock()
169
mock_get_loader.return_value = mock_loader