From 64633da4f64a651f8c1f4c769207ae003ee9f0af Mon Sep 17 00:00:00 2001 From: Alex Yang Date: Mon, 14 Sep 2026 10:15:08 -0600 Subject: [PATCH 1/4] fix(tests): skip tests/prims_ts instead of erroring when CUTLASS DSL < 4.7 Prims-TS imports `cutlass.experimental.primitives` and `cutlass.experimental.task_scheduling`, which exist only in CUTLASS DSL 4.7+. `tests/prims_ts/conftest.py` called `require_cutlass_dsl_experimental()` at module scope, so on an environment resolved against our declared floor (`nvidia-cutlass-dsl>=4.6.2a0`) it raised during collection: ERROR tests/prims_ts - RuntimeError: Prims-TS requires the CUTLASS DSL wheel. Install the pinned release-branch wheel before using Prims-TS kernels. That is a collection error, so it takes out the whole directory. Raising the floor is not the fix. FlashInfer's lower bound is aligned with vLLM's and SGLang's on purpose -- CuTe DSL is a diamond dependency, and pinning up would make FlashInfer unsatisfiable in those stacks. 4.7-only features are meant to be gated at runtime and enabled when the environment provides them, which is what `sm120_fmha.py` and `kda.py` already do. Prims-TS was the outlier. The blanket require was also over-broad: 20 of the 21 test modules here import without the wheel, including `test_moe_bf16_support.py` and `test_moe_fp8_block_support.py`, which monkeypatch `is_prims_ts_available` and exist precisely to exercise the unavailable path. Skip collected items rather than skipping at module level: a module-level skip collects nothing, so a run scoped to this directory would exit with pytest's "no tests collected" code 5, trading one red for another. Verified that the mechanism here reports skips and exits 0. This also makes `pytest_report_header`'s unavailable branch reachable for the first time -- it could never fire, because the require above it raised first. AI-assisted: analysis and patch prepared with Claude Code. Co-Authored-By: Claude Opus 5 --- tests/prims_ts/conftest.py | 53 ++++++++++++++++++++++++++------------ 1 file changed, 37 insertions(+), 16 deletions(-) diff --git a/tests/prims_ts/conftest.py b/tests/prims_ts/conftest.py index 256321e289a..5e721bf7836 100644 --- a/tests/prims_ts/conftest.py +++ b/tests/prims_ts/conftest.py @@ -16,35 +16,56 @@ 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() +# Every module here except this one imports without the wheel; this one reaches +# the vendored kernels at import time, so it cannot be collected at all and has +# to be dropped before collection rather than skipped during it. +if not HAS_PRIMS_TS_RUNTIME: + collect_ignore = ["test_batched_gemm_captured_schedule_tasks.py"] -def _has_module(name: str) -> bool: - try: - return importlib.util.find_spec(name) is not None - except ModuleNotFoundError: - return False +def pytest_collection_modifyitems(config, items): + """Skip this directory when the installed DSL cannot provide Prims-TS. -HAS_PRIMS_TS_RUNTIME = all( - _has_module(name) - for name in ( - "cutlass", - "cutlass.cute", - "cutlass.experimental.primitives", - "cutlass.experimental.task_scheduling", + 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 + + 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: + if Path(str(item.path)).parent == here: + item.add_marker(skip_prims_ts) @pytest.fixture(autouse=True) From 9b5a3c812d8acb188c37edb1bae99adbc981e612 Mon Sep 17 00:00:00 2001 From: Alex Yang Date: Tue, 15 Sep 2026 21:16:06 -0600 Subject: [PATCH 2/4] test(prims_ts): enforce the collect_ignore list under a simulated DSL 4.6 Review follow-up on two real gaps (thanks @bkryu, and CodeRabbit raised the first independently). 1. `collect_ignore` was a load-bearing claim with nothing enforcing it. A test module that imports the vendored kernels at module scope cannot be marked during collection, only dropped before it -- so adding one without listing it silently reintroduces #5213 (`ERROR tests/prims_ts`) on every pre-4.7 lane. Adds a meta-check that runs `pytest tests/prims_ts --collect-only` under a stub emulating CUTLASS DSL 4.6 and asserts exit 0. Exit 0 is a stricter assertion 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. The stub resolves `cutlass.experimental` as an empty package while leaving `.primitives` / `.task_scheduling` missing. A finder that fails `cutlass.experimental` outright does not simulate 4.6: on the 4.7 wheel it breaks `import cutlass` itself, which is a different failure. Verified the stub blocks 4.7's eager chain with an ImportError, which `ensure_cutlass_dsl_experimental` catches via `except BaseException`. 2. The skip marker matched only direct children of `tests/prims_ts/`, so a future subdirectory would escape it. Now uses `here in Path(...).parents`. The test lives in tests/test_prims_ts_import_isolation.py rather than under tests/prims_ts/, so it still runs on lanes where that directory is skipped. AI-assisted: prepared with Claude Code. Co-Authored-By: Claude Opus 5 --- tests/prims_ts/conftest.py | 9 ++- tests/test_prims_ts_import_isolation.py | 82 +++++++++++++++++++++++++ 2 files changed, 90 insertions(+), 1 deletion(-) diff --git a/tests/prims_ts/conftest.py b/tests/prims_ts/conftest.py index 5e721bf7836..d818aff5122 100644 --- a/tests/prims_ts/conftest.py +++ b/tests/prims_ts/conftest.py @@ -39,6 +39,11 @@ # the vendored kernels at import time, so it cannot be collected at all and has # to be dropped before collection rather than skipped during it. if not HAS_PRIMS_TS_RUNTIME: + # Hand-maintained and load-bearing: a module that imports the vendored kernels + # eagerly cannot be marked during collection, only dropped before it. Adding one + # without listing it here silently reintroduces #5213, so + # tests/test_prims_ts_import_isolation.py enforces this list under a simulated + # CUTLASS DSL 4.6. collect_ignore = ["test_batched_gemm_captured_schedule_tasks.py"] @@ -64,7 +69,9 @@ def pytest_collection_modifyitems(config, items): ) ) for item in items: - if Path(str(item.path)).parent == here: + # `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) diff --git a/tests/test_prims_ts_import_isolation.py b/tests/test_prims_ts_import_isolation.py index 26487503cd5..da2211f3944 100644 --- a/tests/test_prims_ts_import_isolation.py +++ b/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}" + ) From 9bdf8268ce2d44a71b75c80b6fa2876779a39b8c Mon Sep 17 00:00:00 2001 From: Alex Yang Date: Tue, 15 Sep 2026 22:14:46 -0600 Subject: [PATCH 3/4] fix(tests): ignore every prims_ts module that imports cutlass.experimental The meta-check added in the previous commit immediately earned its keep: it proves collect_ignore was incomplete, so the fix it guards was too. `flashinfer/prims_ts/batched_gemm/batched_gemm_config.py` does an unguarded module-level `from cutlass.experimental import primitives`. The requirement is therefore transitive, and my earlier survey only looked one import deep -- it found the single module that reaches `batched_gemm_kernel` directly and missed the eleven that reach `batched_gemm_config`. On a real pre-4.7 lane those eleven would still have raised during collection, i.e. #5213 would have persisted for most of the directory. Full set is all ten `test_batched_gemm_*` modules plus test_moe_bf16_support and test_moe_nvfp4_support. Uses a glob for the former since new ones are the likely growth; over-ignoring is harmless because the whole directory is skipped on such a lane anyway. Verified by replaying the conftest's ignore rules against a transitive module-level import closure: zero uncollectable modules remain un-ignored, and seven modules still collect -- so the run yields items, and exit 0 rather than pytest's "no tests collected" exit 5. AI-assisted: prepared with Claude Code. Co-Authored-By: Claude Opus 5 --- tests/prims_ts/conftest.py | 28 +++++++++++++++++++--------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/tests/prims_ts/conftest.py b/tests/prims_ts/conftest.py index d818aff5122..be126d4e73b 100644 --- a/tests/prims_ts/conftest.py +++ b/tests/prims_ts/conftest.py @@ -35,16 +35,26 @@ # ``pytest_report_header`` instead of raising. HAS_PRIMS_TS_RUNTIME = ensure_cutlass_dsl_experimental() -# Every module here except this one imports without the wheel; this one reaches -# the vendored kernels at import time, so it cannot be collected at all and has -# to be dropped before collection rather than skipped during it. +# 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: - # Hand-maintained and load-bearing: a module that imports the vendored kernels - # eagerly cannot be marked during collection, only dropped before it. Adding one - # without listing it here silently reintroduces #5213, so - # tests/test_prims_ts_import_isolation.py enforces this list under a simulated - # CUTLASS DSL 4.6. - collect_ignore = ["test_batched_gemm_captured_schedule_tasks.py"] + 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): From 96537b6fd9aee29ac9161c9e192825cb225fd73b Mon Sep 17 00:00:00 2001 From: Alex Yang Date: Wed, 16 Sep 2026 14:22:41 -0600 Subject: [PATCH 4/4] fix(tests): add test_moe_compile_cache to the prims_ts collect_ignore The guard test caught a gap in its own fix: `tests/prims_ts` still errored during collection on a pre-4.7 CUTLASS DSL, on `test_moe_compile_cache.py`. The original audit walked one import level. `test_moe_compile_cache` does `from flashinfer.prims_ts.batched_gemm import batched_gemm_run`, and the unguarded `from cutlass.experimental import primitives` is three hops further down (`batched_gemm_run` -> `batched_gemm_quant` -> `batched_gemm_config`), so a module with no GEMM in its name still needs the DSL at import time. That put the real closure at 13 modules rather than 12. Recomputed the closure with an AST walk that resolves relative imports and follows every import-time edge; the list is now exactly the 13 modules that reach `cutlass.experimental` before an item exists, with nothing over-ignored. AI-assisted. Co-Authored-By: Claude Opus 5 --- tests/prims_ts/conftest.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/prims_ts/conftest.py b/tests/prims_ts/conftest.py index be126d4e73b..dfd510dea78 100644 --- a/tests/prims_ts/conftest.py +++ b/tests/prims_ts/conftest.py @@ -42,7 +42,10 @@ # 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. +# modules plus the three MoE modules below. Follow the chain all the way down -- +# ``test_moe_compile_cache`` is three hops from anything named after a GEMM +# (``batched_gemm_run`` -> ``batched_gemm_quant`` -> ``batched_gemm_config``), +# and stopping at the first hop is what left it off this list originally. # # 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 @@ -53,6 +56,7 @@ collect_ignore_glob = ["test_batched_gemm_*.py"] collect_ignore = [ "test_moe_bf16_support.py", + "test_moe_compile_cache.py", "test_moe_nvfp4_support.py", ]