fix(tests): skip tests/prims_ts instead of erroring when CUTLASS DSL < 4.7

tests/prims_ts/conftest.py

@@ -16,35 +16,73 @@
 Shared test setup for vendored Prims-TS kernels.
 """
 
-import importlib.util
 from pathlib import Path
 
 import pytest
 
 from flashinfer.prims_ts.cutlass_dsl import (
+    ensure_cutlass_dsl_experimental,
     get_cutlass_dsl_bootstrap_error,
-    require_cutlass_dsl_experimental,
 )
 
-require_cutlass_dsl_experimental()
+# Prims-TS needs ``cutlass.experimental.primitives`` and
+# ``cutlass.experimental.task_scheduling``, which exist only in CUTLASS DSL 4.7+.
+# FlashInfer's declared floor is deliberately lower so it stays co-installable
+# with vLLM and SGLang, so an environment that cannot provide them is supported
+# and must degrade rather than fail -- the same contract the other 4.7-gated
+# features honor (``flashinfer/attention/cute_dsl/sm120_fmha.py``,
+# ``flashinfer/kda.py``). ``ensure_...`` records the import failure for
+# ``pytest_report_header`` instead of raising.
+HAS_PRIMS_TS_RUNTIME = ensure_cutlass_dsl_experimental()
 
+# Modules that reach ``cutlass.experimental`` at IMPORT time cannot be marked
+# during collection -- they raise before an item exists -- so they have to be
+# dropped before collection instead.
+#
+# The reach is transitive, not direct: ``batched_gemm_config`` does an unguarded
+# module-level ``from cutlass.experimental import primitives``, so every module
+# importing it inherits the requirement. That is all ten ``test_batched_gemm_*``
+# modules plus the two MoE support modules below.
+#
+# The glob is deliberate: new ``test_batched_gemm_*`` modules are the likely
+# growth, and over-ignoring is harmless here (on a pre-4.7 lane the whole
+# directory is skipped anyway). ``test_prims_ts_directory_still_collects_without_
+# cutlass_experimental`` in tests/test_prims_ts_import_isolation.py enforces that
+# this stays complete.
+if not HAS_PRIMS_TS_RUNTIME:
+    collect_ignore_glob = ["test_batched_gemm_*.py"]
+    collect_ignore = [
+        "test_moe_bf16_support.py",
+        "test_moe_nvfp4_support.py",
+    ]
+
+
+def pytest_collection_modifyitems(config, items):
+    """Skip this directory when the installed DSL cannot provide Prims-TS.
 
-def _has_module(name: str) -> bool:
-    try:
-        return importlib.util.find_spec(name) is not None
-    except ModuleNotFoundError:
-        return False
+    Marking collected items rather than skipping at module level is deliberate:
+    a module-level skip collects nothing, and a run scoped to this directory
+    would then exit with pytest's "no tests collected" code 5 -- trading a red
+    for a red. Marked items are collected first, so the run reports skips and
+    exits 0.
+    """
 
+    del config
+    if HAS_PRIMS_TS_RUNTIME:
+        return
 
-HAS_PRIMS_TS_RUNTIME = all(
-    _has_module(name)
-    for name in (
-        "cutlass",
-        "cutlass.cute",
-        "cutlass.experimental.primitives",
-        "cutlass.experimental.task_scheduling",
+    here = Path(__file__).parent
+    skip_prims_ts = pytest.mark.skip(
+        reason=(
+            "Prims-TS needs cutlass.experimental from CUTLASS DSL 4.7+; the "
+            f"installed wheel cannot provide it: {get_cutlass_dsl_bootstrap_error()!r}"
+        )
     )
-)
+    for item in items:
+        # `in ... parents` rather than `== parent`: a future subdirectory of
+        # tests/prims_ts/ would otherwise escape the skip marker.
+        if here in Path(str(item.path)).parents:
+            item.add_marker(skip_prims_ts)
 
 
 @pytest.fixture(autouse=True)

tests/test_prims_ts_import_isolation.py

@@ -114,3 +114,85 @@ def test_prims_ts_bootstrap_scopes_work_tile_info_customization():
 assert WorkTileInfo.__init__ is original_init
 """
     )
+
+
+# Emulates a CUTLASS DSL 4.6 wheel: `cutlass.experimental` resolves as an empty
+# package so the eager chain inside `import cutlass` is satisfied, while every
+# submodule Prims-TS needs (`.primitives`, `.task_scheduling`) stays missing. A
+# finder that fails `cutlass.experimental` outright would not simulate 4.6 -- on
+# the 4.7 wheel it breaks `import cutlass` itself, which is a different failure.
+_DSL46_STUB = """
+import importlib.abc
+import importlib.machinery
+import sys
+
+
+class _EmptyLoader(importlib.abc.Loader):
+    def create_module(self, spec):
+        return None
+
+    def exec_module(self, module):
+        module.__path__ = []
+
+
+class _Stub(importlib.abc.MetaPathFinder):
+    def find_spec(self, fullname, path=None, target=None):
+        if fullname == "cutlass.experimental":
+            return importlib.machinery.ModuleSpec(
+                fullname, _EmptyLoader(), is_package=True
+            )
+        return None
+
+
+sys.meta_path.insert(0, _Stub())
+"""
+
+
+def test_prims_ts_directory_still_collects_without_cutlass_experimental(tmp_path):
+    """``tests/prims_ts`` must COLLECT cleanly on a DSL that lacks 4.7.
+
+    ``collect_ignore`` in ``tests/prims_ts/conftest.py`` is hand-maintained: a new
+    test module that imports the vendored kernels at module scope cannot be marked
+    during collection, only dropped before it, and forgetting to list it silently
+    reintroduces #5213 (``ERROR tests/prims_ts``) on every pre-4.7 lane.
+
+    Exit code 0 is the assertion that matters, and it is stricter than it looks:
+    a missed module raises during collection (exit 2), and a module-level skip --
+    the obvious alternative fix -- collects nothing and exits 5.
+    """
+
+    stub = tmp_path / "dsl46_stub.py"
+    stub.write_text(_DSL46_STUB)
+
+    env = os.environ.copy()
+    env.pop("PYTEST_ADDOPTS", None)
+    env["PYTHONPATH"] = os.pathsep.join(
+        filter(None, (str(tmp_path), str(_REPO_ROOT), env.get("PYTHONPATH")))
+    )
+    env["FLASHINFER_WORKSPACE_BASE"] = str(tmp_path / "workspace")
+
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-m",
+            "pytest",
+            "tests/prims_ts",
+            "--collect-only",
+            "-q",
+            "--color=no",
+            "-p",
+            "dsl46_stub",
+        ],
+        cwd=_REPO_ROOT,
+        env=env,
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+
+    assert result.returncode == 0, (
+        "tests/prims_ts failed to collect under a simulated CUTLASS DSL 4.6.\n"
+        "If this is a newly added module that imports the Prims-TS kernels at "
+        "module scope, add it to collect_ignore in tests/prims_ts/conftest.py.\n"
+        f"exit={result.returncode}\n{result.stdout}\n{result.stderr}"
+    )