test_lora_utils.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project


from vllm.lora.utils import is_in_target_modules, is_supported_lora_module


class TestIsSupportedLoraModule:
    """Tests for is_supported_lora_module (model-definition check)."""

    def test_suffix_match(self):
        assert is_supported_lora_module(
            "model.layers.0.self_attn.o_proj", ["o_proj", "q_proj"]
        )

    def test_no_match(self):
        assert not is_supported_lora_module(
            "model.layers.0.self_attn.o_proj", ["q_proj", "k_proj"]
        )

    def test_exact_match(self):
        assert is_supported_lora_module("o_proj", ["o_proj"])

    def test_regex_suffix_matching(self):
        """Regex anchors to end — partial suffix should not match."""
        assert not is_supported_lora_module("model.layers.0.self_attn.o_proj", ["proj"])

    def test_empty_supported_modules(self):
        assert not is_supported_lora_module("model.layers.0.self_attn.o_proj", [])

    def test_multiple_supported_modules(self):
        supported = ["q_proj", "k_proj", "v_proj", "o_proj"]
        assert is_supported_lora_module("model.layers.0.self_attn.v_proj", supported)
        assert not is_supported_lora_module("model.layers.0.mlp.gate_proj", supported)


class TestIsInTargetModules:
    """Tests for is_in_target_modules (deployment-time filter)."""

    def test_none_allows_all(self):
        assert is_in_target_modules("model.layers.0.self_attn.o_proj", None)

    def test_suffix_in_target(self):
        assert is_in_target_modules(
            "model.layers.0.self_attn.o_proj", ["o_proj", "q_proj"]
        )

    def test_suffix_not_in_target(self):
        assert not is_in_target_modules(
            "model.layers.0.self_attn.o_proj", ["q_proj", "k_proj"]
        )

    def test_empty_target_modules(self):
        assert not is_in_target_modules("model.layers.0.self_attn.o_proj", [])

    def test_exact_name_match(self):
        assert is_in_target_modules("dense1", ["dense1", "dense2"])

    def test_exact_name_no_match(self):
        assert not is_in_target_modules("dense3", ["dense1", "dense2"])

    def test_packed_parent_matches_child_target_modules(self):
        assert is_in_target_modules(
            "model.layers.0.mlp.gate_up_proj",
            ["gate_proj", "up_proj"],
            {"gate_up_proj": ["gate_proj", "up_proj"]},
        )

    def test_packed_child_matches_parent_target_modules(self):
        assert is_in_target_modules(
            "model.layers.0.mlp.gate_proj",
            ["gate_up_proj"],
            {"gate_up_proj": ["gate_proj", "up_proj"]},
        )

    def test_fused_parent_matches_child_target_modules(self):
        assert is_in_target_modules(
            "model.layers.0.self_attn.fused_qkv_a_proj",
            ["q_a_proj", "kv_a_proj_with_mqa"],
            {"fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"]},
        )