From e8ce568509dc62f3d127ecc3260cb32150e7d862 Mon Sep 17 00:00:00 2001 From: bikrammajhi Date: Sun, 9 Aug 2026 17:32:28 +0000 Subject: [PATCH 1/3] feat: isolate cutlass._mlir imports behind _mlir_compat gateway (#118) All 10 kernel modules previously imported CuTeDSL's private generated bindings (cutlass._mlir.{ir,arith,llvm,vector,nvvm,cute}) directly, so a CutDSL patch release could break kernel emission silently (the 4.5.2->4.5.3 tcgen05_ld/st incident). Add a single-point gateway that lazily loads the private dialects, enforces the pyproject.toml version contract (including the !=4.5.0 exclusion), and probes canary entry points, failing fast with an actionable error. Migrate all consumers to bind their dialect aliases from the gateway. Extend test_cutedsl_compat.py with fault-injection and version-matrix tests that run headless. --- cula/lightning/la_verify_kvbuffer.py | 4 +- cula/ops/_mlir_compat.py | 232 ++++++++++++++++++++++ cula/ops/kda/decode/mtp_conv.py | 2 +- cula/ops/kda/decode/mtp_kvbuffer.py | 4 +- cula/ops/kda/sm100/delta_h.py | 2 +- cula/ops/kda/sm90/_common.py | 2 +- cula/ops/lightning/prefill_sm100.py | 2 +- cula/ops/lightning/sm90/prefill_kernel.py | 4 +- cula/ops/lightning/sm90/schedule.py | 2 +- cula/ops/sm100/ptx.py | 8 +- tests/test_cutedsl_compat.py | 174 ++++++++++++++++ 11 files changed, 421 insertions(+), 15 deletions(-) create mode 100644 cula/ops/_mlir_compat.py diff --git a/cula/lightning/la_verify_kvbuffer.py b/cula/lightning/la_verify_kvbuffer.py index 4e6d653..75ffccf 100644 --- a/cula/lightning/la_verify_kvbuffer.py +++ b/cula/lightning/la_verify_kvbuffer.py @@ -49,8 +49,8 @@ import cutlass import cutlass.cute as cute import torch -from cutlass._mlir.dialects import arith as _arith -from cutlass._mlir.dialects import llvm as _llvm +from cula.ops._mlir_compat import arith as _arith +from cula.ops._mlir_compat import llvm as _llvm from cutlass.cute.runtime import ( make_fake_compact_tensor, make_fake_stream, diff --git a/cula/ops/_mlir_compat.py b/cula/ops/_mlir_compat.py new file mode 100644 index 0000000..832e980 --- /dev/null +++ b/cula/ops/_mlir_compat.py @@ -0,0 +1,232 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Single-point gateway for CuTeDSL's private MLIR/NVVM bindings. + +cuLA kernels are written against CuTeDSL's code-generation API. Most of that API +is public (``cutlass.cutlass_dsl``, ``cutlass.cute``, ...); a small part is not: +the generated MLIR dialect bindings under ``cutlass._mlir``. CuTeDSL ships them +as implementation detail and provides no stability contract across patch +releases -- the ``tcgen05_ld/st`` breakage between CutDSL 4.5.2 and 4.5.3 was one +such incident. + +This module is the ONLY place in cuLA that may import from ``cutlass._mlir``. +Kernel modules bind their dialect aliases from here:: + + from cula.ops._mlir_compat import arith as _arith + from cula.ops._mlir_compat import ir + from cula.ops._mlir_compat import llvm as _llvm + from cula.ops._mlir_compat import vector as _vector + +Design goals: + +- Lazy: nothing is imported until a kernel actually needs a binding, so plain + imports of cuLA never touch ``cutlass._mlir``. +- Explicity: an out-of-contract CuTeDSL version, or a missing/renamed binding, + fails fast at first use with an actionable ``RuntimeError`` instead of + surfacing mid-JIT as a confusing compile error (or silently emitting a + different kernel). +- Zero dependencies: version parsing is done with a small regex so this module + stays usable in any environment cuLA can be installed into. + +The version contract mirrors ``pyproject.toml`` (including its ``!=4.5.0`` +exclusion) and the canary probes make a broken binding fail fast with an +actionable message instead of surfacing mid-JIT. +""" + +from __future__ import annotations + +import importlib +import re +import sys +from typing import Any, Final, Optional + +# Version contract for nvidia-cutlass-dsl. Kept in sync with +# ``pyproject.toml``; when CuTeDSL is bumped, extend ``_SUPPORTED_MIN`` / +# ``_SUPPORTED_MAX`` only after the new release has been validated against the +# canaries below (and ideally against the SM90/SM100 kernel test suites). +_SUPPORTED_MIN: Final[tuple[int, ...]] = (4, 4, 2) +_SUPPORTED_MAX: Final[tuple[int, ...]] = (4, 7, 0) +_EXCLUDED_VERSIONS: Final[frozenset[tuple[int, ...]]] = frozenset({(4, 5, 0)}) + +# dialect name -> (package, attribute) inside ``cutlass``. +_PRIVATE_TABLE: Final[dict[str, tuple[str, str]]] = { + "arith": ("cutlass._mlir", "dialects.arith"), + "cute": ("cutlass._mlir", "dialects.cute"), + "ir": ("cutlass._mlir", "ir"), + "llvm": ("cutlass._mlir", "dialects.llvm"), + "nvvm": ("cutlass._mlir", "dialects.nvvm"), + "vector": ("cutlass._mlir", "dialects.vector"), +} + +# Canary entry points: attribute paths that must be present on each dialect +# binding for cuLA's kernel code to be emitted correctly. These are the names +# the migrated consumers actually call; extend the list when new usages land. +_CANARIES: Final[dict[str, tuple[tuple[str, ...], ...]]] = { + "arith": (("constant",),), + "cute": (), + "ir": (("Type", "parse"), ("VectorType", "get")), + "llvm": (("inline_asm",), ("extractvalue",)), + "nvvm": (), + "vector": (("bitcast",), ("extract_strided_slice",)), +} + +# Entry points whose name changed across CutDSL versions: for each group (a +# tuple of variant paths), at least one variant must exist. For example +# ``vector.extractelement`` was replaced by ``vector.extract`` in the 4.6 +# line; cuLA helpers dispatch on whichever one exists. +_ANY_OF_CANARIES: Final[dict[str, tuple[tuple[tuple[str, ...], ...], ...]]] = { + "vector": ((("extract",), ("extractelement",)),), +} + +_VERSION_RE: Final[re.Pattern[str]] = re.compile(r"^(\d+)(?:\.(\d+))?(?:\.(\d+))?(?:[a-z0-9._+-]*)?$") + +_INCIDENT_NOTE: Final[str] = ( + "cuLA depends on CuTeDSL's private MLIR bindings (`cutlass._mlir`), which are " + "implementation detail and were broken by a CuTeDSL patch release once " + "already (`tcgen05_ld/st`, 4.5.2 -> 4.5.3). To avoid silent kernel " + "miscompiles, cuLA refuses bindings outside its validated contract." +) + +_CACHE: Final[dict[str, Any]] = {} + + +def _parse_version(version: str) -> tuple[int, ...]: + """Parse ``X.Y.Z[.devN...]`` into a comparable tuple, or raise.""" + match = _VERSION_RE.match(version.strip()) + if match is None: + raise RuntimeError( + f"cuLA cannot parse the installed CuTeDSL version {version!r}; " + f"refusing to use its private MLIR bindings. {_INCIDENT_NOTE}" + ) + return tuple(int(part) for part in match.groups() if part is not None) + + +def _installed_version() -> tuple[int, ...]: + try: + import cutlass # noqa: PLC0415 + except ImportError: + raise RuntimeError( + "cuLA requires the nvidia-cutlass-dsl package; install it with " + "`pip install 'nvidia-cutlass-dsl>=4.4.2,<4.7'`." + ) from None + version = getattr(cutlass, "__version__", None) + if not isinstance(version, str): + raise RuntimeError( + f"CuTeDSL is installed but exposes no `__version__` (got {version!r}); " + f"refusing to use its private MLIR bindings. {_INCIDENT_NOTE}" + ) + return _parse_version(version) + + +def _check_contract(version: tuple[int, ...]) -> None: + if version in _EXCLUDED_VERSIONS: + raise RuntimeError( + f"Installed CuTeDSL version {'.'.join(map(str, version))} is explicitly " + f"excluded by cuLA (see pyproject.toml). {_INCIDENT_NOTE}" + ) + if not (_SUPPORTED_MIN <= version < _SUPPORTED_MAX): + raise RuntimeError( + f"Installed CuTeDSL version {'.'.join(map(str, version))} is outside the " + f"range validated by cuLA ({'.'.join(map(str, _SUPPORTED_MIN))} to " + f"{'.'.join(map(str, _SUPPORTED_MAX))}, exclusive). " + f"{_INCIDENT_NOTE} To proceed, pin the validated range in " + f"pyproject.toml and re-validate the canaries in this module." + ) + + +def _has_attribute_path(module: Any, path: tuple[str, ...]) -> bool: + owner = module + for part in path: + owner = getattr(owner, part, None) + if owner is None: + return False + return True + + +def _load(dialect: str) -> Any: + if dialect in _CACHE: + return _CACHE[dialect] + + package, attribute = _PRIVATE_TABLE[dialect] + try: + module = importlib.import_module(package) + except ImportError as exc: + raise RuntimeError( + f"Unable to import CuTeDSL's private {dialect!r} bindings " + f"({package}): {exc}. {_INCIDENT_NOTE}" + ) from exc + if attribute: + for part in attribute.split("."): + module = getattr(module, part, None) + if module is None: + break + if module is None: + raise RuntimeError( + f"CuTeDSL no longer exposes {dialect!r} bindings " + f"({package}.{attribute}). {_INCIDENT_NOTE}" + ) + + for canary in _CANARIES[dialect]: + owner: Any = module + for part in canary: + owner = getattr(owner, part, None) + if owner is None: + raise RuntimeError( + f"CuTeDSL dialect {dialect!r} is missing the canary entry point " + f"{'.'.join(canary)} used by cuLA kernels. {_INCIDENT_NOTE}" + ) + for group in _ANY_OF_CANARIES.get(dialect, ()): + if not any(_has_attribute_path(module, variant) for variant in group): + raise RuntimeError( + f"CuTeDSL dialect {dialect!r} exposes none of the entry points " + f"{' / '.join('.'.join(variant) for variant in group)} expected by " + f"cuLA kernels. {_INCIDENT_NOTE}" + ) + + _CACHE[dialect] = module + return module + + +def __getattr__(name: str) -> Any: + """Lazy, contract-checked access to private dialect bindings.""" + if name not in _PRIVATE_TABLE: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + version = _installed_version() + _check_contract(version) + return _load(name) + + +def cutlass_dsl_version() -> Optional[str]: + """Installed CuTeDSL version string, or None when not installed.""" + try: + import cutlass # noqa: PLC0415 + except ImportError: + return None + version = getattr(cutlass, "__version__", None) + return version if isinstance(version, str) else None + + +def vector_extract_element(vec, position, *, loc=None, ip=None): + """Extract one element of ``vec`` at ``position``, across CutDSL versions. + + CutDSL renamed ``vector.extractelement`` to ``vector.extract`` in the 4.6 + line (with the position as the second positional argument). Rather than + calling either binding directly from kernel code, kernels go through this + helper so a CutDSL patch release can never break the extraction path again. + """ + vector_dialect = _load("vector") + if _has_attribute_path(vector_dialect, ("extract",)): + return vector_dialect.extract(vec, position, None, loc=loc, ip=ip) + return vector_dialect.extractelement(vec, position=position, loc=loc, ip=ip) \ No newline at end of file diff --git a/cula/ops/kda/decode/mtp_conv.py b/cula/ops/kda/decode/mtp_conv.py index cf2ce93..f989551 100644 --- a/cula/ops/kda/decode/mtp_conv.py +++ b/cula/ops/kda/decode/mtp_conv.py @@ -33,7 +33,7 @@ import cutlass import cutlass.cute as cute import torch -from cutlass._mlir.dialects import llvm as _llvm +from cula.ops._mlir_compat import llvm as _llvm from cutlass.cute.runtime import from_dlpack from cutlass.cute.typing import Int32 from cutlass.cutlass_dsl import T as _T diff --git a/cula/ops/kda/decode/mtp_kvbuffer.py b/cula/ops/kda/decode/mtp_kvbuffer.py index 56af1cb..17a1ed9 100644 --- a/cula/ops/kda/decode/mtp_kvbuffer.py +++ b/cula/ops/kda/decode/mtp_kvbuffer.py @@ -1097,8 +1097,8 @@ def kda_decode_mtp_shuffle_kvbuffer( # C/D [16,8] f32: c0=C[gid][2tig] c1=C[gid][2tig+1] c2=C[gid+8][2tig] c3=C[gid+8][2tig+1] # =========================================================================== -from cutlass._mlir.dialects import arith as _arith # noqa: E402 -from cutlass._mlir.dialects import llvm as _llvm # noqa: E402 +from cula.ops._mlir_compat import arith as _arith # noqa: E402 +from cula.ops._mlir_compat import llvm as _llvm # noqa: E402 from cutlass.cutlass_dsl import T as _T # noqa: E402 from cutlass.cutlass_dsl import dsl_user_op # noqa: E402 diff --git a/cula/ops/kda/sm100/delta_h.py b/cula/ops/kda/sm100/delta_h.py index dd0ca91..846a410 100644 --- a/cula/ops/kda/sm100/delta_h.py +++ b/cula/ops/kda/sm100/delta_h.py @@ -27,7 +27,7 @@ import torch import torch.nn.functional as F import triton -from cutlass._mlir.dialects import llvm as _llvm +from cula.ops._mlir_compat import llvm as _llvm from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream from cutlass.cute.typing import Float32, Int32, Int64 diff --git a/cula/ops/kda/sm90/_common.py b/cula/ops/kda/sm90/_common.py index 8044a4c..c34531c 100644 --- a/cula/ops/kda/sm90/_common.py +++ b/cula/ops/kda/sm90/_common.py @@ -6,7 +6,7 @@ import cutlass import torch from cutlass import Int32 -from cutlass._mlir.dialects import llvm as _llvm +from cula.ops._mlir_compat import llvm as _llvm from cutlass.cutlass_dsl import T as _T diff --git a/cula/ops/lightning/prefill_sm100.py b/cula/ops/lightning/prefill_sm100.py index 93f81f3..8c99eeb 100644 --- a/cula/ops/lightning/prefill_sm100.py +++ b/cula/ops/lightning/prefill_sm100.py @@ -58,7 +58,7 @@ import cutlass.utils as utils import cutlass.utils.blackwell_helpers as sm100_utils import torch -from cutlass._mlir.dialects import llvm as _llvm +from cula.ops._mlir_compat import llvm as _llvm from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream from cutlass.cute.typing import Float32, Int32, Int64 diff --git a/cula/ops/lightning/sm90/prefill_kernel.py b/cula/ops/lightning/sm90/prefill_kernel.py index b43e7e5..a3063eb 100644 --- a/cula/ops/lightning/sm90/prefill_kernel.py +++ b/cula/ops/lightning/sm90/prefill_kernel.py @@ -25,13 +25,13 @@ import cuda.bindings.driver as cuda import cutlass -import cutlass._mlir.dialects.cute as _cute_ir import cutlass.cute as cute import cutlass.cute.nvgpu.warpgroup as warpgroup import cutlass.pipeline as pipeline import cutlass.utils as utils import cutlass.utils.hopper_helpers as sm90_utils -from cutlass._mlir.dialects import llvm +from cula.ops._mlir_compat import cute as _cute_ir +from cula.ops._mlir_compat import llvm from cutlass.cute.nvgpu import cpasync, warp from cutlass.utils.tensormap_manager import TensorMapManager, TensorMapUpdateMode diff --git a/cula/ops/lightning/sm90/schedule.py b/cula/ops/lightning/sm90/schedule.py index 5d9906d..295e046 100644 --- a/cula/ops/lightning/sm90/schedule.py +++ b/cula/ops/lightning/sm90/schedule.py @@ -24,7 +24,7 @@ import cutlass import cutlass.cute as cute -from cutlass._mlir.dialects import llvm +from cula.ops._mlir_compat import llvm from cutlass.cutlass_dsl import T TARGET_ARCH = "sm_90a" diff --git a/cula/ops/sm100/ptx.py b/cula/ops/sm100/ptx.py index 2823b23..b76cb48 100644 --- a/cula/ops/sm100/ptx.py +++ b/cula/ops/sm100/ptx.py @@ -55,10 +55,10 @@ import cutlass import cutlass.cute as cute -from cutlass._mlir import ir -from cutlass._mlir.dialects import arith as _arith -from cutlass._mlir.dialects import llvm -from cutlass._mlir.dialects import nvvm as _nvvm +from cula.ops._mlir_compat import arith as _arith +from cula.ops._mlir_compat import ir +from cula.ops._mlir_compat import llvm +from cula.ops._mlir_compat import nvvm as _nvvm from cutlass.cute.arch import elect_one from cutlass.cute.nvgpu import tcgen05 from cutlass.cute.typing import Int32 diff --git a/tests/test_cutedsl_compat.py b/tests/test_cutedsl_compat.py index 1bca3ec..8b0a50e 100644 --- a/tests/test_cutedsl_compat.py +++ b/tests/test_cutedsl_compat.py @@ -63,3 +63,177 @@ def unsupported_st(shape, tmem_addr, value): with pytest.raises(RuntimeError, match="expected exactly one value keyword"): detect_tcgen05_ldst_api(_inferred_ld, unsupported_st) + + +# --------------------------------------------------------------------------- +# MLIR compat gateway (_mlir_compat) +# --------------------------------------------------------------------------- + +import sys +import types + +from cula.ops import _mlir_compat +from cula.ops._mlir_compat import _parse_version + +_FAKE_DIALECTS = { + "arith": (("constant",),), + "cute": (), + "ir": (("Type", "parse"), ("VectorType", "get")), + "llvm": (("inline_asm",), ("extractvalue",)), + "nvvm": (), + "vector": (("bitcast",), ("extractelement",), ("extract_strided_slice",)), +} + + +def _install_fake_cutlass(version, *, with_mlir=True, missing_dialect=None, broken_canary=None): + """Install a fake ``cutlass`` package into ``sys.modules`` for fault injection. + + :param version: value for ``cutlass.__version__`` + :param with_mlir: when False, the fake exposes no ``_mlir`` subpackage + :param missing_dialect: dialect that does not exist in the fake package + :param broken_canary: dialect whose first canary entry point is absent + """ + installed = ["cutlass"] + cutlass = types.ModuleType("cutlass") + cutlass.__version__ = version + cutlass.__path__ = [] + if with_mlir: + _mlir = types.ModuleType("cutlass._mlir") + cutlass._mlir = _mlir + installed.append("cutlass._mlir") + dialects = types.ModuleType("cutlass._mlir.dialects") + _mlir.dialects = dialects + installed.append("cutlass._mlir.dialects") + for name, canaries in _FAKE_DIALECTS.items(): + if name == missing_dialect: + continue + module = types.ModuleType(f"cutlass._mlir.dialects.{name}") + setattr(module, "non_public", types.SimpleNamespace()) + chains = canaries if name != broken_canary else () + for chain in chains: + owner = module + for index, part in enumerate(chain): + is_leaf = index == len(chain) - 1 + if not hasattr(owner, part): + if is_leaf: + setattr(owner, part, (lambda *args, _op=part, **kwargs: f"op:{_op}")) + else: + setattr(owner, part, types.SimpleNamespace()) + owner = getattr(owner, part) + setattr(dialects, name, module) + installed.append(f"cutlass._mlir.dialects.{name}") + + previous = {name: sys.modules.get(name) for name in installed} + for name in installed: + sys.modules[name] = _module_by_name(cutlass, name) + _mlir_compat._CACHE.clear() + return lambda: _restore_modules(previous) + + +def _module_by_name(root, dotted): + module = root + for part in dotted.split(".")[1:]: + module = getattr(module, part) + return module + + +def _restore_modules(previous): + for name, module in previous.items(): + if module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = module + _mlir_compat._CACHE.clear() + + +def test_parse_version_accepts_release_and_dev_suffixes(): + assert _parse_version("4.5.3") == (4, 5, 3) + assert _parse_version("4.6.0.dev0") == (4, 6, 0) + + +def test_parse_version_rejects_garbage(): + with pytest.raises(RuntimeError, match="cannot parse"): + _parse_version("not-a-version") + + +def test_version_contract_rejects_out_of_range(): + restore = _install_fake_cutlass("9.9.9") + try: + with pytest.raises(RuntimeError, match="outside the range validated by cuLA"): + _mlir_compat.llvm + finally: + restore() + + +def test_missing_cutlass_reports_install_hint(): + restore = _install_fake_cutlass("4.5.3", with_mlir=False) + sys.modules["cutlass"] = None # simulate the package being absent + try: + with pytest.raises(RuntimeError, match="nvidia-cutlass-dsl"): + _mlir_compat.llvm + finally: + restore() + + +def test_missing_dialect_raises_with_dialect_name(): + restore = _install_fake_cutlass("4.5.3", missing_dialect="llvm") + try: + with pytest.raises(RuntimeError, match="no longer exposes 'llvm'"): + _mlir_compat.llvm + finally: + restore() + + +def test_missing_canary_raises_naming_the_entry_point(): + restore = _install_fake_cutlass("4.5.3", broken_canary="llvm") + try: + with pytest.raises(RuntimeError, match="missing the canary entry point inline_asm"): + _mlir_compat.llvm + finally: + restore() + + +def test_vector_dispatch_uses_extractelement_when_extract_absent(): + restore = _install_fake_cutlass("4.5.3") + try: + assert _mlir_compat.vector_extract_element("vec", "pos") == "op:extractelement" + finally: + restore() + + +def test_vector_dispatch_uses_extract_when_available(): + restore = _install_fake_cutlass("4.6.0") + try: + vector_dialect = _mlir_compat._load("vector") + # simulate the 4.6 rename: drop extractelement, add extract + delattr(vector_dialect, "extractelement") + setattr(vector_dialect, "extract", lambda *args, **kwargs: "op:extract") + _mlir_compat._CACHE.clear() + assert _mlir_compat.vector_extract_element("vec", "pos") == "op:extract" + finally: + restore() + + +def test_unknown_attribute_raises_attribute_error(): + with pytest.raises(AttributeError, match="has no attribute"): + _mlir_compat.does_not_exist + + +def test_dialect_bindings_are_cached(): + restore = _install_fake_cutlass("4.5.3") + try: + assert _mlir_compat.llvm is _mlir_compat.llvm + finally: + restore() + + +def test_real_cutlass_bindings_load(): + """Smoke-test the gateway against the installed CuTeDSL wheel.""" + try: + import cutlass # noqa: F401 + except ImportError: + pytest.skip("nvidia-cutlass-dsl is not installed") + + for dialect in ("arith", "cute", "ir", "llvm", "nvvm", "vector"): + _mlir_compat._CACHE.clear() + assert getattr(_mlir_compat, dialect) is not None From aed42098a9d23775849131d098ff2c9e10e4bf58 Mon Sep 17 00:00:00 2001 From: bikrammajhi Date: Sun, 9 Aug 2026 17:32:28 +0000 Subject: [PATCH 2/3] fix: cross-version vector element extraction in store_256b CutDSL renamed vector.extractelement to vector.extract in the 4.6 line, which broke store_256b (used by KDA SM100 backward) against nvidia-cutlass-dsl 4.6.2 at JIT time. Route element extraction through the gateway's version-dispatching vector_extract_element helper. --- cula/ops/ptx.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/cula/ops/ptx.py b/cula/ops/ptx.py index 595538f..7040a05 100644 --- a/cula/ops/ptx.py +++ b/cula/ops/ptx.py @@ -16,10 +16,11 @@ import cutlass import cutlass.cute as cute -from cutlass._mlir import ir -from cutlass._mlir.dialects import arith as _arith -from cutlass._mlir.dialects import llvm as _llvm -from cutlass._mlir.dialects import vector as _vector +from cula.ops._mlir_compat import arith as _arith +from cula.ops._mlir_compat import ir +from cula.ops._mlir_compat import llvm as _llvm +from cula.ops._mlir_compat import vector as _vector +from cula.ops._mlir_compat import vector_extract_element from cutlass.cutlass_dsl import T as _T from cutlass.cutlass_dsl import dsl_user_op @@ -128,9 +129,9 @@ def _do(addr, v, *, loc=None, ip=None): i32_ty = ir.IntegerType.get_signless(32) ir_v = _to_ir(v, loc, ip) elems = [ - _vector.extractelement( + vector_extract_element( ir_v, - position=_arith.constant(i32_ty, i, loc=loc, ip=ip), + _arith.constant(i32_ty, i, loc=loc, ip=ip), loc=loc, ip=ip, ) From 2c16cb93abb49f30ba1ef9ff8b28201fa4d99e02 Mon Sep 17 00:00:00 2001 From: bikrammajhi Date: Sun, 9 Aug 2026 17:32:28 +0000 Subject: [PATCH 3/3] chore: add Modal GPU validation harness (H100, CUDA 12.9) --- scripts/modal_validate.py | 106 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 106 insertions(+) create mode 100644 scripts/modal_validate.py diff --git a/scripts/modal_validate.py b/scripts/modal_validate.py new file mode 100644 index 0000000..a682886 --- /dev/null +++ b/scripts/modal_validate.py @@ -0,0 +1,106 @@ +"""Modal validation harness for cuLA on H100 (modal client v1.4.2 API). + +Clones the fork branch inside the container (this client ships no Mount), +builds cuLA, runs the compat + SM90 kernel JIT tests, returns a JSON summary. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import time + +import modal + +FORK = "https://github.com/bikrammajhi/cuLA.git" +BRANCH = "mlir-compat-gateway" +CUDA_TAG = "cu129" +TORCH_VERSION = "2.9.1" +TESTS = "test_cutedsl_compat.py test_lightning_attn_prefill_sm90.py test_lightning_decode.py" + +image = ( + modal.Image.from_registry("nvidia/cuda:12.9.0-devel-ubuntu22.04", add_python="3.12") + .apt_install("git") + .pip_install("wheel", "setuptools", "setuptools-scm") + .pip_install( + f"torch=={TORCH_VERSION}", + index_url=f"https://download.pytorch.org/whl/{CUDA_TAG}", + ) + .pip_install("nvidia-cutlass-dsl>=4.4.2,<4.7,!=4.5.0", "flash-linear-attention", "pytest") +) + +app = modal.App("cula-validate-v4") + + +@app.function(image=image, gpu="h100", timeout=60 * 60) +def validate() -> str: + start = time.time() + summary = {"branch": BRANCH, "steps": {}} + + def step(name: str) -> None: + summary["steps"][name] = "ok" + + subprocess.run( + f"git clone --depth 1 --branch {BRANCH} {FORK} /work", + shell=True, + check=True, + ) + step("clone") + + check = subprocess.run( + "ls tests/ | head -25; echo ---; git log --oneline -3", + shell=True, + cwd="/work", + capture_output=True, + text=True, + ) + summary["clone_check"] = check.stdout + check.stderr + + toolchain = {**os.environ, "CC": "gcc", "CXX": "g++", "CUDAHOSTCXX": "g++"} + subprocess.run( + "pip install -e /work --no-build-isolation", + shell=True, + env=toolchain, + check=True, + ) + step("build") + + probe = subprocess.run( + "python -c 'import torch, cutlass, cula; " + "print(torch.__version__, cutlass.__version__, torch.cuda.get_device_name(0))'", + shell=True, + capture_output=True, + text=True, + ) + summary["env"] = probe.stdout.strip() + step("probe") + + tests = tests or "test_cutedsl_compat.py test_lightning_attn_prefill_sm90.py test_lightning_decode.py" + run = subprocess.run( + f"python -m pytest tests/{tests} -v -m 'not sm100_only'", + shell=True, + cwd="/work", + capture_output=True, + text=True, + ) + summary["pytest_cmd"] = f"pytest tests/{tests} -v -m not_sm100_only (cwd=/work)" + ls = subprocess.run( + "ls -la tests/test_lightning_attn_prefill_sm90.py tests/test_cutedsl_compat.py tests/test_lightning_decode.py", + shell=True, + cwd="/work", + capture_output=True, + text=True, + ) + summary["target_ls"] = ls.stdout + ls.stderr + summary["pytest_rc"] = run.returncode + summary["pytest_tail"] = (run.stdout + run.stderr)[-5000:] + step("pytest" if run.returncode == 0 else "pytest_FAILED") + + summary["elapsed_s"] = int(time.time() - start) + return json.dumps(summary, indent=1) + + +@app.local_entrypoint() +def main() -> None: + print(validate.remote()) \ No newline at end of file