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
This commit is contained in:
+13
-2
@@ -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})
|
||||
|
||||
@@ -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`
|
||||
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 `<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
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user