fix: report true rank in logs via dist-aware helpers

- get_rank/get_world_size fall back to RANK/WORLD_SIZE env instead of hardcoded 0/1, matching torchrun's env-before-init contract
- log filter reuses the helpers so initialized groups show real ranks; local-spawn children previously logged rank=0/8

astrai/extension/fp8.py

@@ -412,11 +412,13 @@ class _LinearFp8(torch.autograd.Function):
         w8 = w if _is_fp8(w.dtype) else quantize(w, sw.reciprocal(), fmt)[0]
         grad_x = mm_fp8(g8, w8, sg * sw).reshape(x.shape)  # g8[m,n] @ w8[n,k]
         grad_w = mm_fp8(g8, x8, sg * sx, trans_a=True)  # g8.T @ x8
-        grad_b = g2.sum(0).to(torch.bfloat16)
+        # bias-free linears must not pay the column-sum
+        # reduce: g2.sum(0) is another full read of the gradient.
+        grad_b = g2.sum(0).to(torch.bfloat16) if ctx.needs_input_grad[2] else None
         if not ctx.is_dynamic:
             meta.g.update(amax_g, fmt)
             meta.g.advance()
-        return grad_x, grad_w, grad_b if ctx.needs_input_grad[2] else None
+        return grad_x, grad_w, grad_b
 
 
 # ---------------------------------------------------------------------------

astrai/logging.py

@@ -1,11 +1,13 @@
 import logging
 import os
 
+from astrai.parallel.setup import get_rank, get_world_size
+
 
 class _DistributedContextFilter(logging.Filter):
     def filter(self, record: logging.LogRecord) -> bool:
-        record.rank = os.environ.get("RANK", "0")
+        record.rank = str(get_rank())
-        record.world_size = os.environ.get("WORLD_SIZE", "1")
+        record.world_size = str(get_world_size())
         return True
 
 

astrai/parallel/setup.py

@@ -30,15 +30,13 @@ def get_current_device():
 def get_world_size() -> int:
     if dist.is_available() and dist.is_initialized():
         return dist.get_world_size()
-    else:
+    return int(os.environ.get("WORLD_SIZE", "1"))
-        return 1
 
 
 def get_rank() -> int:
     if dist.is_available() and dist.is_initialized():
         return dist.get_rank()
-    else:
+    return int(os.environ.get("RANK", "0"))
-        return 0
 
 
 @contextmanager

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})

csrc/kernels/fp8/gemm.cuh

