sparse-delta-memory 2.0.0__py3-none-any.whl

sparse_delta_memory/cuda/_build.py

@@ -0,0 +1,254 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+"""_build: lazy, content-addressed JIT builds of the package's CUDA extensions.
+
+The build directory and the imported module are keyed only by the extension name. A
+process that loses the build lock imports whatever shared object the winner produced
+without checking that it came from the same sources. Every extension here is therefore
+named ``<base>_<sha256(build inputs)[:16]>``, where the build inputs are the sources,
+the headers, the exported bindings, all compiler flags, the torch/CUDA/Python versions,
+the CUDA toolkit location and the target GPU architectures. Identical inputs share one
+cache entry; anything else gets its own directory, so a stale ``.so`` is never reused
+and two source trees can share one ``TORCH_EXTENSIONS_DIR``.
+
+Each entry of ``cuda_sources`` is its own translation unit, so ninja compiles them in
+parallel (``MAX_JOBS`` caps the job count). ``headers`` are written next to the sources,
+so any unit can ``#include "<name>"``. The toolkit's CCCL directory (CUB, Thrust,
+libcu++) is on the include path.
+
+Builds happen on the first call, never at import. Call each op once eagerly (for
+example in a warmup step) before capturing a CUDA graph so that no compilation is
+triggered during capture."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import shutil
+import sys
+import tempfile
+import threading
+import warnings
+from collections.abc import Mapping, Sequence
+from types import ModuleType
+
+import torch
+from torch.utils import cpp_extension
+
+from . import _image
+
+__all__ = ["load_cuda_extension"]
+
+_LOCK = threading.Lock()
+
+
+def _as_tuple(value: str | Sequence[str]) -> tuple[str, ...]:
+    return (value,) if isinstance(value, str) else tuple(value)
+
+
+def _require_toolchain() -> str:
+    """Return CUDA_HOME, or raise a readable error if the toolkit is unusable."""
+    if not torch.cuda.is_available() or torch.version.cuda is None:
+        raise RuntimeError(
+            "The SDM CUDA kernels need a CUDA build of PyTorch and a visible GPU."
+        )
+    cuda_home = cpp_extension.CUDA_HOME
+    if cuda_home is None:
+        raise RuntimeError(
+            "The SDM CUDA kernels are JIT-compiled and need the CUDA toolkit, but "
+            "PyTorch could not locate it. Set CUDA_HOME to a toolkit that matches "
+            f"torch.version.cuda={torch.version.cuda} (it must contain bin/nvcc)."
+        )
+    nvcc = os.environ.get("PYTORCH_NVCC") or os.path.join(cuda_home, "bin", "nvcc")
+    if not (os.path.exists(nvcc) or shutil.which(nvcc)):
+        raise RuntimeError(
+            f"nvcc was not found at {nvcc!r}. Set CUDA_HOME to a full CUDA toolkit "
+            f"matching torch.version.cuda={torch.version.cuda}."
+        )
+    return cuda_home
+
+
+def toolkit() -> tuple:
+    """The toolkit an extension build uses (cpp_extension.CUDA_HOME, PYTORCH_NVCC): part of _image's build environment."""
+    return cpp_extension.CUDA_HOME, os.environ.get("PYTORCH_NVCC")
+
+
+def _cccl_include_paths(cuda_home: str) -> tuple[str, ...]:
+    """CUDA 13 moved CUB/Thrust/libcu++ to ``include/cccl``; older toolkits keep them
+    directly under ``include``, which PyTorch already adds."""
+    path = os.path.join(cuda_home, "include", "cccl")
+    return (path,) if os.path.isdir(path) else ()
+
+
+def _target_architectures() -> dict[str, object]:
+    """The inputs PyTorch uses to choose the -gencode flags."""
+    capabilities = sorted(
+        {torch.cuda.get_device_capability(i) for i in range(torch.cuda.device_count())}
+    )
+    return {
+        "TORCH_CUDA_ARCH_LIST": os.environ.get("TORCH_CUDA_ARCH_LIST"),
+        "capabilities": [list(c) for c in capabilities],
+    }
+
+
+def _header_items(headers: Mapping[str, str] | None) -> tuple[tuple[str, str], ...]:
+    items = tuple(sorted((headers or {}).items()))
+    for file_name, _ in items:
+        if not file_name or os.path.basename(file_name) != file_name:
+            raise ValueError(f"header names must be plain file names, got {file_name!r}")
+    return items
+
+
+def _extension_name(
+    base: str,
+    *,
+    cpp_sources: str | Sequence[str],
+    cuda_sources: str | Sequence[str],
+    functions: Sequence[str],
+    headers: Mapping[str, str] | None = None,
+    extra_cflags: Sequence[str] = (),
+    extra_cuda_cflags: Sequence[str] = (),
+    extra_ldflags: Sequence[str] = (),
+    extra_include_paths: Sequence[str] = (),
+) -> str:
+    """Return ``<base>_<16 hex digits>`` identifying one exact build."""
+    if not base.isidentifier():
+        raise ValueError(f"extension base name must be an identifier, got {base!r}")
+    cuda_home = _require_toolchain()
+    identity = {
+        "sources": {"cpp": _as_tuple(cpp_sources), "cuda": _as_tuple(cuda_sources)},
+        "headers": _header_items(headers),
+        "functions": tuple(functions),
+        "flags": {
+            "cflags": tuple(extra_cflags),
+            "cuda_cflags": tuple(extra_cuda_cflags),
+            "ldflags": tuple(extra_ldflags),
+            "include_paths": tuple(extra_include_paths),
+        },
+        "torch": torch.__version__,
+        "torch_cuda": torch.version.cuda,
+        "python": sys.implementation.cache_tag,
+        "cuda_home": os.path.realpath(cuda_home),
+        "arch": _target_architectures(),
+    }
+    nvcc_env = {k: os.environ[k] for k in ("NVCC_APPEND_FLAGS", "NVCC_PREPEND_FLAGS") if os.environ.get(k)}
+    if nvcc_env:   # nvcc's environment flags shape the targets: a build under them is its own build (_image)
+        identity["nvcc_env"] = nvcc_env
+    payload = json.dumps(identity, sort_keys=True, separators=(",", ":"))
+    return f"{base}_{hashlib.sha256(payload.encode('utf-8')).hexdigest()[:16]}"
+
+
+def _warn_if_locked(build_dir: str) -> None:
+    """A build killed mid-compile leaves a lock file that makes later loads wait forever."""
+    lock = os.path.join(build_dir, "lock")
+    if os.path.exists(lock):
+        warnings.warn(
+            f"{lock} exists. If no other process is currently building this "
+            "extension, delete that file; otherwise this load will wait for it.",
+            stacklevel=3,
+        )
+
+
+def _write_if_changed(path: str, text: str) -> str:
+    """Write atomically, and leave an identical file untouched so ninja skips it."""
+    try:
+        with open(path, encoding="utf-8") as f:
+            if f.read() == text:
+                return path
+    except FileNotFoundError:
+        pass
+    fd, tmp = tempfile.mkstemp(dir=os.path.dirname(path), prefix=".tmp_")
+    with os.fdopen(fd, "w", encoding="utf-8") as f:
+        f.write(text)
+    os.replace(tmp, path)
+    return path
+
+
+def _bindings_source(cpp_sources: tuple[str, ...], functions: tuple[str, ...]) -> str:
+    """The C++ unit holding the declarations and the pybind11 module."""
+    defs = [f'  m.def("{f}", torch::wrap_pybind_function({f}), "{f}");' for f in functions]
+    return "\n".join(
+        ["#include <torch/extension.h>", *cpp_sources,
+         "PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {", *defs, "}", ""]
+    )
+
+
+def load_cuda_extension(
+    name: str,
+    *,
+    cpp_sources: str | Sequence[str],
+    cuda_sources: str | Sequence[str],
+    functions: Sequence[str],
+    headers: Mapping[str, str] | None = None,
+    extra_cflags: Sequence[str] = (),
+    extra_cuda_cflags: Sequence[str] = (),
+    extra_ldflags: Sequence[str] = (),
+    extra_include_paths: Sequence[str] = (),
+    verbose: bool = False,
+) -> ModuleType:
+    """Build (on first use) and import a JIT-compiled CUDA extension.
+
+    Args:
+        name: identifier prefix of the extension; the content hash is appended.
+        cpp_sources: C++ declarations of the bound functions.
+        cuda_sources: CUDA translation unit(s); each string is compiled separately.
+        functions: names to expose through pybind11; each must be declared in
+            ``cpp_sources``. The loaded module is checked to export all of them.
+        headers: ``{file name: text}`` of headers the units include as
+            ``#include "<file name>"``.
+        extra_cflags / extra_cuda_cflags / extra_ldflags / extra_include_paths:
+            forwarded to the compiler and linker and included in the content hash.
+        verbose: forwarded to ``torch.utils.cpp_extension.load``.
+
+    Returns the imported module. Callers cache it: each call hashes the inputs and loads
+    the extension.
+    """
+    cpp_sources, cuda_sources = _as_tuple(cpp_sources), _as_tuple(cuda_sources)
+    functions = tuple(functions)
+    header_items = _header_items(headers)
+    if not functions:
+        raise ValueError("load_cuda_extension needs at least one bound function")
+    if not cuda_sources:
+        raise ValueError("load_cuda_extension needs at least one CUDA source")
+    full_name = _extension_name(
+        name,
+        cpp_sources=cpp_sources,
+        cuda_sources=cuda_sources,
+        functions=functions,
+        headers=headers,
+        extra_cflags=extra_cflags,
+        extra_cuda_cflags=extra_cuda_cflags,
+        extra_ldflags=extra_ldflags,
+        extra_include_paths=extra_include_paths,
+    )
+    with _LOCK:
+        build_dir = cpp_extension._get_build_directory(full_name, verbose)
+        _warn_if_locked(build_dir)
+        for file_name, text in header_items:
+            _write_if_changed(os.path.join(build_dir, file_name), text)
+        sources = [_write_if_changed(os.path.join(build_dir, "main.cpp"),
+                                     _bindings_source(cpp_sources, functions))]
+        sources += [_write_if_changed(os.path.join(build_dir, f"unit{i}.cu"), text)
+                    for i, text in enumerate(cuda_sources)]
+        module = cpp_extension.load(
+            name=full_name,
+            sources=sources,
+            extra_cflags=list(extra_cflags),
+            extra_cuda_cflags=list(extra_cuda_cflags),
+            extra_ldflags=list(extra_ldflags),
+            extra_include_paths=[*_cccl_include_paths(cpp_extension.CUDA_HOME),
+                                 *extra_include_paths],
+            build_directory=build_dir,
+            verbose=verbose,
+            with_cuda=True,
+        )
+        missing = [f for f in functions if not callable(getattr(module, f, None))]
+        if missing:
+            raise RuntimeError(
+                f"CUDA extension {full_name!r} loaded from "
+                f"{getattr(module, '__file__', None)!r} does not export {missing}. "
+                "The cached build is incomplete; delete its directory and retry."
+            )
+    if name != _image.PROBE_BASE:
+        _image.require_supported_image()   # before any kernel of the extension runs
+    return module

