build: parametrize CUDA version for wheels and docker

- Add cu128/cu130 build matrix to release workflow
- Parametrize Dockerfile and docker-compose with CUDA_TAG build arg
- Allow csrc/ and setup.py in docker context via .dockerignore
- Add nvcc/torch CUDA version mismatch preflight warning in setup.py
- Add cuda_toolkit_version() helper in csrc/build.py
- Use at::IntArrayRef explicitly to fix ATen overload ambiguity
- Guard kernels with CUDART_VERSION >= 11020 check
- Remove invalid [tool.pip] section from pyproject.toml

.dockerignore

@@ -5,5 +5,7 @@
 !astrai/
 !scripts/
 !docs/
+!csrc/
+!setup.py
 !pyproject.toml
 !README.md

.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

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

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
 

csrc/kernels/attn_entry_utils.cuh

@@ -23,8 +23,8 @@ using bf16 = __nv_bfloat16;
 template<typename P>
 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();
 }

csrc/kernels/attn_mma_utils.cuh

@@ -3,6 +3,12 @@
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
 
+// 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.
 //

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"

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__" }
 

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]