@@ -942,11 +942,19 @@ void launch_with_smem(int smem_bytes, dim3 grid, dim3 block,
 // that stall, and the trade flipped across the whole measured band: 1280^3
 // (1.09 waves) big 163T vs small 100T, 4096x512x4096 (1.4 waves) big 139T
 // vs small 113T. The small CTA now only serves the genuinely sub-wave band
-// below 5/8 of a wave (and m <= 64, where a 128-row CTA wastes half its
+// below 5/8 of a wave, the narrow-M/N edge, and the divisibility rule below.
-// rows); inside [5/8, 1] waves the big CTA was already the measured winner
+inline bool prefer_small_cta(int64_t tiles_128, int64_t m, int64_t n) {
-// (63-tile rect +8%).
+    if (m <= 64 || n <= 64) return true;
-inline bool prefer_small_cta(int64_t tiles_128, int64_t m) {
+    const bool big_div = (m % 128 == 0) && (n % 128 == 0);
-    if (m <= 64) return true;
+    const bool small_div = (m % 64 == 0) && (n % 64 == 0);
+    // Divisibility: a 128x128 CTA that is NOT exactly tiled (m or n not a
+    // multiple of 128) runs its edge tiles on the predicated generic path,
+    // and with a single in-flight wave the runtime is the slowest CTA — the
+    // edge tiles drag the whole shape down (1088^3: 76T vs 93T with the
+    // 64x64 CTA, whose grid tiles exactly and overlaps waves; measured
+    // sweep, perf 5.1). When 64 divides both dims, the 64x64 small CTA
+    // wins the non-128-divisible band by 23..67%.
+    if (!big_div && small_div) return true;
     return tiles_128 < device_sm_count() * 5 / 8;
 }
 
@@ -962,7 +970,7 @@ void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
     // per-matrix tiles (see prefer_small_cta).
     const int64_t tiles_128 =
         (int64_t)p.batch * ((p.m + 127) / 128) * ((p.n + 127) / 128);
-    if (prefer_small_cta(tiles_128, p.m)) {
+    if (prefer_small_cta(tiles_128, p.m, p.n)) {
         dim3 grid((p.n + 63) / 64, (p.m + 63) / 64, p.batch);
         // Full-ring small CTAs — ONE __syncthreads per k-tile, cuBLAS's
         // barrier structure (the lean ring traded a second barrier for a

docker-compose.yml

@@ -1,5 +1,6 @@
 services:
   server:
+    image: astrai:latest
     build:
       context: .
       dockerfile: Dockerfile
@@ -20,7 +21,7 @@ services:
         reservations:
           devices:
             - driver: nvidia
-              count: 1
+              count: all
               capabilities: [gpu]
     healthcheck:
       test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
@@ -31,6 +32,7 @@ services:
     restart: unless-stopped
 
   server-cpu:
+    image: astrai:latest
     profiles: [cpu]
     build:
       context: .
@@ -54,6 +56,7 @@ services:
     restart: unless-stopped
 
   trainer:
+    image: astrai:latest
     profiles: [train]
     build:
       context: .

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
 

docs/developer/docker-serving.md

@@ -17,7 +17,7 @@ scripts/serve.sh             preflight, Compose wrapper, lifecycle
         └── server.py --config /run/astrai/serve.yaml
 ```
 
-`scripts/tools/serve_runtime.py` reads `runtime:` plus the two container-side
+`scripts/docker/serve_runtime.py` reads `runtime:` plus the two container-side
 values Compose needs (`server.port` for the port mapping, `server.device` for
 the preflight GPU check). `scripts/tools/server.py --config` reads `server:`.
 Explicit CLI arguments to `server.py` override `server:` YAML values.
@@ -54,8 +54,9 @@ server:
 - `runtime.gpu.enabled: true` (default) selects the `server` service with an
   NVIDIA device reservation; `false` selects `server-cpu` (no GPU passthrough).
   When disabled, `server.device` must be `cpu`.
-- `runtime.gpu.devices` is either `all` or a single-device list such as `[0]`;
+- `runtime.gpu.devices` is `all` (default) or a single-device list such as `[0]`;
-  the list becomes `CUDA_VISIBLE_DEVICES`. Compose passes `count: 1`.
+  the list becomes `CUDA_VISIBLE_DEVICES`. Compose passes `count: all`; the
+  env var performs the only filtering.
 - `environment` values are explicitly passed to the serving container. Keep
   host-specific settings here; they are not universal defaults.
 - `server.device` must agree with `runtime.gpu.enabled`; `preflight` enforces it.
@@ -89,9 +90,9 @@ bash scripts/serve.sh status [CONFIG]
 
 `preflight` validates Docker, the model directory
 (`config.json` + `model.safetensors`), GPU/device consistency, and the
-rendered Compose configuration. `up` starts the container detached and
+rendered Compose configuration. `up` starts the container detached; `run`
-rebuilds the image when the code changed (`--build`); `run` keeps it in the
+keeps it in the foreground. Both reuse the existing image; run
-foreground. The wrapper manages a fixed container name
+`bash scripts/serve.sh build [CONFIG]` after code changes. The wrapper manages a fixed container name
 (`astrai-server` or `astrai-server-<job_name>`); the plain
 `docker compose up -d` / `docker compose --profile cpu up -d` path keeps
 working with defaults (port 8000, `./params`).
@@ -99,7 +100,7 @@ working with defaults (port 8000, `./params`).
 ## Hard Rules
 
 1. Keep Docker settings in `runtime` and server settings in `server`.
-2. Filter GPUs once: the `server` service reserves one device; a `devices`
+2. Filter GPUs once: Compose passes `count: all`; a `devices`
    list becomes `CUDA_VISIBLE_DEVICES`.
 3. `runtime.gpu.enabled: false` requires `server.device: cpu`.
 4. In Docker, `server.port` must match the published container port (default

docs/developer/docker-training.md

@@ -18,7 +18,7 @@ scripts/train.sh             preflight, Compose wrapper, lifecycle, timer
               └── train.py --config /run/astrai/train.yaml
 ```
 
-The two parsers deliberately own different sections. `scripts/tools/train_runtime.py`
+The two parsers deliberately own different sections. `scripts/docker/train_runtime.py`
 reads only `runtime`; `scripts/tools/train.py` reads only
 `model/data/parallel/training/ckpt/log`. Explicit trainer arguments after `--`
 override training YAML values.
@@ -72,8 +72,8 @@ runtime:
 | the selected YAML | `/run/astrai/train.yaml` | read-only |
 
 Training configuration must therefore use `data_root_path: /data`. The source
-code is baked into `/app`; `start` uses `--build`, so code changes rebuild the
+code is baked into `/app`; `start` reuses the existing image, so run
-image when necessary.
+`bash scripts/train.sh build [CONFIG]` after code changes.
 
 ## Operations
 

scripts/docker/serve_runtime.py

@@ -78,9 +78,10 @@ def load_runtime(config_path: str) -> dict[str, str]:
         raise ValueError("runtime.gpu.enabled must be a boolean")
 
     devices = gpu.get("devices", "all")
+    visible_devices = None
     if gpu_enabled:
         if devices == "all":
-            visible_devices = ""
+            pass
         elif isinstance(devices, list) and len(devices) == 1:
             text = str(devices[0])
             if not text.isdigit():
@@ -93,7 +94,6 @@ def load_runtime(config_path: str) -> dict[str, str]:
                 "runtime.gpu.devices must be 'all' or a single-device list such as [0]"
             )
     else:
-        visible_devices = ""
         if devices != "all":
             raise ValueError(
                 "runtime.gpu.devices is ignored when runtime.gpu.enabled is false"
@@ -110,9 +110,10 @@ def load_runtime(config_path: str) -> dict[str, str]:
         "SERVE_PARAM_DIR": _path(paths.get("param", "./params"), "param", path.parent),
         "SERVE_GPU_ENABLED": "true" if gpu_enabled else "false",
         "SERVE_DEVICE": device,
-        "CUDA_VISIBLE_DEVICES": visible_devices,
         "CUDA_TAG": str(container.get("cuda_tag", "cu128")),
     }
+    if visible_devices is not None:
+        values["CUDA_VISIBLE_DEVICES"] = visible_devices
 
     for name, value in environment.items():
         if not isinstance(name, str) or not ENV_NAME.fullmatch(name):

scripts/docker/train_runtime.py

@@ -53,9 +53,9 @@ def load_runtime(config_path: str) -> dict[str, str]:
         )
 
     devices = gpu.get("devices", "all")
+    visible_devices = None
     if devices == "all":
         gpu_count = "all"
-        visible_devices = ""
     elif isinstance(devices, list) and devices:
         normalized = []
         for device in devices:
@@ -102,7 +102,6 @@ def load_runtime(config_path: str) -> dict[str, str]:
             paths.get("checkpoints"), "checkpoints", path.parent
         ),
         "TRAIN_GPU_COUNT": gpu_count,
-        "CUDA_VISIBLE_DEVICES": visible_devices,
         "TRAIN_PARALLEL_MODE": parallel_mode,
         "CUDA_TAG": str(container.get("cuda_tag", "cu128")),
         "TRAIN_IPC_MODE": str(container.get("ipc", "host")),
@@ -111,6 +110,8 @@ def load_runtime(config_path: str) -> dict[str, str]:
         "CHECKPOINT_KEEP_LAST": str(container.get("checkpoint_keep_last", 5)),
         "TRAIN_MAX_DURATION_SECONDS": str(max_seconds),
     }