sparse_delta_memory/cuda/_device.py

@@ -0,0 +1,158 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+"""_device: the routing predicate (backend(): which kernels and plans serve a device) and the device properties that the
+launch-shape rules scale with: SM count and L2 size, cached per device index.
+
+The launch-shape rules were tuned on GB300 (152 SMs, 135528448-byte L2), where they reduce to the tuned constants. They
+choose launch shapes only, never the arithmetic, so no output depends on them."""
+
+import contextlib
+import functools
+import os
+import threading
+
+import torch
+
+
+def _index(device) -> int:
+    """CUDA device index of `device` (None, an int, a string or a torch.device; unindexed = the current device)."""
+    if device is None:
+        return torch.cuda.current_device()
+    if isinstance(device, int):
+        return device
+    d = device if isinstance(device, torch.device) else torch.device(device)
+    return d.index if d.index is not None else torch.cuda.current_device()
+
+
+@functools.lru_cache(maxsize=None)
+def _props(index: int):
+    p = torch.cuda.get_device_properties(index)
+    return int(p.multi_processor_count), int(p.L2_cache_size)
+
+
+def num_sms(device=None) -> int:
+    """Streaming multiprocessors of `device` (default: the current device)."""
+    return _FORCED[2] if _FORCED else _props(_index(device))[0]
+
+
+def l2_bytes(device=None) -> int:
+    """L2 cache size in bytes of `device` (default: the current device)."""
+    return _FORCED[3] if _FORCED and _FORCED[3] is not None else _props(_index(device))[1]
+
+
+def prev_pow2(x: int) -> int:
+    """Largest power of two <= x (x >= 1)."""
+    return 1 << (int(x).bit_length() - 1)
+
+
+# The ONE routing predicate (the supported hardware matrix; no knob). "modern": H100 / B200 / GB300 (compute capability
+# 9.0 / 10.0 / 10.3 with at least 232,448 bytes of opt-in shared memory per block): the fast kernels as tuned (TMA, PDL,
+# clusters, hot/cold Phase 1). "a100": A100 (8.0, 166,912 bytes): the same kernels through the portable (cp.async) side of
+# their arch guards, no PDL, no clusters, the ordered Phase 1, A100 tile plans. "rtx": RTX 30 / 40 / 50 (8.6 / 8.9 / 12.0,
+# 101,376 bytes): the same portable code at RTX limits (a native sm_120 build compiles the same device branches as sm_86
+# / 89: sdm_copy.cuh). Every other target (below 8.0, or a capability or shared-memory limit not listed): one clear error
+# naming the supported GPUs. Read from torch's device properties only, so a test that patches them forces a row.
+MODERN_CAPS = ((9, 0), (10, 0), (10, 3))
+MODERN_SMEM = 232448
+A100_CAPS, A100_SMEM = ((8, 0),), 166912
+RTX_CAPS, RTX_SMEM = ((8, 6), (8, 9), (12, 0)), 101376
+UNSUPPORTED = ("sparse_delta_memory supports NVIDIA H100 / B200 / GB300 (compute capability 9.0 / 10.0 / 10.3), A100 (8.0) "
+               "and RTX 30 / 40 / 50 (8.6 / 8.9 / 12.0); device {i} ({name}) has compute capability {major}.{minor} with "
+               "{smem} bytes of opt-in shared memory per block")
+
+# The device-side mirror of backend() lives in sdm_copy.cuh (SDM_PORTABLE_ARCH: every device pass but sm_90 / 100 / 103),
+# with the portable copy layer; kernels that guard on the architecture include it (headers()).
+_COPY_HEADER = "sdm_copy.cuh"
+_PLAN_HEADER = "sdm_plancap.cuh"   # host-side launch capture of the backward's step plans (bwd_steps_cuda)
+
+
+@functools.cache
+def headers() -> dict:
+    """{file name: text} of the shared headers for load_cuda_extension(headers=...)."""
+    out = {}
+    for name in (_COPY_HEADER, _PLAN_HEADER):
+        with open(os.path.join(os.path.dirname(__file__), name), encoding="utf-8") as f:
+            out[name] = f.read()
+    return out
+
+
+def capability(device=None) -> tuple:
+    """Compute capability (major, minor) of `device`."""
+    return _FORCED[0] if _FORCED else tuple(torch.cuda.get_device_capability(_index(device)))
+
+
+def _optin(device=None) -> int:
+    """The device's own opt-in shared memory per block (bytes): the classification input of backend()."""
+    return _FORCED[1] if _FORCED else int(torch.cuda.get_device_properties(_index(device)).shared_memory_per_block_optin)
+
+
+def backend(device=None) -> str:
+    """'modern' | 'a100' | 'rtx' for `device` (see MODERN_CAPS); any other target raises RuntimeError(UNSUPPORTED)."""
+    i = _index(device)
+    cc, optin = capability(i), _optin(i)
+    if cc in MODERN_CAPS and optin >= MODERN_SMEM:
+        return "modern"
+    if cc in A100_CAPS and optin == A100_SMEM:
+        return "a100"
+    if cc in RTX_CAPS and optin == RTX_SMEM:
+        return "rtx"
+    raise RuntimeError(UNSUPPORTED.format(i=i, name=torch.cuda.get_device_name(i), major=cc[0], minor=cc[1], smem=optin))
+
+
+def cuda_backend(device=None) -> bool:
+    """The package's CUDA kernels serve `device`: True on every backend (elsewhere, backend()'s clear error)."""
+    return bool(backend(device))
+
+
+def pdl_ok(device=None) -> bool:
+    """Programmatic dependent launch may be requested on `device`: the modern backend only."""
+    return backend(device) == "modern"
+
+
+def tbind(kernel, grid, ref):
+    """kernel[grid] for an SDM Triton launch (`ref`: one of its CUDA operands)."""
+    return kernel[grid]
+
+
+# Shared memory per block every plan of a class assumes (modern: 232,448 on any modern part; A100 / RTX exact).
+CLASS_SMEM = {"modern": MODERN_SMEM, "a100": A100_SMEM, "rtx": RTX_SMEM}
+
+
+def smem_cap(device=None) -> int:
+    """Opt-in shared memory per block the plans of `device`'s class use (bytes; CLASS_SMEM)."""
+    return CLASS_SMEM[backend(device)]
+
+
+def smem_per_sm(device=None) -> int:
+    """Shared memory per SM of `device`'s class (bytes): smem_cap() plus the 1 KB reserved per block."""
+    return smem_cap(device) + RESERVED_SMEM
+
+
+RESERVED_SMEM = 1024   # shared memory the system reserves per resident block (sm_80+)
+
+# _force_row overrides capability(), _optin(), num_sms() and l2_bytes() for every device (process-wide; contexts nest).
+_FORCED = None
+_FORCE_STACK = []
+_FORCE_LOCK = threading.Lock()
+
+
+@contextlib.contextmanager
+def _force_row(capability, smem_optin, sm_count, l2_bytes=None):
+    """TEST ONLY: while active, every device reports this capability, opt-in shared memory, SM count and (given) L2, so a
+    test reproduces a hardware-matrix row on another GPU. Contexts must nest (an exit out of order raises)."""
+    global _FORCED
+    entry = (tuple(int(c) for c in capability), int(smem_optin), int(sm_count),
+             None if l2_bytes is None else int(l2_bytes))
+    token = object()
+    with _FORCE_LOCK:
+        _FORCE_STACK.append((token, entry))
+        _FORCED = entry
+    top = True
+    try:
+        yield
+    finally:
+        with _FORCE_LOCK:
+            top = bool(_FORCE_STACK) and _FORCE_STACK[-1][0] is token
+            _FORCE_STACK[:] = [e for e in _FORCE_STACK if e[0] is not token]
+            _FORCED = _FORCE_STACK[-1][1] if _FORCE_STACK else None
+    if not top:   # outside the finally block: an exception in flight (e.g. ImageRefusal) is never replaced
+        raise RuntimeError("_force_row: forcing contexts must nest (this one exited out of order)")

