[plugin][UT] add test cases(not model level) for ATOM OOT

Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

Author: zejunchen-zejun

tests/plugin/test_plugin_config_translation.py

@@ -0,0 +1,82 @@
+import pytest
+
+import atom.plugin.config as plugin_config
+
+
+class _Obj:
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+
+class _FakeConfig:
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+
+class _FakeCompilationConfig:
+    def __init__(self, level, use_cudagraph, cudagraph_mode):
+        self.level = level
+        self.use_cudagraph = use_cudagraph
+        self.cudagraph_mode = cudagraph_mode
+
+
+def _patch_atom_config_module(monkeypatch):
+    import atom.config as atom_config_module
+
+    monkeypatch.setattr(atom_config_module, "Config", _FakeConfig, raising=False)
+    monkeypatch.setattr(
+        atom_config_module, "CompilationConfig", _FakeCompilationConfig, raising=False
+    )
+
+
+def test_generate_from_vllm_translates_core_fields(monkeypatch):
+    _patch_atom_config_module(monkeypatch)
+    monkeypatch.setenv("ATOM_DISABLE_VLLM_PLUGIN_ATTENTION", "0")
+
+    vllm_cfg = _Obj(
+        model_config=_Obj(model="m1", max_model_len=4096),
+        scheduler_config=_Obj(max_num_batched_tokens=2048, max_num_seqs=8),
+        cache_config=_Obj(
+            gpu_memory_utilization=0.5,
+            block_size=16,
+            num_gpu_blocks=1024,
+            cache_dtype="auto",
+            enable_prefix_caching=True,
+        ),
+        parallel_config=_Obj(
+            rank=1, tensor_parallel_size=2, enable_expert_parallel=False
+        ),
+        compilation_config=_Obj(mode=3),
+        quant_config=_Obj(name="q"),
+    )
+
+    cfg = plugin_config._generate_atom_config_from_vllm_config(vllm_cfg)
+
+    assert cfg.model == "m1"
+    assert cfg.max_num_batched_tokens == 2048
+    assert cfg.max_num_seqs == 8
+    assert cfg.max_model_len == 4096
+    assert cfg.tensor_parallel_size == 2
+    assert cfg.enforce_eager is True
+    assert cfg.compilation_config.level == 3
+    assert cfg.plugin_config.is_plugin_mode is True
+    assert cfg.plugin_config.is_vllm is True
+    assert cfg.plugin_config.is_sglang is False
+    assert cfg.plugin_config.vllm_use_atom_attention is True
+
+
+def test_generate_atom_config_requires_plugin_mode(monkeypatch):
+    import atom.plugin.config as config_module
+    import atom.plugin as plugin_module
+    import atom.config as atom_config_module
+
+    monkeypatch.setattr(plugin_module, "is_vllm", lambda: False, raising=False)
+    monkeypatch.setattr(plugin_module, "is_sglang", lambda: False, raising=False)
+    monkeypatch.setattr(
+        atom_config_module, "set_current_atom_config", lambda _cfg: None, raising=False
+    )
+
+    with pytest.raises(ValueError, match="running in plugin mode"):
+        config_module.generate_atom_config_for_plugin_mode(config=None)

tests/plugin/test_plugin_env_flags.py

@@ -0,0 +1,56 @@
+import importlib
+import importlib.util
+import sys
+import types
+
+import pytest
+
+
+def test_disable_vllm_plugin_flag_disables_platform(monkeypatch):
+    # ATOM_DISABLE_VLLM_PLUGIN takes precedence:
+    # when it is 1, vLLM should not get ATOM platform/attention at all.
+    for disable_attention in ("0", "1"):
+        monkeypatch.setenv("ATOM_DISABLE_VLLM_PLUGIN", "1")
+        monkeypatch.setenv("ATOM_DISABLE_VLLM_PLUGIN_ATTENTION", disable_attention)
+
+        import atom.plugin.vllm.platform as platform_module
+        import atom.plugin.vllm.register as register_module
+
+        importlib.reload(platform_module)
+        importlib.reload(register_module)
+
+        assert platform_module.ATOMPlatform is None
+        assert register_module.register_platform() is None
+
+
+@pytest.mark.skipif(
+    importlib.util.find_spec("vllm") is None,
+    reason="vllm is not installed in current test environment",
+)
+def test_disable_vllm_plugin_attention_fallbacks_to_non_atom_backend(monkeypatch):
+    rocm_module = types.ModuleType("vllm.platforms.rocm")
+
+    class _RocmPlatform:
+        @classmethod
+        def get_attn_backend_cls(cls, selected_backend, attn_selector_config):
+            return "vllm.default.backend"
+
+    rocm_module.RocmPlatform = _RocmPlatform
+
+    monkeypatch.setitem(sys.modules, "vllm", types.ModuleType("vllm"))
+    monkeypatch.setitem(
+        sys.modules, "vllm.platforms", types.ModuleType("vllm.platforms")
+    )
+    monkeypatch.setitem(sys.modules, "vllm.platforms.rocm", rocm_module)
+    monkeypatch.setenv("ATOM_DISABLE_VLLM_PLUGIN", "0")
+    monkeypatch.setenv("ATOM_DISABLE_VLLM_PLUGIN_ATTENTION", "1")
+
+    import atom.plugin.vllm.platform as platform_module
+
+    importlib.reload(platform_module)
+
+    result = platform_module.ATOMPlatform.get_attn_backend_cls(
+        selected_backend="x",
+        attn_selector_config=types.SimpleNamespace(use_mla=True),
+    )
+    assert result == "vllm.default.backend"

