From 4147235c764b6e8fcdc1f629dcca8f6c5d83577e Mon Sep 17 00:00:00 2001 From: SieDeta Date: Tue, 9 Jun 2026 19:07:13 +0700 Subject: [PATCH 1/4] test: marker taxonomy + skip guards + categorize (PR1) Introduce the test-suite categorization foundation from the test-suite refactor plan (PR #1 of 6). - pyproject.toml: register the full marker set (cpu, gpu, multi_gpu, e2e, clickhouse, vllm, hf, ring_native, slow, manual, numeric); default collection excludes `manual` and skips tests/tools, tests/ring, .venv, integration, build via norecursedirs. - tests/_requirements.py: CPU-importable skip guards that fail closed with precise reasons (require_cuda, require_gpus, require_clickhouse, require_vllm, require_model_cache, require_nvcc). - Mark every test: pure-CPU contract tests -> `cpu`; GPU/E2E suites -> explicit gpu/vllm/clickhouse/e2e/hf marks. test_gpt2_parity pulls real gpt2 weights -> `hf` + require_model_cache, not cpu. - Split test_producer_chunked_schema: op-registration tests stay `cpu`, CUDA device-smoke tests become `gpu`. - test_per_hook_isolation: source-patch unit classes -> `cpu`; the slow GPU sweep -> `gpu` (keeps `slow`). - Relocate manual shell wrappers to tests/tools/ (+ README); fix their repo-root path computations and cross-references for the new depth. Default CPU command: python -m pytest -m "not gpu and not e2e and not manual" -q Verified: markers register without warnings; the CPU selection deselects 32 GPU/E2E tests. Pre-existing (reproduced on base main, unrelated to this PR): monitoring/ring_transport.py references HOOK_TYPE_ROUTER_LOGITS (defined only in csrc enum, missing on the Python side), and the vendored transformers submodule is unbuilt -- these block a fully green CPU run and must be fixed separately. Co-Authored-By: Claude Opus 4.8 --- pyproject.toml | 15 ++ tests/_requirements.py | 128 ++++++++++++++++++ tests/test_adapter_protocol.py | 4 + tests/test_config.py | 2 + tests/test_e2e_correctness_vs_hf.py | 7 + tests/test_gpt2_parity.py | 5 + tests/test_hf_eos_strip.py | 2 + tests/test_hook_spec_flags.py | 3 + tests/test_moe_v1_routing_hooks.py | 4 + tests/test_no_graph_breaks.py | 2 + tests/test_per_hook_isolation.py | 4 + tests/test_producer_chunked_schema.py | 29 ++-- ...est_ring_transport_no_framework_strings.py | 2 + tests/test_tp_shapes.py | 2 + tests/test_vllm_identical.py | 7 + tests/test_vllm_rowcnt.py | 7 + tests/tools/README.md | 35 +++++ tests/{ => tools}/identical_vllm.sh | 2 +- .../run_qwen2_moe_vllm_pipeline.sh | 2 +- tests/{ => tools}/run_regression.sh | 12 +- tests/{ => tools}/run_tp_compare_hf.sh | 2 +- tests/{ => tools}/run_tp_compare_vllm.sh | 2 +- tests/{ => tools}/verify_hf.sh | 0 tests/{ => tools}/verify_vllm.sh | 2 +- 24 files changed, 261 insertions(+), 19 deletions(-) create mode 100644 tests/_requirements.py create mode 100644 tests/tools/README.md rename tests/{ => tools}/identical_vllm.sh (95%) rename tests/{ => tools}/run_qwen2_moe_vllm_pipeline.sh (98%) rename tests/{ => tools}/run_regression.sh (77%) rename tests/{ => tools}/run_tp_compare_hf.sh (95%) rename tests/{ => tools}/run_tp_compare_vllm.sh (96%) rename tests/{ => tools}/verify_hf.sh (100%) rename tests/{ => tools}/verify_vllm.sh (97%) diff --git a/pyproject.toml b/pyproject.toml index 3b1539913..e495d7def 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,21 @@ monitoring = [ include = ["monitoring*", "benchmark*", "integration"] [tool.pytest.ini_options] +# `manual` tools and native-ring sources are not collected by default. The +# default suite is the CPU contract/unit tests: +# python -m pytest -m "not gpu and not e2e and not manual" -q +addopts = "-ra -m 'not manual'" +norecursedirs = ["tests/tools", "tests/ring", ".venv", "integration", "build"] markers = [ + "cpu: pure-CPU contract/unit test; the default suite (no CUDA / ClickHouse / vLLM runtime / model weights / native extension build required)", + "gpu: requires a CUDA device", + "multi_gpu: requires >= 2 CUDA devices (TP / EP / routing)", + "e2e: end-to-end pipeline through the native backend + host engine", + "clickhouse: requires a reachable ClickHouse instance", + "vllm: requires the vLLM runtime importable", + "hf: requires HuggingFace weights / model cache", + "ring_native: native CUDA ring tests built via tests/ring/Makefile (needs nvcc)", "slow: tests that take more than ~30 s (per-hook isolation full sweep, large E2E sweeps); skipped by default unless `-m slow` is passed", + "manual: investigation / tooling, not a regression gate; not collected by default", + "numeric: per-hook numeric-difference study (drift vs the unhooked baseline)", ] diff --git a/tests/_requirements.py b/tests/_requirements.py new file mode 100644 index 000000000..bb40407ab --- /dev/null +++ b/tests/_requirements.py @@ -0,0 +1,128 @@ +"""Resource skip-guards for the test suite. + +Every GPU / E2E / native test should fail *closed with a precise reason* when a +prerequisite is missing, never with an import error or an opaque crash. Each +helper here returns a ``pytest.mark.skipif`` marker, so it can be used either as +a decorator:: + + from tests._requirements import require_cuda + + @require_cuda() + def test_kernel(): + ... + +or composed into a module-level mark list alongside a category marker:: + + pytestmark = [pytest.mark.gpu, require_cuda()] + +This module must stay importable on a CPU-only box with no CUDA, ClickHouse, +vLLM, model weights, or native build toolchain present. All heavy imports +(``torch`` in particular) are deferred into the helper bodies so that merely +importing this file costs nothing. +""" +from __future__ import annotations + +import importlib.util +import os +import shutil +import socket + +import pytest + +__all__ = [ + "require_cuda", + "require_gpus", + "require_clickhouse", + "require_vllm", + "require_model_cache", + "require_nvcc", +] + + +def _cuda_device_count() -> int: + """CUDA device count, or 0 if torch/CUDA is unavailable. Never raises.""" + try: + import torch + except Exception: + return 0 + try: + return torch.cuda.device_count() if torch.cuda.is_available() else 0 + except Exception: + return 0 + + +def require_cuda(): + """Skip unless at least one CUDA device is visible.""" + return pytest.mark.skipif( + _cuda_device_count() < 1, reason="no CUDA device available" + ) + + +def require_gpus(n: int): + """Skip unless at least ``n`` CUDA devices are visible.""" + have = _cuda_device_count() + return pytest.mark.skipif( + have < n, reason=f"needs >= {n} CUDA device(s), found {have}" + ) + + +def require_clickhouse(host: str | None = None, port: int | None = None): + """Skip unless a ClickHouse TCP port is reachable. + + Host/port default to the ``DMX_DB_HOST`` / ``DMX_DB_PORT`` env vars (and + finally ``127.0.0.1:9000``), matching the runners' connection defaults. + """ + host = host or os.environ.get("DMX_DB_HOST", "127.0.0.1") + port = int(port if port is not None else os.environ.get("DMX_DB_PORT", "9000")) + reachable = False + try: + with socket.create_connection((host, port), timeout=1.0): + reachable = True + except OSError: + reachable = False + return pytest.mark.skipif( + not reachable, reason=f"ClickHouse unreachable at {host}:{port}" + ) + + +def require_vllm(): + """Skip unless the vLLM runtime is importable.""" + available = importlib.util.find_spec("vllm") is not None + return pytest.mark.skipif(not available, reason="vLLM not importable") + + +def _model_in_cache(model: str) -> bool: + """Best-effort check that ``model`` is a local path or a cached HF repo.""" + # Explicit local path (a checkpoint dir). + if os.path.sep in model and os.path.exists(model): + return True + # HuggingFace hub cache layout: ``models----``. + cache_root = os.environ.get( + "HF_HUB_CACHE", + os.path.join( + os.environ.get( + "HF_HOME", os.path.expanduser("~/.cache/huggingface") + ), + "hub", + ), + ) + repo_dir = "models--" + model.replace("/", "--") + return os.path.isdir(os.path.join(cache_root, repo_dir)) + + +def require_model_cache(model: str): + """Skip unless ``model`` (HF repo id or local path) is already on disk. + + Keeps the default/offline suites from triggering a network download. + """ + return pytest.mark.skipif( + not _model_in_cache(model), + reason=f"model {model!r} not found in local cache", + ) + + +def require_nvcc(): + """Skip unless the CUDA compiler ``nvcc`` is on PATH (native ring tests).""" + return pytest.mark.skipif( + shutil.which("nvcc") is None, reason="nvcc not found on PATH" + ) diff --git a/tests/test_adapter_protocol.py b/tests/test_adapter_protocol.py index 11151beb4..e9d41414f 100644 --- a/tests/test_adapter_protocol.py +++ b/tests/test_adapter_protocol.py @@ -16,9 +16,13 @@ import dataclasses +import pytest + from monitoring.adaptor_base import BackendAdaptor from monitoring.step_context import StepContext +pytestmark = pytest.mark.cpu + # --------------------------------------------------------------------------- # Test doubles diff --git a/tests/test_config.py b/tests/test_config.py index d5bd177a7..3c7ab4da1 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -15,6 +15,8 @@ from monitoring.config import CaptureSchedule, MonitoringConfig +pytestmark = pytest.mark.cpu + # ============================================================================= # CaptureSchedule diff --git a/tests/test_e2e_correctness_vs_hf.py b/tests/test_e2e_correctness_vs_hf.py index 31cb562f1..8f2080940 100644 --- a/tests/test_e2e_correctness_vs_hf.py +++ b/tests/test_e2e_correctness_vs_hf.py @@ -76,6 +76,13 @@ import pytest import torch +pytestmark = [ + pytest.mark.gpu, + pytest.mark.e2e, + pytest.mark.clickhouse, + pytest.mark.hf, +] + from monitoring.clickhouse_reader import CHClickhouseDriverReadOnly from monitoring.segment_merger import merge_segments, parse_internal_id diff --git a/tests/test_gpt2_parity.py b/tests/test_gpt2_parity.py index 08ad887e5..2cac0705a 100644 --- a/tests/test_gpt2_parity.py +++ b/tests/test_gpt2_parity.py @@ -5,6 +5,11 @@ from transformers.models.gpt2_p.modeling_gpt2 import HookedGPT2Model from transformers.models.gpt2_p.modeling_gpt2 import GPT2LMHeadModel as HFModifiedGPT2 +from tests._requirements import require_model_cache + +# Runs on CPU but pulls real gpt2 weights via from_pretrained -> `hf`, not `cpu`. +pytestmark = [pytest.mark.hf, require_model_cache("gpt2")] + @pytest.fixture(scope="module") def gpt2_tokenizer(): diff --git a/tests/test_hf_eos_strip.py b/tests/test_hf_eos_strip.py index 85034d8b2..ee58af598 100644 --- a/tests/test_hf_eos_strip.py +++ b/tests/test_hf_eos_strip.py @@ -21,6 +21,8 @@ import pytest import torch +pytestmark = pytest.mark.cpu + from integration.hf_adapter import HFAdaptor diff --git a/tests/test_hook_spec_flags.py b/tests/test_hook_spec_flags.py index c081f12ac..f95510a2a 100644 --- a/tests/test_hook_spec_flags.py +++ b/tests/test_hook_spec_flags.py @@ -11,11 +11,14 @@ from unittest.mock import MagicMock +import pytest import torch import torch.nn as nn from monitoring.ring_transport import HookSpec, ModelShapeConfig, RingTransport +pytestmark = pytest.mark.cpu + def test_hook_spec_flag_defaults_false(): spec = HookSpec(hook_type=0, module=nn.Identity()) diff --git a/tests/test_moe_v1_routing_hooks.py b/tests/test_moe_v1_routing_hooks.py index 271092f59..1959c94fd 100644 --- a/tests/test_moe_v1_routing_hooks.py +++ b/tests/test_moe_v1_routing_hooks.py @@ -2,6 +2,8 @@ import json +import pytest + from transformers import Qwen2MoeConfig from transformers.models.qwen2_moe_compare.modeling_qwen2_moe import CompareQwen2MoeForCausalLM from transformers.models.qwen2_moe_p.modeling_qwen2_moe import HookedQwen2MoeForCausalLM @@ -19,6 +21,8 @@ ) from tests.ref_disk_worker import _ARCH_REMAP as _REF_ARCH_REMAP +pytestmark = pytest.mark.cpu + def test_moe_v1_routing_hook_types_registered() -> None: assert _id_by_short["router_logits"] == HOOK_TYPE_ROUTER_LOGITS diff --git a/tests/test_no_graph_breaks.py b/tests/test_no_graph_breaks.py index 0b9b4b6f4..e640cccb0 100644 --- a/tests/test_no_graph_breaks.py +++ b/tests/test_no_graph_breaks.py @@ -34,6 +34,8 @@ import pytest +pytestmark = [pytest.mark.gpu, pytest.mark.vllm, pytest.mark.hf] + REPO_ROOT = Path(__file__).resolve().parent.parent HF_GPT2_MODEL = "gpt2" diff --git a/tests/test_per_hook_isolation.py b/tests/test_per_hook_isolation.py index 97b77735b..dedcd1940 100644 --- a/tests/test_per_hook_isolation.py +++ b/tests/test_per_hook_isolation.py @@ -89,6 +89,7 @@ def allocate(self): """).strip("\n") +@pytest.mark.cpu class TestPatcherCorrectness: """Verify the regex + line-by-line patching logic.""" @@ -126,6 +127,7 @@ def test_indentation_preserved_on_commented_lines(self): assert indent == 8, f"unexpected indent on: {line!r}" +@pytest.mark.cpu class TestPatcherRoundTrip: """Verify the on-disk patch / unpatch context manager.""" @@ -163,6 +165,7 @@ def test_stale_backup_raises(self, tmp_path, monkeypatch): isolate_hook.patch("test", "fake", "q") +@pytest.mark.cpu class TestRealCompareModelsContainAllExpectedHooks: """The smoke cells assume specific hooks have a `.copy_()` line in the real _compare files. If a hook is missing the smoke fails opaquely, @@ -488,6 +491,7 @@ def _run_rollout( ) +@pytest.mark.gpu @pytest.mark.slow @pytest.mark.parametrize( "framework,model_key,hook,mode", SMOKE_CELLS, diff --git a/tests/test_producer_chunked_schema.py b/tests/test_producer_chunked_schema.py index bcc77f843..0508eb364 100644 --- a/tests/test_producer_chunked_schema.py +++ b/tests/test_producer_chunked_schema.py @@ -21,18 +21,35 @@ def setup_module(module): # noqa: D401 -- pytest hook _load_extension() # ensure .so loaded -> three ring ops registered +# --- Registration / wiring: no CUDA device required (CPU default suite) ------ + +@pytest.mark.cpu def test_producer_op_registered(): assert torch.ops.ring.producer.default is not None +@pytest.mark.cpu def test_producer_prefix_op_registered(): assert torch.ops.ring.producer_prefix.default is not None +@pytest.mark.cpu def test_producer_chunked_op_registered(): assert torch.ops.ring.producer_chunked.default is not None +@pytest.mark.cpu +def test_hook_point_strip_attrs_default_to_static(): + """HookPoint instances default to the static path.""" + from monitoring.hook_points import HookPoint + hp = HookPoint() + assert hp._strip_tensor is None + assert hp._strip_row_bytes == 0 + + +# --- Device smoke: allocate CUDA tensors and dispatch the op (GPU only) ------- + +@pytest.mark.gpu def test_producer_static_smoke(): """Static op accepts (Tensor(a!), Tensor, int, int). C++ impl early-returns when no engine is active, so this is a pure schema @@ -44,6 +61,7 @@ def test_producer_static_smoke(): torch.ops.ring.producer(ring_payload, x, 0, 0) +@pytest.mark.gpu def test_producer_prefix_smoke(): """Prefix op accepts (Tensor(a!), Tensor, Tensor, int, int, int).""" if not torch.cuda.is_available(): @@ -54,6 +72,7 @@ def test_producer_prefix_smoke(): torch.ops.ring.producer_prefix(ring_payload, x, row_count, 8, 0, 0) +@pytest.mark.gpu def test_producer_chunked_smoke(): """Chunked op accepts (Tensor(a!), Tensor, Tensor, int, int).""" if not torch.cuda.is_available(): @@ -64,14 +83,7 @@ def test_producer_chunked_smoke(): torch.ops.ring.producer_chunked(ring_payload, x, chunk_bytes, 0, 0) -def test_hook_point_strip_attrs_default_to_static(): - """HookPoint instances default to the static path.""" - from monitoring.hook_points import HookPoint - hp = HookPoint() - assert hp._strip_tensor is None - assert hp._strip_row_bytes == 0 - - +@pytest.mark.gpu def test_hook_point_strip_attrs_settable_for_prefix_mode(): """Setting _strip_tensor + _strip_row_bytes > 0 selects prefix mode.""" if not torch.cuda.is_available(): @@ -85,6 +97,7 @@ def test_hook_point_strip_attrs_settable_for_prefix_mode(): assert hp._strip_row_bytes == 8 +@pytest.mark.gpu def test_hook_point_strip_attrs_settable_for_chunked_mode(): """Setting _strip_tensor + _strip_row_bytes == 0 selects chunked mode.""" if not torch.cuda.is_available(): diff --git a/tests/test_ring_transport_no_framework_strings.py b/tests/test_ring_transport_no_framework_strings.py index cb5acf064..9bc8c7d82 100644 --- a/tests/test_ring_transport_no_framework_strings.py +++ b/tests/test_ring_transport_no_framework_strings.py @@ -25,6 +25,8 @@ import pytest +pytestmark = pytest.mark.cpu + RING_TRANSPORT = ( Path(__file__).resolve().parent.parent / "monitoring" / "ring_transport.py" ) diff --git a/tests/test_tp_shapes.py b/tests/test_tp_shapes.py index d56f6889c..c6a3ea441 100644 --- a/tests/test_tp_shapes.py +++ b/tests/test_tp_shapes.py @@ -6,6 +6,8 @@ import pytest import torch +pytestmark = pytest.mark.cpu + from monitoring.ring_transport import ( HOOK_TYPE_RESID_PRE, HOOK_TYPE_LN1, diff --git a/tests/test_vllm_identical.py b/tests/test_vllm_identical.py index 92f153ed6..0049a05f0 100644 --- a/tests/test_vllm_identical.py +++ b/tests/test_vllm_identical.py @@ -40,6 +40,13 @@ import pytest import torch +pytestmark = [ + pytest.mark.gpu, + pytest.mark.vllm, + pytest.mark.clickhouse, + pytest.mark.e2e, +] + _MODEL_REF_FILES = { "gpt2": "gpt2_ref.py", "qwen2_moe": "qwen2_moe_ref.py", diff --git a/tests/test_vllm_rowcnt.py b/tests/test_vllm_rowcnt.py index 13f1b1ceb..09231df3d 100644 --- a/tests/test_vllm_rowcnt.py +++ b/tests/test_vllm_rowcnt.py @@ -40,6 +40,13 @@ import pytest import torch +pytestmark = [ + pytest.mark.gpu, + pytest.mark.vllm, + pytest.mark.clickhouse, + pytest.mark.e2e, +] + _MODEL_ALIASES = { "gpt2": "gpt2", "qwen2_moe": "Qwen/Qwen1.5-MoE-A2.7B", diff --git a/tests/tools/README.md b/tests/tools/README.md new file mode 100644 index 000000000..e6225cd19 --- /dev/null +++ b/tests/tools/README.md @@ -0,0 +1,35 @@ +# tests/tools — manual analysis & release-sweep scripts + +These are **manual** entry points: debugging aids, transport/correctness sweeps, +and release-candidate regression wrappers. They are intentionally **not** part of +the pytest regression gates — `pyproject.toml` lists `tests/tools` under +`norecursedirs`, so pytest never discovers anything here. + +Run them by hand, from the **repository root**, when you want a full sweep or are +debugging a specific backend. They require a GPU (most also need ClickHouse and +the vLLM runtime); they are not CPU-safe. + +| Script | What it does | +|---|---| +| `run_regression.sh` | Full release sweep: CPU unit tests + HF/vLLM transport correctness across models/modes/TP. Calls the `run_tp_compare_*` wrappers. | +| `run_tp_compare_hf.sh` | Single HF transport-correctness run (`.copy_()` buffers vs ClickHouse) for one model/mode/TP. | +| `run_tp_compare_vllm.sh` | Single vLLM transport-correctness run for one model/mode/TP. | +| `run_qwen2_moe_vllm_pipeline.sh` | Qwen2-MoE / EP vLLM ref → monitored → compare pipeline. | +| `identical_vllm.sh` | Wrapper around the vLLM bitwise-identical pytest check. | +| `verify_vllm.sh` | vLLM row-count + identical verification sweep across ring sizes. | +| `verify_hf.sh` | HF E2E correctness sweep across ring sizes. | + +Example: + +```bash +# from the repo root +LD_PRELOAD=/path/to/libstdc++.so.6 CUDA_VISIBLE_DEVICES=0,1 \ + bash tests/tools/run_regression.sh +``` + +> Native CUDA ring tests live separately under `tests/ring/` (built via its +> `Makefile`, marker `ring_native`, needs `nvcc`) and are likewise excluded from +> default pytest discovery. + +As the configurable E2E matrix (`tests/e2e_matrix`) lands, these hardcoded +wrappers are expected to be superseded by matrix invocations. diff --git a/tests/identical_vllm.sh b/tests/tools/identical_vllm.sh similarity index 95% rename from tests/identical_vllm.sh rename to tests/tools/identical_vllm.sh index 444d6839a..997779513 100755 --- a/tests/identical_vllm.sh +++ b/tests/tools/identical_vllm.sh @@ -11,7 +11,7 @@ # - LD_PRELOAD for libstdc++ (caller's responsibility) # # Usage: -# LD_PRELOAD=/path/to/libstdc++.so.6 bash tests/identical_vllm.sh +# LD_PRELOAD=/path/to/libstdc++.so.6 bash tests/tools/identical_vllm.sh set -e export VLLM_DISABLE_COMPILE_CACHE=1 diff --git a/tests/run_qwen2_moe_vllm_pipeline.sh b/tests/tools/run_qwen2_moe_vllm_pipeline.sh similarity index 98% rename from tests/run_qwen2_moe_vllm_pipeline.sh rename to tests/tools/run_qwen2_moe_vllm_pipeline.sh index 7cc8437ff..89c79ce3b 100755 --- a/tests/run_qwen2_moe_vllm_pipeline.sh +++ b/tests/tools/run_qwen2_moe_vllm_pipeline.sh @@ -1,7 +1,7 @@ #!/bin/bash set -euo pipefail -ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" CONDA_SH="${CONDA_SH:-/home/sixian/miniforge3/etc/profile.d/conda.sh}" CONDA_ENV_NAME="${CONDA_ENV_NAME:-ring_offload}" diff --git a/tests/run_regression.sh b/tests/tools/run_regression.sh similarity index 77% rename from tests/run_regression.sh rename to tests/tools/run_regression.sh index d179b2384..02c3ab69a 100755 --- a/tests/run_regression.sh +++ b/tests/tools/run_regression.sh @@ -1,6 +1,6 @@ #!/bin/bash # Full regression test suite. -# Usage: LD_PRELOAD=... CUDA_VISIBLE_DEVICES=0,1 bash tests/run_regression.sh +# Usage: LD_PRELOAD=... CUDA_VISIBLE_DEVICES=0,1 bash tests/tools/run_regression.sh set -e PASS=0 @@ -23,7 +23,7 @@ run_test() { fi } -PROJECT_ROOT=$(cd "$(dirname "$0")/.." && pwd) +PROJECT_ROOT=$(cd "$(dirname "$0")/../.." && pwd) cd "$PROJECT_ROOT" # --- Unit tests --- @@ -34,9 +34,9 @@ run_test "unit: test_tp_shapes + test_config" \ for model in qwen3 gpt2; do for mode in eager cudagraph; do run_test "vllm: $model $mode tp=1" \ - bash tests/run_tp_compare_vllm.sh "$model" "$mode" 1 + bash tests/tools/run_tp_compare_vllm.sh "$model" "$mode" 1 run_test "vllm: $model $mode tp=2" \ - bash tests/run_tp_compare_vllm.sh "$model" "$mode" 2 + bash tests/tools/run_tp_compare_vllm.sh "$model" "$mode" 2 done done @@ -44,13 +44,13 @@ done for model in gpt2 qwen3; do for mode in eager cudagraph; do run_test "hf: $model $mode tp=1" \ - bash tests/run_tp_compare_hf.sh "$model" "$mode" 1 + bash tests/tools/run_tp_compare_hf.sh "$model" "$mode" 1 done done # HF TP=2 (qwen3 only — gpt2 lacks tp_plan support) for mode in eager cudagraph; do run_test "hf: qwen3 $mode tp=2" \ - bash tests/run_tp_compare_hf.sh qwen3 "$mode" 2 + bash tests/tools/run_tp_compare_hf.sh qwen3 "$mode" 2 done # --- Summary --- diff --git a/tests/run_tp_compare_hf.sh b/tests/tools/run_tp_compare_hf.sh similarity index 95% rename from tests/run_tp_compare_hf.sh rename to tests/tools/run_tp_compare_hf.sh index df4f13ed0..745130895 100755 --- a/tests/run_tp_compare_hf.sh +++ b/tests/tools/run_tp_compare_hf.sh @@ -1,6 +1,6 @@ #!/bin/bash # HF transport correctness test: single run with compare model. -# Usage: bash tests/run_tp_compare_hf.sh [model] [mode] [tp] +# Usage: bash tests/tools/run_tp_compare_hf.sh [model] [mode] [tp] # model: gpt2 (default) or qwen3 # mode: eager (default) or cudagraph # tp: 1 (default) or 2 diff --git a/tests/run_tp_compare_vllm.sh b/tests/tools/run_tp_compare_vllm.sh similarity index 96% rename from tests/run_tp_compare_vllm.sh rename to tests/tools/run_tp_compare_vllm.sh index 06a27801e..f23143921 100755 --- a/tests/run_tp_compare_vllm.sh +++ b/tests/tools/run_tp_compare_vllm.sh @@ -1,6 +1,6 @@ #!/bin/bash # Transport correctness test: single run, compare .copy_() buffers vs ClickHouse. -# Usage: bash tests/run_tp_compare_vllm.sh [model] [mode] [tp] +# Usage: bash tests/tools/run_tp_compare_vllm.sh [model] [mode] [tp] # model: qwen3 (default) or gpt2 # mode: eager (default) or cudagraph # tp: 1 (default) or 2 diff --git a/tests/verify_hf.sh b/tests/tools/verify_hf.sh similarity index 100% rename from tests/verify_hf.sh rename to tests/tools/verify_hf.sh diff --git a/tests/verify_vllm.sh b/tests/tools/verify_vllm.sh similarity index 97% rename from tests/verify_vllm.sh rename to tests/tools/verify_vllm.sh index 877831985..7f65a78a0 100755 --- a/tests/verify_vllm.sh +++ b/tests/tools/verify_vllm.sh @@ -17,7 +17,7 @@ # - LD_PRELOAD for libstdc++ (caller's responsibility) # # Usage: -# LD_PRELOAD=/path/to/libstdc++.so.6 bash tests/verify_vllm.sh +# LD_PRELOAD=/path/to/libstdc++.so.6 bash tests/tools/verify_vllm.sh set -e export VLLM_DISABLE_COMPILE_CACHE=1 From 1a6c4009bddf15ad730069f5c339675376afc51f Mon Sep 17 00:00:00 2001 From: SieDeta Date: Tue, 9 Jun 2026 19:15:49 +0700 Subject: [PATCH 2/4] test: marker taxonomy + skip guards + categorize --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index d7715a391..cf2859f68 100644 --- a/.gitignore +++ b/.gitignore @@ -14,6 +14,7 @@ build/ venv/ env/ ENV/ +.venv/ # IDE .vscode/ From bd5088c5e78a511bb8a8994fe1992cd3b5320b45 Mon Sep 17 00:00:00 2001 From: Samfisheryu Date: Sat, 13 Jun 2026 15:39:35 -0400 Subject: [PATCH 3/4] Fix CPU-only test markers --- monitoring/ring_transport.py | 8 +-- pyproject.toml | 8 ++- tests/test_gpt2_parity.py | 20 ++++-- tests/test_moe_v1_routing_hooks.py | 99 ++++++++++++++++----------- tests/test_producer_chunked_schema.py | 13 ++-- 5 files changed, 88 insertions(+), 60 deletions(-) diff --git a/monitoring/ring_transport.py b/monitoring/ring_transport.py index 4b4197d17..929f2ce25 100644 --- a/monitoring/ring_transport.py +++ b/monitoring/ring_transport.py @@ -172,8 +172,8 @@ class HookSpec: # META_FLAG_ALLOW_MISMATCH; consumer recomputes dim-0 from actual bytes. allow_token_cnt_mismatch: bool = False # True when this spec's shape has dim-0 = total_tokens in the framework's - # layout (vLLM flat: total_tokens; HF batched: batch * q_len when q_len is - # the variable axis). Adapters that enable a padding-strip mode use this + # packed-flat layout, or batch * q_len in the batched layout when q_len is + # the variable axis. Adapters that enable a padding-strip mode use this # flag to mark prefix-eligible specs. Static property; ignored when no # adapter activates strip. dim0_is_actual_tokens: bool = False @@ -420,8 +420,8 @@ def __init__(self, ring_engine: Any) -> None: # 2. fits after flushing the ring -> flush_and_wait + reserve_one + ring # 3. single tensor > ring -> flush_and_wait + submit_cpu_direct # Owned by adaptor_base.before_forward (per-batch reassignment based - # on prepare_step result and dynamic-spec presence). Consumers - # (vLLM dispatch wrapper, HookPoint.forward) read only. + # on prepare_step result and dynamic-spec presence). Dispatch + # wrappers and HookPoint.forward read only. self.force_eager: bool = False # New-path state diff --git a/pyproject.toml b/pyproject.toml index e495d7def..229b7275f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,18 +33,20 @@ include = ["monitoring*", "benchmark*", "integration"] [tool.pytest.ini_options] # `manual` tools and native-ring sources are not collected by default. The -# default suite is the CPU contract/unit tests: -# python -m pytest -m "not gpu and not e2e and not manual" -q +# lightweight PR gate should select true CPU tests explicitly: +# python -m pytest -m "cpu" -q addopts = "-ra -m 'not manual'" norecursedirs = ["tests/tools", "tests/ring", ".venv", "integration", "build"] markers = [ - "cpu: pure-CPU contract/unit test; the default suite (no CUDA / ClickHouse / vLLM runtime / model weights / native extension build required)", + "cpu: pure-CPU contract/unit test (no CUDA / ClickHouse / vLLM runtime / model weights / native extension build required)", "gpu: requires a CUDA device", "multi_gpu: requires >= 2 CUDA devices (TP / EP / routing)", "e2e: end-to-end pipeline through the native backend + host engine", "clickhouse: requires a reachable ClickHouse instance", "vllm: requires the vLLM runtime importable", "hf: requires HuggingFace weights / model cache", + "native_backend: requires the DMI native backend .so importable/built; no GPU necessarily", + "framework_fork: requires vendored/modified framework forks importable", "ring_native: native CUDA ring tests built via tests/ring/Makefile (needs nvcc)", "slow: tests that take more than ~30 s (per-hook isolation full sweep, large E2E sweeps); skipped by default unless `-m slow` is passed", "manual: investigation / tooling, not a regression gate; not collected by default", diff --git a/tests/test_gpt2_parity.py b/tests/test_gpt2_parity.py index 2cac0705a..4d80bbd4a 100644 --- a/tests/test_gpt2_parity.py +++ b/tests/test_gpt2_parity.py @@ -1,24 +1,32 @@ -import torch import pytest -from transformers import AutoTokenizer, GPT2Config -from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel as HFOriginalGPT2 -from transformers.models.gpt2_p.modeling_gpt2 import HookedGPT2Model -from transformers.models.gpt2_p.modeling_gpt2 import GPT2LMHeadModel as HFModifiedGPT2 +import torch from tests._requirements import require_model_cache # Runs on CPU but pulls real gpt2 weights via from_pretrained -> `hf`, not `cpu`. -pytestmark = [pytest.mark.hf, require_model_cache("gpt2")] +pytestmark = [pytest.mark.hf, pytest.mark.framework_fork, require_model_cache("gpt2")] + + +def _import_gpt2_modules(): + try: + from transformers import AutoTokenizer + from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel as HFOriginalGPT2 + from transformers.models.gpt2_p.modeling_gpt2 import GPT2LMHeadModel as HFModifiedGPT2 + except ImportError as exc: + pytest.skip(f"modified transformers fork required: {exc}") + return AutoTokenizer, HFOriginalGPT2, HFModifiedGPT2 @pytest.fixture(scope="module") def gpt2_tokenizer(): + AutoTokenizer, _, _ = _import_gpt2_modules() tokenizer = AutoTokenizer.from_pretrained("gpt2") tokenizer.pad_token = tokenizer.eos_token return tokenizer def build_models(seed: int = 0): + _, HFOriginalGPT2, HFModifiedGPT2 = _import_gpt2_modules() torch.manual_seed(seed) hf_original = HFOriginalGPT2.from_pretrained("gpt2") torch.manual_seed(seed) diff --git a/tests/test_moe_v1_routing_hooks.py b/tests/test_moe_v1_routing_hooks.py index 1959c94fd..d8e2110f7 100644 --- a/tests/test_moe_v1_routing_hooks.py +++ b/tests/test_moe_v1_routing_hooks.py @@ -1,37 +1,48 @@ from __future__ import annotations +from functools import lru_cache import json +from types import SimpleNamespace import pytest -from transformers import Qwen2MoeConfig -from transformers.models.qwen2_moe_compare.modeling_qwen2_moe import CompareQwen2MoeForCausalLM -from transformers.models.qwen2_moe_p.modeling_qwen2_moe import HookedQwen2MoeForCausalLM - -from integration.model_shape import _make_model_shape_from_hf_config -from integration.vllm_adapter import _ARCH_REMAP -from integration.vllm.vllm.model_executor.models.enable_ref_hooks import enable_ref_hooks -from integration.vllm.vllm.model_executor.models.registry import _TEXT_GENERATION_MODELS -from monitoring.ring_transport import ( - HOOK_TYPE_ROUTER_LOGITS, - HOOK_TYPE_TOPK_IDS, - HOOK_TYPE_TOPK_WEIGHTS, - _compute_hook_shape, - _id_by_short, -) -from tests.ref_disk_worker import _ARCH_REMAP as _REF_ARCH_REMAP - -pytestmark = pytest.mark.cpu +pytestmark = pytest.mark.framework_fork + + +@lru_cache(maxsize=1) +def _mods() -> SimpleNamespace: + try: + from transformers import Qwen2MoeConfig + from transformers.models.qwen2_moe_compare.modeling_qwen2_moe import CompareQwen2MoeForCausalLM + from transformers.models.qwen2_moe_p.modeling_qwen2_moe import HookedQwen2MoeForCausalLM + + from integration.model_shape import _make_model_shape_from_hf_config + from integration.vllm_adapter import _ARCH_REMAP + from integration.vllm.vllm.model_executor.models.enable_ref_hooks import enable_ref_hooks + from integration.vllm.vllm.model_executor.models.registry import _TEXT_GENERATION_MODELS + from monitoring.ring_transport import ( + HOOK_TYPE_ROUTER_LOGITS, + HOOK_TYPE_TOPK_IDS, + HOOK_TYPE_TOPK_WEIGHTS, + _compute_hook_shape, + _id_by_short, + ) + from tests.ref_disk_worker import _ARCH_REMAP as _REF_ARCH_REMAP + except ImportError as exc: + pytest.skip(f"modified framework forks required: {exc}") + return SimpleNamespace(**locals()) def test_moe_v1_routing_hook_types_registered() -> None: - assert _id_by_short["router_logits"] == HOOK_TYPE_ROUTER_LOGITS - assert _id_by_short["topk_ids"] == HOOK_TYPE_TOPK_IDS - assert _id_by_short["topk_weights"] == HOOK_TYPE_TOPK_WEIGHTS + m = _mods() + assert m._id_by_short["router_logits"] == m.HOOK_TYPE_ROUTER_LOGITS + assert m._id_by_short["topk_ids"] == m.HOOK_TYPE_TOPK_IDS + assert m._id_by_short["topk_weights"] == m.HOOK_TYPE_TOPK_WEIGHTS def test_moe_v1_routing_shapes_from_qwen2_moe_config() -> None: - cfg = Qwen2MoeConfig( + m = _mods() + cfg = m.Qwen2MoeConfig( hidden_size=64, intermediate_size=128, num_hidden_layers=2, @@ -41,45 +52,49 @@ def test_moe_v1_routing_shapes_from_qwen2_moe_config() -> None: num_experts_per_tok=4, vocab_size=128, ) - model_shape = _make_model_shape_from_hf_config(cfg) + model_shape = m._make_model_shape_from_hf_config(cfg) assert model_shape is not None q_len = 17 kv_dim = 17 - assert _compute_hook_shape( - HOOK_TYPE_ROUTER_LOGITS, model_shape, batch=0, q_len=q_len, kv_dim=kv_dim + assert m._compute_hook_shape( + m.HOOK_TYPE_ROUTER_LOGITS, model_shape, batch=0, q_len=q_len, kv_dim=kv_dim ) == [q_len, 60] - assert _compute_hook_shape( - HOOK_TYPE_TOPK_IDS, model_shape, batch=0, q_len=q_len, kv_dim=kv_dim + assert m._compute_hook_shape( + m.HOOK_TYPE_TOPK_IDS, model_shape, batch=0, q_len=q_len, kv_dim=kv_dim ) == [q_len, 4] - assert _compute_hook_shape( - HOOK_TYPE_TOPK_WEIGHTS, model_shape, batch=0, q_len=q_len, kv_dim=kv_dim + assert m._compute_hook_shape( + m.HOOK_TYPE_TOPK_WEIGHTS, model_shape, batch=0, q_len=q_len, kv_dim=kv_dim ) == [q_len, 4] def test_vllm_adapter_remaps_qwen2_moe_to_hooked_variant() -> None: - assert _ARCH_REMAP["Qwen2MoeForCausalLM"] == "Qwen2MoePForCausalLM" + m = _mods() + assert m._ARCH_REMAP["Qwen2MoeForCausalLM"] == "Qwen2MoePForCausalLM" def test_vllm_compare_model_is_registered() -> None: - assert _TEXT_GENERATION_MODELS["Qwen2MoeCompareForCausalLM"] == ( + m = _mods() + assert m._TEXT_GENERATION_MODELS["Qwen2MoeCompareForCausalLM"] == ( "qwen2_moe_compare", "Qwen2MoeCompareForCausalLM", ) def test_vllm_ref_model_is_registered() -> None: - assert _TEXT_GENERATION_MODELS["Qwen2MoeRefForCausalLM"] == ( + m = _mods() + assert m._TEXT_GENERATION_MODELS["Qwen2MoeRefForCausalLM"] == ( "qwen2_moe_ref", "Qwen2MoeRefForCausalLM", ) - assert _REF_ARCH_REMAP["Qwen2MoeForCausalLM"] == "Qwen2MoeRefForCausalLM" + assert m._REF_ARCH_REMAP["Qwen2MoeForCausalLM"] == "Qwen2MoeRefForCausalLM" def test_hf_hooked_qwen2_moe_exposes_routing_hook_specs() -> None: - model = HookedQwen2MoeForCausalLM( - Qwen2MoeConfig( + m = _mods() + model = m.HookedQwen2MoeForCausalLM( + m.Qwen2MoeConfig( hidden_size=64, intermediate_size=128, moe_intermediate_size=64, @@ -94,14 +109,15 @@ def test_hf_hooked_qwen2_moe_exposes_routing_hook_specs() -> None: ) ) emitted = {spec.hook_type for spec in model.get_hook_specs()} - assert HOOK_TYPE_ROUTER_LOGITS in emitted - assert HOOK_TYPE_TOPK_IDS in emitted - assert HOOK_TYPE_TOPK_WEIGHTS in emitted + assert m.HOOK_TYPE_ROUTER_LOGITS in emitted + assert m.HOOK_TYPE_TOPK_IDS in emitted + assert m.HOOK_TYPE_TOPK_WEIGHTS in emitted def test_hf_compare_qwen2_moe_exposes_compare_api() -> None: - model = CompareQwen2MoeForCausalLM( - Qwen2MoeConfig( + m = _mods() + model = m.CompareQwen2MoeForCausalLM( + m.Qwen2MoeConfig( hidden_size=64, intermediate_size=128, moe_intermediate_size=64, @@ -120,12 +136,13 @@ def test_hf_compare_qwen2_moe_exposes_compare_api() -> None: def test_qwen2_moe_ref_preset_adds_routing_hooks(tmp_path) -> None: + m = _mods() model_file = tmp_path / "qwen2_moe_ref.py" model_file.write_text("class Dummy:\n pass\n", encoding="utf-8") out_dir = tmp_path / "out" cfg_out = tmp_path / "ref_config.json" - enable_ref_hooks( + m.enable_ref_hooks( model_file=str(model_file), hooks="vllm-full", max_len=128, diff --git a/tests/test_producer_chunked_schema.py b/tests/test_producer_chunked_schema.py index 0508eb364..ac557604e 100644 --- a/tests/test_producer_chunked_schema.py +++ b/tests/test_producer_chunked_schema.py @@ -16,29 +16,30 @@ from monitoring._native_engine import _load_extension +pytestmark = pytest.mark.native_backend + def setup_module(module): # noqa: D401 -- pytest hook - _load_extension() # ensure .so loaded -> three ring ops registered + try: + _load_extension() # ensure .so loaded -> three ring ops registered + except ImportError as exc: + pytest.skip(f"DMI native backend required: {exc}", allow_module_level=True) -# --- Registration / wiring: no CUDA device required (CPU default suite) ------ +# --- Registration / wiring: no CUDA device required, but native backend needed. -@pytest.mark.cpu def test_producer_op_registered(): assert torch.ops.ring.producer.default is not None -@pytest.mark.cpu def test_producer_prefix_op_registered(): assert torch.ops.ring.producer_prefix.default is not None -@pytest.mark.cpu def test_producer_chunked_op_registered(): assert torch.ops.ring.producer_chunked.default is not None -@pytest.mark.cpu def test_hook_point_strip_attrs_default_to_static(): """HookPoint instances default to the static path.""" from monitoring.hook_points import HookPoint From 2c133decc4d97f32eba574c3f82c37e605cff797 Mon Sep 17 00:00:00 2001 From: Samfisheryu Date: Sat, 13 Jun 2026 16:00:37 -0400 Subject: [PATCH 4/4] Move native-dependent tests out of CPU gate --- tests/test_adapter_protocol.py | 23 +++++++-- tests/test_hf_eos_strip.py | 21 ++++++-- tests/test_hook_spec_flags.py | 15 +++++- tests/test_tp_shapes.py | 87 ++++++++++++++++++++++++---------- 4 files changed, 109 insertions(+), 37 deletions(-) diff --git a/tests/test_adapter_protocol.py b/tests/test_adapter_protocol.py index e9d41414f..faa6a0e03 100644 --- a/tests/test_adapter_protocol.py +++ b/tests/test_adapter_protocol.py @@ -10,7 +10,8 @@ set transport.force_eager from (result == 2) OR needs_eager. -> set_step_context -> pre_push_all_metas -No GPU / native engine required. +No GPU required; imports the native hook-definition layer through +BackendAdaptor's ring_transport dependency. """ from __future__ import annotations @@ -18,10 +19,22 @@ import pytest -from monitoring.adaptor_base import BackendAdaptor -from monitoring.step_context import StepContext - -pytestmark = pytest.mark.cpu +try: + from monitoring.adaptor_base import BackendAdaptor + from monitoring.step_context import StepContext + _NATIVE_IMPORT_ERROR = None +except ImportError as exc: + BackendAdaptor = object + StepContext = None + _NATIVE_IMPORT_ERROR = exc + +pytestmark = [ + pytest.mark.native_backend, + pytest.mark.skipif( + _NATIVE_IMPORT_ERROR is not None, + reason=f"DMI native backend required: {_NATIVE_IMPORT_ERROR}", + ), +] # --------------------------------------------------------------------------- diff --git a/tests/test_hf_eos_strip.py b/tests/test_hf_eos_strip.py index ee58af598..8eb994580 100644 --- a/tests/test_hf_eos_strip.py +++ b/tests/test_hf_eos_strip.py @@ -2,7 +2,9 @@ Phase 4.B verification gate. Pure unit test against the strip logic -- no DB, no GPU, no real model. Reuses the FakeEngine pattern from -``tests/test_adapter_protocol.py``. +``tests/test_adapter_protocol.py``. HFAdaptor imports the native hook +definition layer through ring_transport, so this is not part of the +no-native-build CPU gate. The strip semantic under test: * Detection runs one step late by construction -- ``input_ids[:, -1]`` @@ -21,9 +23,20 @@ import pytest import torch -pytestmark = pytest.mark.cpu - -from integration.hf_adapter import HFAdaptor +try: + from integration.hf_adapter import HFAdaptor + _NATIVE_IMPORT_ERROR = None +except ImportError as exc: + HFAdaptor = None + _NATIVE_IMPORT_ERROR = exc + +pytestmark = [ + pytest.mark.native_backend, + pytest.mark.skipif( + _NATIVE_IMPORT_ERROR is not None, + reason=f"DMI native backend required: {_NATIVE_IMPORT_ERROR}", + ), +] # --------------------------------------------------------------------------- diff --git a/tests/test_hook_spec_flags.py b/tests/test_hook_spec_flags.py index f95510a2a..b90b67ea4 100644 --- a/tests/test_hook_spec_flags.py +++ b/tests/test_hook_spec_flags.py @@ -15,9 +15,20 @@ import torch import torch.nn as nn -from monitoring.ring_transport import HookSpec, ModelShapeConfig, RingTransport +try: + from monitoring.ring_transport import HookSpec, ModelShapeConfig, RingTransport + _NATIVE_IMPORT_ERROR = None +except ImportError as exc: + HookSpec = ModelShapeConfig = RingTransport = None + _NATIVE_IMPORT_ERROR = exc -pytestmark = pytest.mark.cpu +pytestmark = [ + pytest.mark.native_backend, + pytest.mark.skipif( + _NATIVE_IMPORT_ERROR is not None, + reason=f"DMI native backend required: {_NATIVE_IMPORT_ERROR}", + ), +] def test_hook_spec_flag_defaults_false(): diff --git a/tests/test_tp_shapes.py b/tests/test_tp_shapes.py index c6a3ea441..ac2a504f5 100644 --- a/tests/test_tp_shapes.py +++ b/tests/test_tp_shapes.py @@ -1,36 +1,71 @@ """Unit tests for TP-aware shape computation in _compute_hook_shape. -No GPU or distributed setup needed — tests use ModelShapeConfig directly. +No GPU or distributed setup needed. These tests still import ring_transport's +native hook-definition layer, so they are not part of the no-native-build CPU +gate. """ import pytest import torch -pytestmark = pytest.mark.cpu - -from monitoring.ring_transport import ( - HOOK_TYPE_RESID_PRE, - HOOK_TYPE_LN1, - HOOK_TYPE_ATTN_OUT, - HOOK_TYPE_RESID_MID, - HOOK_TYPE_LN2, - HOOK_TYPE_MLP_IN, - HOOK_TYPE_MLP_OUT, - HOOK_TYPE_Q, - HOOK_TYPE_K, - HOOK_TYPE_V, - HOOK_TYPE_Z, - HOOK_TYPE_ATTN_SCORES, - HOOK_TYPE_MLP_POST, - HOOK_TYPE_RESID_FINAL, - HOOK_TYPE_EMBED, - HOOK_TYPE_POS_EMBED, - HOOK_TYPE_FINAL_LN, - HOOK_TYPE_TOKEN_IDS, - HOOK_TYPE_FINAL_LOGITS, - ModelShapeConfig, - _compute_hook_shape, -) +try: + from monitoring.ring_transport import ( + HOOK_TYPE_RESID_PRE, + HOOK_TYPE_LN1, + HOOK_TYPE_ATTN_OUT, + HOOK_TYPE_RESID_MID, + HOOK_TYPE_LN2, + HOOK_TYPE_MLP_IN, + HOOK_TYPE_MLP_OUT, + HOOK_TYPE_Q, + HOOK_TYPE_K, + HOOK_TYPE_V, + HOOK_TYPE_Z, + HOOK_TYPE_ATTN_SCORES, + HOOK_TYPE_MLP_POST, + HOOK_TYPE_RESID_FINAL, + HOOK_TYPE_EMBED, + HOOK_TYPE_POS_EMBED, + HOOK_TYPE_FINAL_LN, + HOOK_TYPE_TOKEN_IDS, + HOOK_TYPE_FINAL_LOGITS, + ModelShapeConfig, + _compute_hook_shape, + ) + _NATIVE_IMPORT_ERROR = None +except ImportError as exc: + ( + HOOK_TYPE_RESID_PRE, + HOOK_TYPE_LN1, + HOOK_TYPE_ATTN_OUT, + HOOK_TYPE_RESID_MID, + HOOK_TYPE_LN2, + HOOK_TYPE_MLP_IN, + HOOK_TYPE_MLP_OUT, + HOOK_TYPE_Q, + HOOK_TYPE_K, + HOOK_TYPE_V, + HOOK_TYPE_Z, + HOOK_TYPE_ATTN_SCORES, + HOOK_TYPE_MLP_POST, + HOOK_TYPE_RESID_FINAL, + HOOK_TYPE_EMBED, + HOOK_TYPE_POS_EMBED, + HOOK_TYPE_FINAL_LN, + HOOK_TYPE_TOKEN_IDS, + HOOK_TYPE_FINAL_LOGITS, + ModelShapeConfig, + _compute_hook_shape, + ) = (None,) * 21 + _NATIVE_IMPORT_ERROR = exc + +pytestmark = [ + pytest.mark.native_backend, + pytest.mark.skipif( + _NATIVE_IMPORT_ERROR is not None, + reason=f"DMI native backend required: {_NATIVE_IMPORT_ERROR}", + ), +] def _cfg(tp_size=1):