ci(fa4): enforce cutlass-dsl/quack dep floors and rebake cu130 image

.github/workflows/ci.yml

@@ -45,4 +45,4 @@ jobs:
         with:
           test-filter: ${{ env.FA4_TEST_FILTER }}
           fa4_image_cu129: "togethercomputer/training-performance:flash-attn-cu12.9-26.03.25@sha256:304a5c3d2b3a75b151cd2a964cd26d444e0d8b5686d63943df13378c9705f943"
-          fa4_image_cu130: "togethercomputer/training-performance:flash-attn-cu13.0-26.04.01@sha256:56e50b056eb4d671410846c3483e843ee7bd0f5b13cb45b6f0d7eb8bd27694a5"
+          fa4_image_cu130: "togethercomputer/training-performance:flash-attn-cu13.0-26.06.10@sha256:f1efd03b9d78cf65d9f8df107d2f6f6d0a464cb8b773fd9364765f22f4772006"

tools/ci/assert_dsl_floor.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python3
+"""Fail loudly if the CI image's deps are below the flash_attn/cute/pyproject.toml floors.
+
+Runs inside the SIF before tests. The FA4 install in run_fa4_ci.py uses --no-deps (to keep the
+SIF's torch/cudnn), so pyproject floors are not enforced at install time. A SIF baked before a
+floor bump therefore keeps a stale dep — e.g. nvidia-cutlass-dsl 4.4.2, which can't convert the
+AuxData JIT arg and dies with a cryptic DSLRuntimeError deep in SM100 kernel launch. This check
+turns that into an actionable "rebake the image" message up front.
+
+Reads the floor from pyproject so there is no hardcoded version here to drift out of sync.
+"""
+
+from __future__ import annotations
+
+import sys
+from importlib.metadata import PackageNotFoundError, version
+
+from packaging.requirements import Requirement
+from packaging.version import Version
+
+try:
+    import tomllib  # Python 3.11+
+except ModuleNotFoundError:  # Python 3.10 (pyproject declares requires-python >=3.10)
+    try:
+        import tomli as tomllib
+    except ModuleNotFoundError:
+        sys.exit(
+            "ERROR: assert_dsl_floor.py needs a TOML parser — use Python 3.11+ (stdlib tomllib) "
+            "or `pip install tomli` on 3.10."
+        )
+
+# Deps whose floor a stale SIF is known to silently violate. Other pyproject deps (torch, einops…)
+# are baked to match the image and not version-sensitive in the same way, so we don't gate on them.
+CHECKED = ("nvidia-cutlass-dsl", "quack-kernels")
+
+
+def main(pyproject_path: str) -> int:
+    with open(pyproject_path, "rb") as f:
+        deps = tomllib.load(f)["project"]["dependencies"]
+    reqs = {r.name: r for r in (Requirement(d) for d in deps) if r.name in CHECKED}
+
+    failures: list[str] = []
+    oks: list[str] = []
+    for name in CHECKED:
+        req = reqs.get(name)
+        if req is None:
+            continue  # not a hard dep in this pyproject — nothing to enforce
+        try:
+            installed = version(name)
+        except PackageNotFoundError:
+            failures.append(f"{name}: not installed (floor {req.specifier})")
+            continue
+        if req.specifier.contains(Version(installed), prereleases=True):
+            oks.append(f"{name}={installed}")
+        else:
+            failures.append(f"{name}: installed {installed} does not satisfy floor {req.specifier}")
+
+    if failures:
+        print("ERROR: CI image deps are below the flash_attn/cute/pyproject.toml floor:", file=sys.stderr)
+        for line in failures:
+            print(f"  - {line}", file=sys.stderr)
+        print(
+            "\nThe SIF was likely baked before a floor bump. Rebake the image "
+            "(tools/ci/docker/build.sh + tag_and_push.sh) and update the digest in "
+            ".github/workflows/ci.yml.",
+            file=sys.stderr,
+        )
+        return 1
+
+    print("DSL floor check OK: " + ", ".join(oks))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1] if len(sys.argv) > 1 else "flash_attn/cute/pyproject.toml"))

tools/ci/run_fa4_ci.py

@@ -106,9 +106,21 @@ def run_step(step: Step, repo_root: Path, base_env: dict[str, str], sif: str, wo
     print(f"=== {step.name} ===")
 
     # Install FA4 from the current repo inside this exec invocation.
+    # --no-deps keeps the SIF's baked torch/cudnn; deps are expected to already satisfy the
+    # pyproject floors via the image (rebuilt with tools/ci/docker/build.sh). We do NOT upgrade
+    # deps here: the --writable-tmpfs overlay is RAM-backed and too small to hold a cutlass-dsl
+    # reinstall (it ENOSPCs and can corrupt the baked torch). Instead assert_dsl_floor.py below
+    # fails loudly if the image is stale, pointing at a rebake.
     # Must be done per-step because --writable-tmpfs creates a fresh overlay each time.
     install_cmd = f"uv pip install --system --break-system-packages --no-deps -q -e {shlex.quote(str(repo_root / 'flash_attn/cute'))}"
 
+    # Guard against a SIF baked with deps below the pyproject floor (the silent --no-deps gap that
+    # otherwise surfaces as a cryptic DSLRuntimeError on the SM100 path). Cheap: reads versions, no install.
+    floor_check_cmd = (
+        f"python3 {shlex.quote(str(repo_root / 'tools/ci/assert_dsl_floor.py'))} "
+        f"{shlex.quote(str(repo_root / 'flash_attn/cute/pyproject.toml'))}"
+    )
+
     # Convert relative test/benchmark paths to absolute so we can run from /tmp.
     # Running from /tmp ensures Python does not insert repo_root into sys.path[0]
     # (which would cause flash_attn/__init__.py to trigger FA2 imports unavailable in the SIF).
@@ -118,7 +130,7 @@ def run_step(step: Step, repo_root: Path, base_env: dict[str, str], sif: str, wo
     ]
     env_exports = " && ".join(f"export {k}={shlex.quote(v)}" for k, v in step.extra_env.items())
     inner_cmd = shlex.join(command)
-    shell_parts = [install_cmd]
+    shell_parts = [install_cmd, floor_check_cmd]
     if env_exports:
         shell_parts.append(env_exports)
     shell_parts.append(f"cd /tmp && {inner_cmd}")