sparse_delta_memory/cuda/_devmemo.py

@@ -0,0 +1,64 @@
+"""Memoized device-class reads: _device's backend() and capability() serve exactly the originals' values (or raise their
+exceptions, never cached) from one bounded memo keyed by (read, device ordinal, forced test row), so a call makes no
+device query after the device's first one."""
+from . import _device as _dev_mod
+from ._device import CLASS_SMEM, RESERVED_SMEM
+
+_DEVCLS: dict = {}   # ("backend" | "capability", device ordinal, _device._FORCED) -> the original's value (successful reads only)
+
+
+def _memo(read, idx, compute):
+    """compute(idx) memoized under (read, idx, _device._FORCED); an exception propagates and is never cached."""
+    hit = _DEVCLS.get((read, idx, _dev_mod._FORCED))
+    if hit is None:
+        with _dev_mod._FORCE_LOCK:
+            key = (read, idx, _dev_mod._FORCED)
+            hit = _DEVCLS.get(key)
+            if hit is None:
+                hit = compute(idx)
+                if len(_DEVCLS) >= 64:
+                    _DEVCLS.clear()
+                _DEVCLS[key] = hit
+    return hit
+
+
+def _raw_capability(idx) -> tuple:
+    return tuple(_dev_mod.capability(idx))
+
+
+def backend(device=None) -> str:
+    """_device.backend(device), memoized."""
+    return _memo("backend", _dev_mod._index(device), _dev_mod.backend)
+
+
+def capability(device=None) -> tuple:
+    """_device.capability(device) on every row, supported or not (a forced row's directly), memoized."""
+    forced = _dev_mod._FORCED
+    if forced:
+        return forced[0]
+    return _memo("capability", _dev_mod._index(device), _raw_capability)
+
+
+def cuda_backend(device=None) -> bool:
+    """_device.cuda_backend(device), memoized."""
+    return bool(backend(device))
+
+
+def pdl_ok(device=None) -> bool:
+    """_device.pdl_ok(device), memoized."""
+    return backend(device) == "modern"
+
+
+def smem_cap(device=None) -> int:
+    """_device.smem_cap(device), memoized."""
+    return CLASS_SMEM[backend(device)]
+
+
+def smem_per_sm(device=None) -> int:
+    """_device.smem_per_sm(device), memoized."""
+    return smem_cap(device) + RESERVED_SMEM
+
+
+def sm90(device=None) -> bool:
+    """`device` is an sm_90 part (H100): capability(device) == (9, 0), memoized."""
+    return capability(device) == (9, 0)

