diff --git a/.dockerignore b/.dockerignore index 9926b40..a52f56f 100644 --- a/.dockerignore +++ b/.dockerignore @@ -5,5 +5,7 @@ !astrai/ !scripts/ !docs/ +!csrc/ +!setup.py !pyproject.toml !README.md diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index c281d9c..485757d 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -26,22 +26,30 @@ jobs: if-no-files-found: error build-cuda-linux: - name: Build CUDA wheel (Linux) + name: Build CUDA wheel (Linux, ${{ matrix.cuda_tag }}) runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + include: + - cuda_tag: "cu128" + cuda_ver: "12.8.0" + - cuda_tag: "cu130" + cuda_ver: "13.0.0" steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 with: python-version: "3.12" - - name: Install torch (CUDA 12.8) + - name: Install torch (${{ matrix.cuda_tag }}) run: | - pip install torch --index-url https://download.pytorch.org/whl/cu128 + pip install torch --index-url https://download.pytorch.org/whl/${{ matrix.cuda_tag }} - - name: Setup CUDA + - name: Setup CUDA (${{ matrix.cuda_ver }}) uses: Jimver/cuda-toolkit@v0.2.35 with: - cuda: "12.8.0" + cuda: "${{ matrix.cuda_ver }}" - name: Build wheel (with CUDA kernels) run: | @@ -49,7 +57,7 @@ jobs: - uses: actions/upload-artifact@v4 with: - name: cuda-wheel-linux + name: cuda-wheel-linux-${{ matrix.cuda_tag }} path: dist/*.whl if-no-files-found: error @@ -66,10 +74,11 @@ jobs: name: pure-wheel path: release-assets/pure - - name: Download CUDA wheel + - name: Download CUDA wheels (all variants) uses: actions/download-artifact@v4 with: - name: cuda-wheel-linux + pattern: cuda-wheel-linux-* + merge-multiple: true path: release-assets/cuda - name: Verify release assets @@ -79,8 +88,7 @@ jobs: pure_wheels=(release-assets/pure/*.whl) cuda_wheels=(release-assets/cuda/*.whl) test "${#pure_wheels[@]}" -eq 1 - test "${#cuda_wheels[@]}" -eq 1 - test "$(basename "${pure_wheels[0]}")" != "$(basename "${cuda_wheels[0]}")" + test "${#cuda_wheels[@]}" -ge 1 - name: Create release & upload assets uses: softprops/action-gh-release@v2 diff --git a/Dockerfile b/Dockerfile index c07428c..a0f3e85 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,8 +1,16 @@ # AstrAI Dockerfile - Multi-stage Build (Optimized) +# +# CUDA version selection: +# docker build -t astrai . +# docker build -t astrai --build-arg CUDA_TAG=cu128 . +# docker build -t astrai --build-arg CUDA_TAG=cu130 . +# Default: cu128 # Build stage - use base image with minimal build tools FROM ubuntu:24.04 AS builder +ARG CUDA_TAG=cu128 + WORKDIR /app # Install Python 3.12 and minimal build dependencies @@ -20,10 +28,12 @@ ENV PATH="/opt/venv/bin:$PATH" # Copy source code and install (deps read from pyproject.toml) COPY astrai/ ./astrai/ +COPY csrc/ ./csrc/ +COPY setup.py . COPY pyproject.toml . RUN pip install --no-cache-dir --upgrade pip \ && pip install --no-cache-dir . \ - --extra-index-url https://download.pytorch.org/whl/cu128 + --extra-index-url "https://download.pytorch.org/whl/${CUDA_TAG}" # Production stage FROM ubuntu:24.04 AS production diff --git a/csrc/build.py b/csrc/build.py index 540f978..7f1d1ab 100644 --- a/csrc/build.py +++ b/csrc/build.py @@ -1,6 +1,32 @@ from pathlib import Path +def cuda_toolkit_version() -> tuple[int, int] | None: + """Return ``(major, minor)`` of the nvcc on PATH, or ``None``. + + Used by ``setup.py`` to detect nvcc/torch CUDA version mismatches + (e.g. nvcc 13.0 with a cu128 torch wheel) which cause cryptic ABI errors. + """ + import shutil + import subprocess + + nvcc = shutil.which("nvcc") + if nvcc is None: + return None + try: + out = subprocess.check_output( + [nvcc, "--version"], stderr=subprocess.STDOUT, text=True + ) + for line in out.splitlines(): + if "release" in line: + ver = line.split("release")[1].split(",")[0].strip() + major, minor = ver.split(".") + return (int(major), int(minor)) + except Exception: + pass + return None + + def _arch_flags() -> list[str]: import torch diff --git a/csrc/kernels/attn_entry_utils.cuh b/csrc/kernels/attn_entry_utils.cuh index 47a4f17..b35730b 100644 --- a/csrc/kernels/attn_entry_utils.cuh +++ b/csrc/kernels/attn_entry_utils.cuh @@ -23,8 +23,8 @@ using bf16 = __nv_bfloat16; template inline void alloc_split_partials(P& p) { auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); - auto o_part = torch::empty({p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt); - auto ml_part = torch::empty({p.batch, p.q_head, MAX_SPLITS, 2}, fopt); + auto o_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt); + auto ml_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, 2}, fopt); p.o_part = (float*)o_part.data_ptr(); p.ml_part = (float*)ml_part.data_ptr(); } diff --git a/csrc/kernels/attn_mma_utils.cuh b/csrc/kernels/attn_mma_utils.cuh index 3ed44f2..34a2ccb 100644 --- a/csrc/kernels/attn_mma_utils.cuh +++ b/csrc/kernels/attn_mma_utils.cuh @@ -3,6 +3,12 @@ #include #include +// Predicated cp.async (4-operand form) requires CUDA 11.2+. +// bf16 mma.sync requires sm_80+ (guarded at build time by ASTRAI_NO_MMA). +#if CUDART_VERSION < 11020 +#error "AstrAI CUDA kernels require CUDA 11.2 or later (CUDART_VERSION >= 11020)." +#endif + // ============================================================================ // KernelTraits — FlashAttention-v2 style compile-time configuration bundle. // diff --git a/docker-compose.yml b/docker-compose.yml index 71cf702..7ab466f 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -3,6 +3,8 @@ services: build: context: . dockerfile: Dockerfile + args: + CUDA_TAG: ${CUDA_TAG:-cu128} user: "${UID:-1000}:${GID:-1000}" ports: - "8000:8000" @@ -29,6 +31,8 @@ services: build: context: . dockerfile: Dockerfile + args: + CUDA_TAG: ${CUDA_TAG:-cu128} user: "${UID:-1000}:${GID:-1000}" ports: - "8000:8000" diff --git a/pyproject.toml b/pyproject.toml index dfacc57..b53ed2b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,9 +36,6 @@ dev = ["pytest==9.0.2", "ruff", "httpx2"] [tool.setuptools.packages.find] where = ["."] -[tool.pip] -extra-index-url = "https://download.pytorch.org/whl/cu128" - [tool.setuptools.dynamic] version = { attr = "astrai.__version__" } diff --git a/setup.py b/setup.py index 2f1af11..3e48e9d 100644 --- a/setup.py +++ b/setup.py @@ -1,5 +1,6 @@ import os import sys +import warnings from pathlib import Path from setuptools import setup @@ -32,7 +33,24 @@ if _should_build(): import torch from torch.utils.cpp_extension import BuildExtension, CUDAExtension - from csrc.build import REGISTRY + from csrc.build import REGISTRY, cuda_toolkit_version + + # Preflight: warn if nvcc major version != torch's bundled CUDA major version. + # A mismatch (e.g. nvcc 13.0 + cu128 torch) causes cryptic ABI/header errors. + nvcc_ver = cuda_toolkit_version() + torch_cuda = torch.version.cuda + if nvcc_ver is not None and torch_cuda is not None: + torch_major = int(torch_cuda.split(".")[0]) + if nvcc_ver[0] != torch_major: + warnings.warn( + f"CUDA version mismatch: nvcc is {nvcc_ver[0]}.{nvcc_ver[1]} " + f"but torch was built with CUDA {torch_cuda}. " + f"This may cause compilation errors. " + f"Install a matching torch wheel: " + f"pip install torch --index-url " + f"https://download.pytorch.org/whl/cu{nvcc_ver[0]}{nvcc_ver[1]}", + stacklevel=2, + ) _torch_lib = torch.utils.cpp_extension.library_paths()[0]