488 lines
17 KiB
Python
488 lines
17 KiB
Python
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# This file is a part of the vllm-ascend project.
|
|
#
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from vllm_ascend.ops.flashcomm2_oshard_manager import (
|
|
Flashcomm2OShardManager,
|
|
flashcomm2_oshard_manager,
|
|
)
|
|
|
|
|
|
class TestFlashcomm2OShardManager:
|
|
@pytest.fixture
|
|
def manager(self):
|
|
return Flashcomm2OShardManager()
|
|
|
|
def test_init(self, manager):
|
|
assert manager._shard_layers == {}
|
|
assert isinstance(manager._shard_layers, dict)
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.flashcomm2_enable")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.o_shard_enable")
|
|
def test_flashcomm2_oshard_enable_both_enabled(
|
|
self,
|
|
mock_o_shard_enable,
|
|
mock_flashcomm2_enable,
|
|
manager,
|
|
):
|
|
mock_flashcomm2_enable.return_value = True
|
|
mock_o_shard_enable.return_value = True
|
|
|
|
result = manager.flashcomm2_oshard_enable()
|
|
|
|
assert result is True
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.flashcomm2_enable")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.o_shard_enable")
|
|
def test_flashcomm2_oshard_enable_flashcomm2_disabled(
|
|
self,
|
|
mock_o_shard_enable,
|
|
mock_flashcomm2_enable,
|
|
manager,
|
|
):
|
|
mock_flashcomm2_enable.return_value = False
|
|
mock_o_shard_enable.return_value = True
|
|
|
|
result = manager.flashcomm2_oshard_enable()
|
|
|
|
assert result is False
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.flashcomm2_enable")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.o_shard_enable")
|
|
def test_flashcomm2_oshard_enable_o_shard_disabled(
|
|
self,
|
|
mock_o_shard_enable,
|
|
mock_flashcomm2_enable,
|
|
manager,
|
|
):
|
|
mock_flashcomm2_enable.return_value = True
|
|
mock_o_shard_enable.return_value = False
|
|
|
|
result = manager.flashcomm2_oshard_enable()
|
|
|
|
assert result is False
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.flashcomm2_enable")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.o_shard_enable")
|
|
def test_flashcomm2_oshard_enable_both_disabled(
|
|
self,
|
|
mock_o_shard_enable,
|
|
mock_flashcomm2_enable,
|
|
manager,
|
|
):
|
|
mock_flashcomm2_enable.return_value = False
|
|
mock_o_shard_enable.return_value = False
|
|
|
|
result = manager.flashcomm2_oshard_enable()
|
|
|
|
assert result is False
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.register_layer_to_shard_weight_series")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.get_shard_weight_group")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.is_hidden_layer")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.extract_layer_index")
|
|
def test_register_layer_hidden_layer(
|
|
self,
|
|
mock_extract_index,
|
|
mock_is_hidden,
|
|
mock_get_group,
|
|
mock_register,
|
|
manager,
|
|
):
|
|
mock_is_hidden.return_value = True
|
|
mock_extract_index.return_value = 5
|
|
mock_group = MagicMock()
|
|
mock_get_group.return_value = mock_group
|
|
|
|
layer = MagicMock()
|
|
layer.prefix = "model.layers.5.self_attn.o_proj"
|
|
|
|
manager.register_layer(layer, prefetch_step=2)
|
|
|
|
assert 5 in manager._shard_layers
|
|
assert manager._shard_layers[5] is layer
|
|
mock_register.assert_called_once_with(
|
|
series_name="o_proj",
|
|
group=mock_group,
|
|
layer=layer,
|
|
prefetch_step=2,
|
|
)
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.register_layer_to_shard_weight_series")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.get_shard_weight_group")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.is_hidden_layer")
|
|
def test_register_layer_non_hidden_layer(
|
|
self,
|
|
mock_is_hidden,
|
|
mock_get_group,
|
|
mock_register,
|
|
manager,
|
|
):
|
|
mock_is_hidden.return_value = False
|
|
|
|
layer = MagicMock()
|
|
layer.prefix = "model.layers.100.self_attn.o_proj"
|
|
|
|
manager.register_layer(layer)
|
|
|
|
assert len(manager._shard_layers) == 0
|
|
mock_register.assert_not_called()
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.register_layer_to_shard_weight_series")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.get_shard_weight_group")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.is_hidden_layer")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.extract_layer_index")
|
|
def test_register_layer_default_prefetch_step(
|
|
self,
|
|
mock_extract_index,
|
|
mock_is_hidden,
|
|
mock_get_group,
|
|
mock_register,
|
|
manager,
|
|
):
|
|
mock_is_hidden.return_value = True
|
|
mock_extract_index.return_value = 0
|
|
mock_group = MagicMock()
|
|
mock_get_group.return_value = mock_group
|
|
|
|
layer = MagicMock()
|
|
layer.prefix = "model.layers.0.self_attn.o_proj"
|
|
|
|
manager.register_layer(layer)
|
|
|
|
mock_register.assert_called_once_with(
|
|
series_name="o_proj",
|
|
group=mock_group,
|
|
layer=layer,
|
|
prefetch_step=1,
|
|
)
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.register_layer_to_shard_weight_series")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.get_shard_weight_group")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.is_hidden_layer")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.extract_layer_index")
|
|
def test_register_layer_overwrites_existing_layer_with_same_index(
|
|
self,
|
|
mock_extract_index,
|
|
mock_is_hidden,
|
|
mock_get_group,
|
|
mock_register,
|
|
manager,
|
|
):
|
|
mock_is_hidden.return_value = True
|
|
mock_extract_index.return_value = 5
|
|
mock_group = MagicMock()
|
|
mock_get_group.return_value = mock_group
|
|
|
|
first_layer = MagicMock()
|
|
first_layer.prefix = "model.layers.5.self_attn.o_proj"
|
|
second_layer = MagicMock()
|
|
second_layer.prefix = "model.layers.5.self_attn.o_proj"
|
|
|
|
manager.register_layer(first_layer)
|
|
manager.register_layer(second_layer)
|
|
|
|
assert len(manager._shard_layers) == 1
|
|
assert manager.get_layer(5) is second_layer
|
|
assert mock_register.call_count == 2
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.register_layer_to_shard_weight_series")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.get_shard_weight_group")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.is_hidden_layer")
|
|
def test_register_layer_missing_prefix_raises(
|
|
self,
|
|
mock_is_hidden,
|
|
mock_get_group,
|
|
mock_register,
|
|
manager,
|
|
):
|
|
mock_is_hidden.return_value = True
|
|
layer = MagicMock(spec=[])
|
|
|
|
with pytest.raises(AttributeError):
|
|
manager.register_layer(layer)
|
|
|
|
assert manager._shard_layers == {}
|
|
mock_get_group.assert_not_called()
|
|
mock_register.assert_not_called()
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.register_layer_to_shard_weight_series")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.get_shard_weight_group")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.is_hidden_layer")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.extract_layer_index")
|
|
def test_register_layer_extract_layer_index_failure_propagates(
|
|
self,
|
|
mock_extract_index,
|
|
mock_is_hidden,
|
|
mock_get_group,
|
|
mock_register,
|
|
manager,
|
|
):
|
|
mock_is_hidden.return_value = True
|
|
mock_extract_index.side_effect = ValueError("invalid layer prefix")
|
|
|
|
layer = MagicMock()
|
|
layer.prefix = "invalid-prefix"
|
|
|
|
with pytest.raises(ValueError, match="invalid layer prefix"):
|
|
manager.register_layer(layer)
|
|
|
|
assert manager._shard_layers == {}
|
|
mock_get_group.assert_not_called()
|
|
mock_register.assert_not_called()
|
|
|
|
def test_get_layer_existing(self, manager):
|
|
layer = MagicMock()
|
|
manager._shard_layers[5] = layer
|
|
|
|
result = manager.get_layer(5)
|
|
|
|
assert result is layer
|
|
|
|
def test_get_layer_non_existing(self, manager):
|
|
result = manager.get_layer(999)
|
|
|
|
assert result is None
|
|
|
|
def test_get_layer_empty_dict(self, manager):
|
|
result = manager.get_layer(0)
|
|
|
|
assert result is None
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.reach_layer_for_shard_weight_series")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.is_hidden_layer")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.extract_layer_index")
|
|
def test_trigger_broadcast_for_layer_success(
|
|
self,
|
|
mock_extract_index,
|
|
mock_is_hidden,
|
|
mock_reach_layer,
|
|
manager,
|
|
):
|
|
mock_extract_index.return_value = 3
|
|
mock_is_hidden.return_value = True
|
|
|
|
layer = MagicMock()
|
|
manager._shard_layers[3] = layer
|
|
|
|
manager.trigger_broadcast_for_layer("model.layers.3.self_attn.o_proj")
|
|
|
|
mock_reach_layer.assert_called_once_with(layer)
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.reach_layer_for_shard_weight_series")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.is_hidden_layer")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.extract_layer_index")
|
|
def test_trigger_broadcast_for_layer_not_hidden(
|
|
self,
|
|
mock_extract_index,
|
|
mock_is_hidden,
|
|
mock_reach_layer,
|
|
manager,
|
|
):
|
|
mock_extract_index.return_value = 3
|
|
mock_is_hidden.return_value = False
|
|
|
|
layer = MagicMock()
|
|
manager._shard_layers[3] = layer
|
|
|
|
manager.trigger_broadcast_for_layer("model.layers.3.self_attn.o_proj")
|
|
|
|
mock_reach_layer.assert_not_called()
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.reach_layer_for_shard_weight_series")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.extract_layer_index")
|
|
def test_trigger_broadcast_for_layer_not_registered(
|
|
self,
|
|
mock_extract_index,
|
|
mock_reach_layer,
|
|
manager,
|
|
):
|
|
mock_extract_index.return_value = 999
|
|
|
|
manager.trigger_broadcast_for_layer("model.layers.999.self_attn.o_proj")
|
|
|
|
mock_reach_layer.assert_not_called()
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.reach_layer_for_shard_weight_series")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.is_hidden_layer")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.extract_layer_index")
|
|
def test_trigger_broadcast_for_layer_not_registered_short_circuits_hidden_check(
|
|
self,
|
|
mock_extract_index,
|
|
mock_is_hidden,
|
|
mock_reach_layer,
|
|
manager,
|
|
):
|
|
mock_extract_index.return_value = 999
|
|
|
|
manager.trigger_broadcast_for_layer("model.layers.999.self_attn.o_proj")
|
|
|
|
mock_is_hidden.assert_not_called()
|
|
mock_reach_layer.assert_not_called()
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.reach_layer_for_shard_weight_series")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.is_hidden_layer")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.extract_layer_index")
|
|
def test_trigger_broadcast_for_layer_empty_manager(
|
|
self,
|
|
mock_extract_index,
|
|
mock_is_hidden,
|
|
mock_reach_layer,
|
|
manager,
|
|
):
|
|
mock_extract_index.return_value = 0
|
|
|
|
manager.trigger_broadcast_for_layer("model.layers.0.self_attn.o_proj")
|
|
|
|
mock_is_hidden.assert_not_called()
|
|
mock_reach_layer.assert_not_called()
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.reach_layer_for_shard_weight_series")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.extract_layer_index")
|
|
def test_trigger_broadcast_for_layer_extract_layer_index_failure_propagates(
|
|
self,
|
|
mock_extract_index,
|
|
mock_reach_layer,
|
|
manager,
|
|
):
|
|
mock_extract_index.side_effect = ValueError("invalid layer prefix")
|
|
|
|
with pytest.raises(ValueError, match="invalid layer prefix"):
|
|
manager.trigger_broadcast_for_layer("invalid-prefix")
|
|
|
|
mock_reach_layer.assert_not_called()
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.post_process_after_loading_for_shard_weight_series")
|
|
def test_post_process_after_loading_with_layers(self, mock_post_process, manager):
|
|
layer1 = MagicMock()
|
|
layer2 = MagicMock()
|
|
manager._shard_layers[0] = layer1
|
|
manager._shard_layers[1] = layer2
|
|
|
|
manager.post_process_after_loading()
|
|
|
|
mock_post_process.assert_called_once()
|
|
called_layer = mock_post_process.call_args[0][0]
|
|
assert called_layer in [layer1, layer2]
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.post_process_after_loading_for_shard_weight_series")
|
|
def test_post_process_after_loading_uses_first_registered_layer(self, mock_post_process, manager):
|
|
first_layer = MagicMock()
|
|
second_layer = MagicMock()
|
|
manager._shard_layers[1] = first_layer
|
|
manager._shard_layers[2] = second_layer
|
|
|
|
manager.post_process_after_loading()
|
|
|
|
mock_post_process.assert_called_once_with(first_layer)
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.post_process_after_loading_for_shard_weight_series")
|
|
def test_post_process_after_loading_empty(self, mock_post_process, manager):
|
|
manager.post_process_after_loading()
|
|
|
|
mock_post_process.assert_not_called()
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.post_process_after_loading_for_shard_weight_series")
|
|
def test_post_process_after_loading_single_layer(self, mock_post_process, manager):
|
|
layer = MagicMock()
|
|
manager._shard_layers[5] = layer
|
|
|
|
manager.post_process_after_loading()
|
|
|
|
mock_post_process.assert_called_once_with(layer)
|
|
|
|
|
|
class TestGlobalInstance:
|
|
def test_global_instance_exists(self):
|
|
assert flashcomm2_oshard_manager is not None
|
|
assert isinstance(flashcomm2_oshard_manager, Flashcomm2OShardManager)
|
|
|
|
def test_global_instance_has_shard_layers(self):
|
|
assert hasattr(flashcomm2_oshard_manager, "_shard_layers")
|
|
assert isinstance(flashcomm2_oshard_manager._shard_layers, dict)
|
|
|
|
|
|
class TestIntegration:
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.register_layer_to_shard_weight_series")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.get_shard_weight_group")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.is_hidden_layer")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.extract_layer_index")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.reach_layer_for_shard_weight_series")
|
|
def test_full_workflow(
|
|
self,
|
|
mock_reach_layer,
|
|
mock_extract_index,
|
|
mock_is_hidden,
|
|
mock_get_group,
|
|
mock_register,
|
|
):
|
|
manager = Flashcomm2OShardManager()
|
|
|
|
mock_is_hidden.return_value = True
|
|
mock_extract_index.side_effect = lambda x: int(x.split(".")[2])
|
|
mock_group = MagicMock()
|
|
mock_get_group.return_value = mock_group
|
|
|
|
layer0 = MagicMock()
|
|
layer0.prefix = "model.layers.0.self_attn.o_proj"
|
|
|
|
layer1 = MagicMock()
|
|
layer1.prefix = "model.layers.1.self_attn.o_proj"
|
|
|
|
manager.register_layer(layer0)
|
|
manager.register_layer(layer1)
|
|
|
|
assert len(manager._shard_layers) == 2
|
|
assert manager.get_layer(0) is layer0
|
|
assert manager.get_layer(1) is layer1
|
|
|
|
manager.trigger_broadcast_for_layer("model.layers.0.self_attn.o_proj")
|
|
mock_reach_layer.assert_called_with(layer0)
|
|
|
|
manager.trigger_broadcast_for_layer("model.layers.1.self_attn.o_proj")
|
|
mock_reach_layer.assert_called_with(layer1)
|
|
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.register_layer_to_shard_weight_series")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.get_shard_weight_group")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.is_hidden_layer")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.extract_layer_index")
|
|
@patch("vllm_ascend.ops.flashcomm2_oshard_manager.post_process_after_loading_for_shard_weight_series")
|
|
def test_register_and_post_process_workflow(
|
|
self,
|
|
mock_post_process,
|
|
mock_extract_index,
|
|
mock_is_hidden,
|
|
mock_get_group,
|
|
mock_register,
|
|
):
|
|
manager = Flashcomm2OShardManager()
|
|
|
|
mock_is_hidden.return_value = True
|
|
mock_extract_index.return_value = 0
|
|
mock_group = MagicMock()
|
|
mock_get_group.return_value = mock_group
|
|
|
|
layer = MagicMock()
|
|
layer.prefix = "model.layers.0.self_attn.o_proj"
|
|
|
|
manager.register_layer(layer)
|
|
|
|
manager.post_process_after_loading()
|
|
|
|
mock_post_process.assert_called_once()
|