feat: isolate cutlass._mlir imports behind compat gateway (#118)

[bikrammajhi]

Summary

Closes #118 — isolates all direct cutlass._mlir usage behind a single compatibility gateway.

CuTeDSL's generated MLIR/NVVM bindings (cutlass._mlir.{ir,arith,llvm,vector,nvvm,cute}) are private implementation detail with no cross-version stability contract — the tcgen05_ld/st breakage between CutDSL 4.5.2 and 4.5.3 was the first confirmed incident. cuLA imported these directly from 10 kernel modules, so a patch release could break kernel emission silently, delayed until import/JIT time.

Changes

1. cula/ops/_mlir_compat.py — the only module that may touch cutlass._mlir:
   • Lazy dialect loading (plain import cula never touches private bindings)
   • Version contract mirroring pyproject.toml (>=4.4.2,<4.7,!=4.5.0), enforced with a fail-fast RuntimeError naming the offending version
   • Canary probes on every dialect's entry points, so a missing/renamed binding fails with an actionable message instead of surfacing mid-JIT
   • vector_extract_element() helper with cross-version dispatch (see below)
2. Migration — all 10 consumer files bind their dialect aliases from the gateway (one import-line swap per file, zero behavioral changes).
3. Bug fix discovered during prep — CutDSL renamed vector.extractelement to vector.extract in the 4.6 line; against nvidia-cutlass-dsl==4.6.2 (resolved by our own pyproject.toml range) store_256b raises AttributeError at JIT time. store_256b is used by the KDA SM100 backward path. The gateway's vector_extract_element dispatches to whichever binding exists, restoring 4.6 compatibility.
4. Tests — 15 headless tests for the gateway: version parsing, contract enforcement, missing dialect/canary fault injection, vector dispatch on both 4.5/4.6 shapes, and a real-wheel smoke test. No GPU required.
5. scripts/modal_validate.py — H100/CUDA 12.9 validation harness (runs the headless suite + SM90 prefill/decode tests in one modal run).

Validation status

• Headless: 15/15 tests/test_cutedsl_compat.py green against nvidia-cutlass-dsl==4.6.2
• All 10 migrated consumer modules import cleanly against the installed wheel
• SM90 suite on Hopper (H100) — running via modal run scripts/modal_validate.py; will attach results

Notes for reviewers

• The old extractelement call in store_256b passed the position as a keyword (position=...); the new vector.extract binding takes it positionally ((source, dynamic_position, static_position)). The dispatch helper deliberately keeps the exact call shapes per branch.
• Canary choices reflect what the migrated consumers actually call; extend _CANARIES/_ANY_OF_CANARIES when new private-API usage lands.

[bikrammajhi]

Aligned with #119 (merged) — this PR builds on top of the tcgen05 signature-detection work; no overlap: #119 owns _cutedsl_compat.py capability probing, this PR owns the remaining cutlass.isolated-mode isolation.

Scope recap

• cula/ops/_mlir_compat.py: new gateway — dialect-presence canaries (arith/cute/ir/llvm/nvvm/vector), lazy cutlass._mlir.\*aliases, version-parse helper, module-level __getattr__ shim, mlir_context_dir().
• 10 op files migrated off direct cutlass._mlir imports.
• cula/ops/ptx.py + cula/ops/sm100/ptx.py: store_256b is now version-agnostic — probes the real API (fixes the CuTeDSL 4.6.2 extractelement→extract break).
• Tests: tests/test_cutedsl_compat.py extends the Detect tcgen05 load/store API by signature #119 file with gateway + JIT-probe coverage; scripts/modal_validate.py runs the compat + SM90 prefill/decode suites headless on an H100 (image pull/build + pytest; JSON summary returned).

Healthy on the minimal local smoke suite; the automated H100 run is in flight — I'll post the summary when it's back.

[icavan]

Requesting changes based on validation of commit 2c16cb9 on an NVIDIA GB200.

Results:

• tests/test_cutedsl_compat.py: 15/15 passed with CuTeDSL 4.5.2 and with an isolated 4.6.2 environment.
• All 10 migrated modules imported successfully with 4.5.2.
• tests/test_lightning_decode.py: 27/27 passed with 4.5.2.
• tests/test_ptx_umma_ws.py::test_ws_ss_tf32: failed on this PR with both 4.5.2 and 4.6.2. The same test passes at the base commit (5161546) with 4.5.2, so this PR introduces a regression in a supported environment while still not fixing the 4.6 path.

The gateway direction is good, but the store_256b compatibility path and validation harness need to be fixed before merge.

[bikrammajhi]

Validation results (addressing all review comments)

All three inline comments are addressed in commits 1abf845 → ff92a1b (head ff92a1b on mlir-compat-gateway).

[P1] vector_extract_element dispatch — fixed

• The helper now takes the Python element index and builds the operand shape each binding expects: prefers extractelement when present (i32 constant constructed internally — this restores the exact pre-gateway 4.5.x path, byte-identical call shape), and otherwise calls extract(vec, [], [position], ...) — the static-position form you verified in a real 4.6.2 JIT. store_256b in cula/ops/ptx.py now passes the plain index.
• Dispatch tests assert the real operand shapes: extractelement(vec, position=i32const) and extract(vec, [], [3]).

[P1] modal_validate.py — fixed

• Module-level TESTS is now a list of complete test paths (tests/...), passed as an argument list to subprocess.run (no shell=True), and a nonzero pytest exit raises inside the Modal function so the run fails.
• The harness is now parametrized by GPU and CuTeDSL version via env vars (CULA_VALIDATE_GPU, CULA_VALIDATE_CUTLASS) and includes test_ptx_umma_ws.py (conftest self-gates SM100 tests — skip on Hopper, run on Blackwell).

[P2] version parsing — fixed

• _parse_version normalizes to a 3-component (major, minor, patch) tuple; 4.5 → (4, 5, 0) now hits the explicit 4.5.0 exclusion and 4.7 → (4, 7, 0) hits the < 4.7.0 upper bound. Boundary tests added for both.

Results

Environment	compat	SM90	SM100 (test_ptx_umma_ws)
H100, CuTeDSL 4.6.2	18/18	prefill+decode green	n/a (conftest skip)
B200, CuTeDSL 4.6.2	green	green	test_ws_ss_tf32, test_ws_ss_f16, test_ws_ss_tf32_collector PASSED (72/72 total)
B200, CuTeDSL 4.5.2	green	green	72/72 total, incl. test_ws_ss_tf32 PASSED

The test you failed on GB200 (test_ws_ss_tf32) now passes on Blackwell with both 4.5.2 and 4.6.2 — no regression in the supported environment, and the 4.6 path is fixed. Headless suite: 18/18. Ruff lint + format clean.

Full harness output is pasted inline in the follow-up comment (Modal dashboard links are account-private, so the results are shared here instead).

[bikrammajhi]

Full validation output (public copy)

The Modal dashboard links above are account-private (Modal app pages are not publicly viewable, and the /apps/{workspace}/... URL resolves only for the owning account), so here is the complete harness output for the two review-critical runs. All three reviewer comments are addressed at head ff92a1b.

B200 + CuTeDSL 4.5.2 (supported environment, previously FAILED on GB200)

{
 "branch": "mlir-compat-gateway",
 "gpu": "B200",
 "cutlass_spec": "nvidia-cutlass-dsl==4.5.2",
 "steps": {"clone": "ok", "build": "ok", "probe": "ok", "pytest": "ok"},
 "env": "2.9.1+cu129 4.5.2 NVIDIA B200",
 "pytest_rc": 0,
 "pytest_tail": "tests/test_ptx_umma_ws.py::test_ws_ss_tf32 PASSED [97%]\ntests/test_ptx_umma_ws.py::test_ws_ss_f16 PASSED [98%]\ntests/test_ptx_umma_ws.py::test_ws_ss_tf32_collector PASSED [100%]\n======================= 72 passed, 60 warnings in 44.30s ========================",
 "elapsed_s": 215
}

B200 + CuTeDSL 4.6.2 (previously FAILED on GB200)

{
 "branch": "mlir-compat-gateway",
 "gpu": "B200",
 "cutlass_spec": "nvidia-cutlass-dsl>=4.4.2,<4.7,!=4.5.0",
 "steps": {"clone": "ok", "build": "ok", "probe": "ok", "pytest": "ok"},
 "env": "2.9.1+cu129 4.6.2 NVIDIA B200",
 "pytest_rc": 0,
 "pytest_tail": "tests/test_ptx_umma_ws.py::test_ws_ss_tf32 PASSED [97%]\ntests/test_ptx_umma_ws.py::test_ws_ss_f16 PASSED [98%]\ntests/test_ptx_umma_ws.py::test_ws_ss_tf32_collector PASSED [100%]\n======================= 72 passed, 56 warnings in 30.37s ========================",
 "elapsed_s": 192
}

H100 + CuTeDSL 4.6.2 (SM90 coverage; SM100 tests self-skip via conftest)

67 passed, 5 deselected in 29.13s — compat 18/18, SM90 prefill + decode green, pytest_rc: 0.

Headless suite: tests/test_cutedsl_compat.py 18/18 locally; ruff lint + format clean (CI runs prek run ruff --all-files).

[bikrammajhi]

Quick status on the pushed head 159d972:

• CI lint fix: the pre-commit ruff hook (v0.15.0, --fix) was rewriting import order in the migrated modules (isort first-party cula grouping). Applied its own fix — pure import reordering, no semantic change. The re-run needs the usual maintainer approval.
• Re-validated the new head on Modal (same harness as before):
  • B200 + CuTeDSL 4.6.2: 72/72 passed (incl. test_ws_ss_tf32, test_ws_ss_f16, test_ws_ss_tf32_collector)
  • B200 + CuTeDSL 4.5.2: 72/72 passed (incl. test_ws_ss_tf32)
  • H100 + CuTeDSL 4.6.2: 67 passed, 5 SM100-only tests self-skipped

[icavan]

LGTM

[fkuner]

/LGTM