tests/plugin/test_plugin_mode_status.py

@@ -0,0 +1,36 @@
+import pytest
+
+from atom.plugin import prepare as plugin_prepare
+
+
+@pytest.fixture(autouse=True)
+def _reset_framework_state():
+    # Autouse fixture: pytest runs this before/after every test.
+    plugin_prepare._set_framework_backbone("atom")
+    yield
+    plugin_prepare._set_framework_backbone("atom")
+
+
+def test_default_mode_is_server_mode():
+    assert plugin_prepare.is_plugin_mode() is False
+    assert plugin_prepare.is_vllm() is False
+    assert plugin_prepare.is_sglang() is False
+
+
+def test_set_framework_to_vllm():
+    plugin_prepare._set_framework_backbone("vllm")
+    assert plugin_prepare.is_plugin_mode() is True
+    assert plugin_prepare.is_vllm() is True
+    assert plugin_prepare.is_sglang() is False
+
+
+def test_set_framework_to_sgl_alias():
+    plugin_prepare._set_framework_backbone("sgl")
+    assert plugin_prepare.is_plugin_mode() is True
+    assert plugin_prepare.is_vllm() is False
+    assert plugin_prepare.is_sglang() is True
+
+
+def test_set_framework_unsupported_raises():
+    with pytest.raises(ValueError, match="Unsupported framework"):
+        plugin_prepare._set_framework_backbone("tensorflow")

tests/plugin/test_plugin_registries.py

@@ -0,0 +1,57 @@
+import sys
+import types
+import importlib
+import importlib.util
+
+import pytest
+
+from atom.plugin import prepare as plugin_prepare
+import atom.plugin.vllm.register as vllm_register
+
+
+@pytest.fixture(autouse=True)
+def _reset_framework_state():
+    plugin_prepare._set_framework_backbone("atom")
+    yield
+    plugin_prepare._set_framework_backbone("atom")
+
+
+@pytest.mark.skipif(
+    importlib.util.find_spec("vllm") is None,
+    reason="vllm is not installed in current test environment",
+)
+def test_register_platform_returns_oot_platform(monkeypatch):
+    rocm_module = types.ModuleType("vllm.platforms.rocm")
+
+    class _RocmPlatform:
+        pass
+
+    rocm_module.RocmPlatform = _RocmPlatform
+    vllm_platforms = types.ModuleType("vllm.platforms")
+    vllm_platforms.current_platform = None
+
+    monkeypatch.setitem(sys.modules, "vllm", types.ModuleType("vllm"))
+    monkeypatch.setitem(sys.modules, "vllm.platforms", vllm_platforms)
+    monkeypatch.setitem(sys.modules, "vllm.platforms.rocm", rocm_module)
+
+    monkeypatch.setenv("ATOM_DISABLE_VLLM_PLUGIN", "0")
+    monkeypatch.setenv("ATOM_DISABLE_VLLM_PLUGIN_ATTENTION", "0")
+
+    import atom.plugin.vllm.platform as platform_module
+
+    importlib.reload(platform_module)
+    importlib.reload(vllm_register)
+
+    platform_path = vllm_register.register_platform()
+    module_name, class_name = platform_path.rsplit(".", 1)
+    vllm_platforms.current_platform = getattr(
+        importlib.import_module(module_name), class_name
+    )
+
+    # get current platform from vllm side and validate it is ATOM platform.
+    assert vllm_platforms.current_platform is platform_module.ATOMPlatform
+
+
+def test_register_platform_can_be_disabled(monkeypatch):
+    monkeypatch.setattr(vllm_register, "disable_vllm_plugin", True, raising=False)
+    assert vllm_register.register_platform() is None

