diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index 5dc291f..d17ea61 100644 --- a/csrc/CMakeLists.txt +++ b/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}) diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index b7cadaf..947a9a3 100644 --- a/docs/developer/cuda_kernels.md +++ b/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`, `FP8Params` POD — no torch | | `fp8/quantize.cuh` | pure-CUDA device code: `fp8_quantize_kernel` (bf16/fp16/fp32 → FP8 + amax, `quant_in_traits` 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` -takes the combined dequant scale (`sa * sb`); the strategy layer passes -`scale.reciprocal()` / `sa * sb` respectively. `amax` is always 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. +Scale semantics: `quantize` takes the quantization *multiplier*; the +strategy layer passes `scale.reciprocal()` and the kernel multiplies by it. +`mm_fp8` takes the combined dequant scale (`sa * sb`). `amax` is always +returned in the original input domain. Python layer (two levels): `astrai/extension/ops/fp8.py` provides stateless -primitives (`quantize` / `mm_fp8`) via `torch.library.custom_op`, and -`astrai/extension/fp8.py` is the strategy layer (`fp8_autocast`, delayed / -dynamic scaling recipes, `fp8_linear_forward/backward` wiring `aten::linear` -on CUDA). See the FP8 section in `AGENTS.md` for full detail. +primitives (`fp8_quantize` / `fp8_gemm`) via `torch.library.custom_op`, with +plain `quantize` / `mm_fp8` wrappers, and `astrai/extension/fp8.py` is the +strategy layer (`fp8_autocast`, delayed / dynamic scaling recipes, +`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 `.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 `.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 diff --git a/setup.py b/setup.py index 3b84717..78d5f12 100644 --- a/setup.py +++ b/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, +)