sparse_delta_memory/cuda/_image.py

@@ -0,0 +1,128 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+"""_image: PTX-only builds for older GPUs are refused on Hopper / Blackwell GPUs, at every CUDA extension load
+(_build.load_cuda_extension), before any kernel of the extension runs. The modern routes (H100 / B200 / GB300) use device
+features only the modern passes compile (sdm_copy.cuh); the image the runtime loaded for a probe kernel, built like every
+extension, must be one (cudaFuncGetAttributes().ptxVersion). ImageRefusal is a BaseException: no `except Exception`
+fallback catches it. An RTX-class device warns once instead."""
+
+import contextlib
+import os
+import threading
+import warnings
+
+import torch
+
+
+class UntestedHardwareWarning(UserWarning):
+    """sparse_delta_memory runs on a GPU class this release has not tested (RTX: compute capability 8.6 / 8.9 / 12.0): the
+    results are unverified. Issued once per process per device, when the first CUDA extension loads."""
+
+
+UNTESTED_MESSAGE = ("sparse_delta_memory {version} has not been tested on RTX-class GPUs: device {index} ({name}, compute "
+                    "capability {major}.{minor}); results are unverified. Tested: A100 / H100 / B200 / GB300 (compute "
+                    "capability 8.0 / 9.0 / 10.0 / 10.3).")
+
+
+class ImageRefusal(BaseException):
+    """The package refuses to run the CUDA image it loaded (sparse_delta_memory.cuda._image). A BaseException, so that no
+    `except Exception` fallback reroutes it: catch it by name. The message names the image, the device and the fix."""
+
+
+MODERN_PTX = (90, 100, 103)   # ptxVersion of the device passes without SDM_PORTABLE_ARCH (sdm_copy.cuh)
+ENV_KEYS = ("TORCH_CUDA_ARCH_LIST", "CUDA_FORCE_PTX_JIT", "CUDA_DISABLE_PTX_JIT", "NVCC_APPEND_FLAGS", "NVCC_PREPEND_FLAGS")
+PROBE_BASE = "sdm_image_probe"
+MESSAGE = ("sparse_delta_memory: PTX-only builds for older GPUs are refused on Hopper / Blackwell GPUs. The CUDA image "
+           "loaded on device {index} ({name}, compute capability {major}.{minor}) was compiled for compute_{ptx} "
+           "(cudaFuncGetAttributes ptxVersion {ptx}; {env}). Rebuild with TORCH_CUDA_ARCH_LIST including your GPU "
+           "(9.0 for H100, 10.0 for B200, 10.3 for GB300), or unset it, and start a new process.")
+
+_PROBE = {}       # build environment -> the probe extension built under it
+_PTX = {}         # (compute capability, build environment) -> the probe's loaded ptxVersion
+_OK = set()       # (device ordinal, forced test row, build environment) checked
+_WARNED = set()   # (device ordinal, forced test row) given the RTX warning
+_TLS = threading.local()   # .depth: _unchecked_build_env in this thread (test only)
+
+
+def build_env() -> tuple:
+    """What the probe's and every extension's images depend on: ENV_KEYS and the toolkit (_build.toolkit)."""
+    from ._build import toolkit
+    return (*(os.environ.get(k) for k in ENV_KEYS), *toolkit())
+
+
+def env_summary() -> str:
+    return ", ".join(f"{k}={os.environ.get(k)!r}" for k in ENV_KEYS)
+
+
+_PROBE_CPP = "int64_t sdm_image_ptx();"
+_PROBE_SOURCE = r"""
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+__global__ void sdm_image_probe_kernel(int* p) { if (p != nullptr) *p = 0; }
+int64_t sdm_image_ptx() {
+    cudaFuncAttributes attr{};
+    const cudaError_t err = cudaFuncGetAttributes(&attr, reinterpret_cast<const void*>(sdm_image_probe_kernel));
+    if (err != cudaSuccess) {
+        (void)cudaGetLastError();
+        TORCH_CHECK(false, "sdm_image_ptx: cudaFuncGetAttributes failed: ", cudaGetErrorString(err));
+    }
+    return static_cast<int64_t>(attr.ptxVersion);
+}
+"""
+
+
+def probe_ptx(index: int, env: tuple) -> int:
+    """The ptxVersion of the image the runtime loads on device `index` for the probe kernel built under build environment
+    `env` (as for every extension built under it), once per (compute capability, environment)."""
+    real = tuple(int(c) for c in torch.cuda.get_device_capability(index))
+    if (real, env) not in _PTX:
+        if env not in _PROBE:
+            from ._build import load_cuda_extension
+            _PROBE[env] = load_cuda_extension(PROBE_BASE, cpp_sources=_PROBE_CPP, cuda_sources=_PROBE_SOURCE,
+                                              functions=("sdm_image_ptx",))
+        with torch.cuda.device(index):
+            _PTX[real, env] = int(_PROBE[env].sdm_image_ptx())
+    return _PTX[real, env]
+
+
+def require_supported_image(device=None) -> None:
+    """Refuses (ImageRefusal) a device whose routes are the modern ones (a forced test row included) while the image the
+    runtime loads there is not a modern pass; an RTX-class device warns once. `device` None = every visible device, the
+    current one first. Run by _build after every extension load; memoized per (device, forced row, build environment: an
+    extension built under another environment is checked again, with a probe built under it), refusals never."""
+    if getattr(_TLS, "depth", 0):
+        return
+    from . import _device
+    env = build_env()
+    if device is None:
+        cur = torch.cuda.current_device()
+        devices = [cur] + [j for j in range(torch.cuda.device_count()) if j != cur]
+    else:
+        devices = [_device._index(device)]
+    for i in devices:
+        key = (i, _device._FORCED, env)
+        if key in _OK:
+            continue
+        routes = _device.backend(i)
+        if routes == "modern" and probe_ptx(i, env) not in MODERN_PTX:
+            major, minor = (int(c) for c in torch.cuda.get_device_capability(i))
+            raise ImageRefusal(MESSAGE.format(index=i, name=torch.cuda.get_device_name(i), major=major, minor=minor,
+                                              ptx=probe_ptx(i, env), env=env_summary()))
+        if routes == "rtx" and key[:2] not in _WARNED:
+            from .. import __version__
+            cc = tuple(_device.capability(i))
+            warnings.warn(UNTESTED_MESSAGE.format(version=__version__, index=i, name=torch.cuda.get_device_name(i),
+                                                  major=cc[0], minor=cc[1]), UntestedHardwareWarning, stacklevel=2)
+            _WARNED.add(key[:2])
+        _OK.add(key)
+
+
+@contextlib.contextmanager
+def _unchecked_build_env():
+    """TEST ONLY: every visible device is checked first; inside, this thread's builds (a test's own module for another
+    target list, called directly, never through the package's routes) are not checked."""
+    require_supported_image()
+    _TLS.depth = getattr(_TLS, "depth", 0) + 1
+    try:
+        yield
+    finally:
+        _TLS.depth -= 1