Merge pull request #26 from Cytosine-code/fix/cuda-kernel-install

fix: build CUDA kernels during editable installation (sm89-gated fp8_ops) plus doc corrections

csrc/CMakeLists.txt

@@ -52,13 +52,16 @@ set(CMAKE_CUDA_ARCHITECTURES "${ASTRAI_CUDA_ARCH}")
 # globally unique across families) and their per-family source paths under
 # kernels/. `loader.py` auto-discovers the .so files in astrai/extension/lib/,
 # so this CMake registry is the single place to register a new kernel.
+#
+# FP8 MMA instructions require sm_89+. Keep the target out of the build on
+# older architectures instead of instantiating templates that cannot compile.
+# The remaining kernels are still useful on sm_80+ (including sm_86).
 set(KERNEL_NAMES
     attn_decode
     attn_prefill
     attn_paged_decode
     attn_paged_prefill
     rotary_emb
-    fp8_ops
 )
 set(KERNEL_SRCS
     attention/decode.cu
@@ -66,9 +69,17 @@ set(KERNEL_SRCS
     attention/paged_decode.cu
     attention/paged_prefill.cu
     rotary/rotary_emb.cu
-    fp8/ops.cu
 )
 
+if(ASTRAI_CUDA_ARCH GREATER_EQUAL 89)
+    list(APPEND KERNEL_NAMES fp8_ops)
+    list(APPEND KERNEL_SRCS fp8/ops.cu)
+else()
+    message(WARNING
+        "FP8 operator disabled: ASTRAI_CUDA_ARCH=${ASTRAI_CUDA_ARCH} "
+        "requires compute capability 89 or newer")
+endif()
+
 list(LENGTH KERNEL_NAMES _kernel_count)
 math(EXPR _kernel_last "${_kernel_count} - 1")
 foreach(i RANGE ${_kernel_last})

docs/developer/cuda_kernels.md

@@ -48,25 +48,20 @@ style as attention, but split into **three** files:
 |------|------|
 | `fp8/common.h` | `FP8Format` enum (E4M3/E5M2), `Fp8GemmTraits<Fmt, BlockM, BlockN, K, Stages>`, `FP8Params` POD — no torch |
 | `fp8/quantize.cuh` | pure-CUDA device code: `fp8_quantize_kernel<Fmt, InT>` (bf16/fp16/fp32 → FP8 + amax, `quant_in_traits<InT>` vectorized unpack) — no torch |
-| `fp8/gemm.cuh` | pure-CUDA device code: `fp8_gemm_kernel` (pre-quantized GEMM, 128×128 CTA / 64×32 warp / multi-stage cp.async, transposed-operand layouts) — no torch |
+| `fp8/gemm.cuh` | pure-CUDA device code: `fp8_gemm_kernel` (pre-quantized GEMM; 64×64 / 128×128 CTA picked at runtime by `prefer_small_cta`, 64×32 warp tiles, multi-stage cp.async, transposed-operand layouts) — no torch |
 | `fp8/ops.cu` | binding only: `check_fp8_device` (sm_89+), param packing, launch dispatch, pybind → module `fp8_ops` |
 
-Scale semantics: `quantize` takes the quantization *multiplier*, `mm_fp8`
+Scale semantics: `quantize` takes the quantization *multiplier*; the
-takes the combined dequant scale (`sa * sb`); the strategy layer passes
+strategy layer passes `scale.reciprocal()` and the kernel multiplies by it.
-`scale.reciprocal()` / `sa * sb` respectively. `amax` is always returned in
+`mm_fp8` takes the combined dequant scale (`sa * sb`). `amax` is always
-the original input domain.
+returned in the original input domain.
-
-`mm_fp8` also accepts 3D (batched) operands through the same signature:
-`grid.z` slices the operands by their batch strides, a size-1 batch
-broadcasts (stride 0), and inner-transposed views (e.g. `x.t()`) fold into
-the kernel's layout tag at zero copy — only genuinely strided operands pay
-a `.contiguous()` copy.
 
 Python layer (two levels): `astrai/extension/ops/fp8.py` provides stateless
-primitives (`quantize` / `mm_fp8`) via `torch.library.custom_op`, and
+primitives (`fp8_quantize` / `fp8_gemm`) via `torch.library.custom_op`, with
-`astrai/extension/fp8.py` is the strategy layer (`fp8_autocast`, delayed /
+plain `quantize` / `mm_fp8` wrappers, and `astrai/extension/fp8.py` is the
-dynamic scaling recipes, `fp8_linear_forward/backward` wiring `aten::linear`
+strategy layer (`fp8_autocast`, delayed / dynamic scaling recipes,
-on CUDA). See the FP8 section in `AGENTS.md` for full detail.
+`fp8_linear_forward/backward` wiring `aten::linear` on CUDA). See the FP8
+section in `AGENTS.md` for full detail.
 
 ## Build System
 
@@ -105,7 +100,9 @@ unset, `setup.py` auto-detects the real GPU capability through
 - **sm_80+** (Ampere and later): enables the tensor-core MMA path
   (`mma.sync.m16n8k16.bf16` for bf16 attention, `mma.sync.m16n8k32` for FP8).
 - **sm_89+**: required for the FP8 family (`fp8_ops`) — FP8 tensor-core
-  instructions only exist on Ada/Hopper and newer.
+  instructions only exist on Ada/Hopper and newer. On older architectures,
+  CMake emits a warning and skips the `fp8_ops` target so the remaining CUDA
+  kernels still build successfully.
 - **`-DASTRAI_NO_MMA`** is a manual escape hatch only — the build never defines
   it automatically. To disable the MMA path, add it to `NVCC_FLAGS` yourself;
   all supported build targets are sm_80+.
@@ -119,7 +116,7 @@ NVCC_FLAGS = -O3 --expt-relaxed-constexpr --use_fast_math
              --ptxas-options=-O3,-v --extra-device-vectorization --threads=16
 ```
 
-Each kernel in `astrai/extension/lib` is compiled as an independent pybind11 module (one `.so` per kernel, named `<kernel>.cpython-*-x86_64-linux-gnu.so`). CMake builds all six kernel targets in parallel via `cmake --build -j N`. The target list is the **single source of truth**: `KERNEL_NAMES` and the parallel `KERNEL_SRCS` list in `csrc/CMakeLists.txt`; `astrai/extension/loader.py` auto-discovers the compiled `.so` files.
+Each kernel in `astrai/extension/lib` is compiled as an independent pybind11 module (one `.so` per kernel, named `<kernel>.cpython-*-x86_64-linux-gnu.so`). CMake builds all registered kernel targets in parallel via `cmake --build -j N` (the five base targets always; `fp8_ops` additionally on sm_89+). The target list is the **single source of truth**: `KERNEL_NAMES` and the parallel `KERNEL_SRCS` list in `csrc/CMakeLists.txt`; `astrai/extension/loader.py` auto-discovers the compiled `.so` files.
 
 ## Python Extension Architecture
 

setup.py

@@ -6,7 +6,9 @@ import warnings
 from pathlib import Path
 
 from setuptools import setup
+from setuptools.command.build import build as _build
 from setuptools.command.build_ext import build_ext as _build_ext
+from setuptools.command.editable_wheel import editable_wheel as _editable_wheel
 
 sys.path.insert(0, str(Path(__file__).parent))
 os.makedirs("astrai/extension/lib", exist_ok=True)
@@ -93,10 +95,40 @@ class _CMakeBuildExt(_build_ext):
         if not arch:
             arch = _detect_cuda_arch()
         if arch:
+            try:
+                if int(str(arch)) < 89:
+                    warnings.warn(
+                        f"FP8 operator disabled: CUDA compute capability {arch} "
+                        "requires 89 or newer.",
+                        stacklevel=2,
+                    )
+            except ValueError:
+                warnings.warn(
+                    f"Could not parse ASTRAI_CUDA_ARCH={arch!r}; "
+                    "FP8 capability will be decided by CMake.",
+                    stacklevel=2,
+                )
             cfg.append(f"-DASTRAI_CUDA_ARCH={arch}")
         subprocess.run(cfg, check=True)
         subprocess.run([cmake, "--build", str(build_dir), "-j", parallel], check=True)
 
+        # After compilation finishes, verify mandatory CUDA kernels to confirm build succeeded.
+        # CMake may report partial‑target success even if some architecture‑specific kernels are skipped.
+        # Prevent editable install from reporting success when critical kernel shared objects are missing.
+        lib_dir = src / "astrai" / "extension" / "lib"
+        required = (
+            "attn_decode",
+            "attn_prefill",
+            "attn_paged_decode",
+            "attn_paged_prefill",
+            "rotary_emb",
+        )
+        missing = [name for name in required if not any(lib_dir.glob(f"{name}.*.so"))]
+        if missing:
+            raise RuntimeError(
+                "CUDA build completed without some required kernel modules!"
+            )
+
 
 def _cuda_toolkit_version():
     import shutil
@@ -148,6 +180,24 @@ class _NullBuildExt(_build_ext):
         pass
 
 
+class _Build(_build):
+    """Run the CMake kernel build as part of setuptools' build lifecycle."""
+
+    def run(self):
+        if _should_build():
+            self.run_command("build_ext")
+        super().run()
+
+
+class _EditableWheel(_editable_wheel):
+    """Run the CMake kernel build for PEP 660 editable installations."""
+
+    def run(self):
+        if _should_build():
+            self.run_command("build_ext")
+        super().run()
+
+
 cmdclass = {}
 
 if _should_build():
@@ -155,4 +205,10 @@ if _should_build():
 else:
     cmdclass["build_ext"] = _NullBuildExt
 
-setup(ext_modules=[], cmdclass=cmdclass)
+cmdclass["build"] = _Build
+cmdclass["editable_wheel"] = _EditableWheel
+
+setup(
+    ext_modules=[],
+    cmdclass=cmdclass,
+)