+    if visible_devices is not None:
+        values["CUDA_VISIBLE_DEVICES"] = visible_devices
 
     for name, value in environment.items():
         if not isinstance(name, str) or not ENV_NAME.fullmatch(name):

scripts/serve.sh

@@ -46,7 +46,7 @@ load_config() {
         die "PyYAML is required on the host (install python3-yaml)"
 
     local exports
-    exports="$(python3 "${ROOT_DIR}/scripts/tools/serve_runtime.py" exports "${CONFIG_FILE}")" ||
+    exports="$(python3 "${ROOT_DIR}/scripts/docker/serve_runtime.py" exports "${CONFIG_FILE}")" ||
         die "Failed to load runtime configuration"
     eval "${exports}"
     if [[ -n "${SERVE_JOB_NAME}" ]]; then
@@ -55,7 +55,12 @@ load_config() {
 }
 
 compose() {
+    if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then
         ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" "${COMPOSE_BASE[@]}" "$@"
+    else
+        ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" \
+            env -u CUDA_VISIBLE_DEVICES "${COMPOSE_BASE[@]}" "$@"
+    fi
 }
 
 container_name() {
@@ -108,7 +113,7 @@ runtime_environment_args() {
     local pair
     while IFS= read -r -d '' pair; do
         RUNTIME_ENV_ARGS+=(--env "${pair}")
-    done < <(python3 "${ROOT_DIR}/scripts/tools/serve_runtime.py" environment "${CONFIG_FILE}")
+    done < <(python3 "${ROOT_DIR}/scripts/docker/serve_runtime.py" environment "${CONFIG_FILE}")
 }
 
 start_server() {
@@ -129,11 +134,11 @@ start_server() {
         "${RUNTIME_ENV_ARGS[@]}"
     )
     if [[ "${foreground}" == "true" ]]; then
-        compose "${PROFILE_ARGS[@]}" run --build --rm --service-ports \
+        compose "${PROFILE_ARGS[@]}" run --rm --service-ports \
             "${run_options[@]}" "$(service_name)" \
             python -m scripts.tools.server --config /run/astrai/serve.yaml "$@"
     else
-        compose "${PROFILE_ARGS[@]}" run -d --build --service-ports \
+        compose "${PROFILE_ARGS[@]}" run -d --service-ports \
             --name "${container}" "${run_options[@]}" "$(service_name)" \
             python -m scripts.tools.server --config /run/astrai/serve.yaml "$@"
         log_info "Server started; run scripts/serve.sh logs ${CONFIG_FILE} to follow it"

scripts/train.sh

@@ -51,14 +51,19 @@ load_config() {
         die "PyYAML is required on the host (install python3-yaml)"
 
     local exports
-    exports="$(python3 "${ROOT_DIR}/scripts/tools/train_runtime.py" exports "${CONFIG_FILE}")" ||
+    exports="$(python3 "${ROOT_DIR}/scripts/docker/train_runtime.py" exports "${CONFIG_FILE}")" ||
         die "Failed to load runtime configuration"
     eval "${exports}"
     validate_job_name "${TRAIN_JOB_NAME}"
 }
 
 compose() {
+    if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then
         ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" "${COMPOSE_BASE[@]}" "$@"
+    else
+        ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" \
+            env -u CUDA_VISIBLE_DEVICES "${COMPOSE_BASE[@]}" "$@"
+    fi
 }
 
 checkpoint_dir() {
@@ -143,7 +148,7 @@ runtime_environment_args() {
     local pair
     while IFS= read -r -d '' pair; do
         RUNTIME_ENV_ARGS+=(--env "${pair}")
-    done < <(python3 "${ROOT_DIR}/scripts/tools/train_runtime.py" environment "${CONFIG_FILE}")
+    done < <(python3 "${ROOT_DIR}/scripts/docker/train_runtime.py" environment "${CONFIG_FILE}")
 }
 
 start_training() {
@@ -164,9 +169,9 @@ start_training() {
         "${RUNTIME_ENV_ARGS[@]}"
     )
     if [[ "${foreground}" == "true" ]]; then
-        compose run --build --rm "${run_options[@]}" trainer "$@"
+        compose run --rm "${run_options[@]}" trainer "$@"
     else
-        compose run -d --build --name "${container}" "${run_options[@]}" trainer "$@"
+        compose run -d --name "${container}" "${run_options[@]}" trainer "$@"
         schedule_timer
         log_info "Training started; run scripts/train.sh logs ${CONFIG_FILE} to follow it"
     fi

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,
+)

tests/test_serve_runtime.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from scripts.tools.serve_runtime import load_runtime
+from scripts.docker.serve_runtime import load_runtime
 
 
 def _write(tmp_path, body: str) -> str:
@@ -26,7 +26,7 @@ def test_runtime_exports_defaults(tmp_path):
     assert runtime["SERVE_CONTAINER_PORT"] == "8000"
     assert runtime["SERVE_PARAM_DIR"] == str((tmp_path / "params").resolve())
     assert runtime["SERVE_GPU_ENABLED"] == "true"
-    assert runtime["CUDA_VISIBLE_DEVICES"] == ""
+    assert "CUDA_VISIBLE_DEVICES" not in runtime
     assert runtime["SERVE_DEVICE"] == "cuda"
     assert runtime["CUDA_TAG"] == "cu128"
     assert runtime["SERVE_JOB_NAME"] == ""