tests/plugin/test_plugin_unsupported_models.py

@@ -0,0 +1,23 @@
+import importlib.util
+import importlib
+import sys
+import types
+
+import pytest
+
+
+# FIXME: remove it later when enabling fallback for unsupported models
+@pytest.mark.skipif(
+    importlib.util.find_spec("vllm") is None,
+    reason="vllm is not installed in current test environment",
+)
+def test_vllm_wrapper_rejects_unsupported_model_arch(monkeypatch):
+    # Avoid importing deep model-loader dependencies during test collection/import.
+    fake_loader = types.ModuleType("atom.model_loader.loader")
+    fake_loader.load_model_in_plugin_mode = lambda **kwargs: set()
+    monkeypatch.setitem(sys.modules, "atom.model_loader.loader", fake_loader)
+
+    model_wrapper = importlib.import_module("atom.plugin.vllm.model_wrapper")
+
+    with pytest.raises(ValueError, match="not supported by ATOM OOT backend"):
+        model_wrapper._get_atom_model_cls("UnknownModelForCausalLM")

tests/plugin/test_plugin_vllm_import_paths.py

@@ -0,0 +1,85 @@
+import importlib.util
+
+import pytest
+
+
+@pytest.mark.skipif(
+    importlib.util.find_spec("vllm") is None,
+    reason="vllm is not installed in current test environment",
+)
+def test_vllm_import_paths_guardrail():
+    """Guardrail for OOT vLLM import paths used by ATOM plugin mode."""
+    # attention.py / paged_attention.py (new path with legacy fallback)
+    try:
+        from vllm.attention.layer import Attention, MLAAttention, AttentionType
+    except ImportError:
+        from vllm.model_executor.layers.attention import Attention, MLAAttention
+        from vllm.v1.attention.backend import AttentionType
+
+    # attention.py
+    from vllm.config import (
+        VllmConfig,
+        get_current_vllm_config,
+        get_layers_from_vllm_config,
+    )
+    from vllm.model_executor.layers.attention.mla_attention import (
+        MLACommonMetadataBuilder,
+        QueryLenSupport,
+    )
+    from vllm.utils.math_utils import cdiv, round_down
+    from vllm.v1.attention.backend import AttentionCGSupport, AttentionMetadataBuilder
+    from vllm.v1.attention.backends.utils import (
+        get_dcp_local_seq_lens,
+        split_decodes_and_prefills,
+        split_decodes_prefills_and_extends,
+    )
+    from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
+    from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
+
+    # model_wrapper.py (core vLLM model interfaces)
+    from vllm.model_executor.models.interfaces import SupportsPP, SupportsQuant
+    from vllm.model_executor.models.interfaces_base import (
+        VllmModel,
+        VllmModelForTextGeneration,
+    )
+    from vllm.model_executor.models.registry import ModelRegistry
+    from vllm.sequence import IntermediateTensors
+
+    # attention_mla.py / platform.py / register.py
+    from vllm import _custom_ops
+    from vllm.distributed.parallel_state import get_dcp_group
+    from vllm.platforms import current_platform
+    from vllm.platforms.rocm import RocmPlatform
+
+    assert all(
+        obj is not None
+        for obj in [
+            Attention,
+            MLAAttention,
+            AttentionType,
+            QueryLenSupport,
+            MLACommonMetadataBuilder,
+            cdiv,
+            round_down,
+            AttentionCGSupport,
+            AttentionMetadataBuilder,
+            get_dcp_local_seq_lens,
+            split_decodes_and_prefills,
+            split_decodes_prefills_and_extends,
+            cp_lse_ag_out_rs,
+            merge_attn_states,
+            VllmConfig,
+            get_current_vllm_config,
+            get_layers_from_vllm_config,
+            SupportsPP,
+            SupportsQuant,
+            VllmModel,
+            VllmModelForTextGeneration,
+            ModelRegistry,
+            IntermediateTensors,
+            _custom_ops,
+            get_dcp_group,
+            current_platform,
+            RocmPlatform,
+        ]
+    )