ci: build & test the CUDA extension on the real GPU arch, fail fast i…

.github/workflows/gpu-ci.yml

@@ -48,9 +48,26 @@ jobs:
           chmod 600 ~/.ssh/id_ed25519
           ssh-keygen -y -f ~/.ssh/id_ed25519 > /dev/null && echo "key OK" || echo "key BROKEN"
 
+      # Follow-up (#191): to test multiple architectures, add a matrix and pass
+      # GPU_ID + TARGET_SM (+ KERNEL_ALIGN_FORCE_SM90 for Hopper) through to the script.
+      # run_gpu_ci.sh reads all three, normalizes TARGET_SM, asserts the pod matches it,
+      # and forwards KERNEL_ALIGN_FORCE_SM90 into the remote build:
+      #   strategy:
+      #     matrix:
+      #       include:
+      #         - { gpu_id: "NVIDIA RTX A4000",      target_sm: "8.6" }   # Ampere
+      #         - { gpu_id: "NVIDIA A100 80GB PCIe", target_sm: "8.0" }
+      #         - { gpu_id: "NVIDIA H100 PCIe",      target_sm: "9.0", force_sm90: "1" }  # build Hopper TMA/WGMMA kernels
+      #         - { gpu_id: "NVIDIA B200",           target_sm: "10.0" }
+      # Per-arch jobs must NOT fall back to a different-capability GPU: the script
+      # fails fast when the pod arch != requested TARGET_SM, so keep fallback within
+      # the same compute capability (or unset it for these jobs).
       - name: Run GPU tests on RunPod
         env:
           RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
           PR_REPO_URL: ${{ github.event.pull_request.head.repo.clone_url }}
           PR_SHA: ${{ github.event.pull_request.head.sha }}
+          # GPU_ID: ${{ matrix.gpu_id }}
+          # TARGET_SM: ${{ matrix.target_sm }}
+          # KERNEL_ALIGN_FORCE_SM90: ${{ matrix.force_sm90 }}
         run: bash ci/run_gpu_ci.sh

ci/run_gpu_ci.sh

@@ -1,13 +1,20 @@
 #!/usr/bin/env bash
 set -uo pipefail
 
-# TP=2
-PRIMARY_GPU_ID="NVIDIA RTX A4000"
-PRIMARY_GPU_COUNT=2
+# TP=2 (override via env; GPU_ID/GPU_COUNT accepted as matrix-friendly aliases)
+PRIMARY_GPU_ID="${PRIMARY_GPU_ID:-${GPU_ID:-NVIDIA RTX A4000}}"
+PRIMARY_GPU_COUNT="${PRIMARY_GPU_COUNT:-${GPU_COUNT:-2}}"
 
 # TP=1
-FALLBACK_GPU_ID="NVIDIA A40"
-FALLBACK_GPU_COUNT=1
+FALLBACK_GPU_ID="${FALLBACK_GPU_ID:-NVIDIA A40}"
+FALLBACK_GPU_COUNT="${FALLBACK_GPU_COUNT:-1}"
+
+# Optional arch override; asserted against the pod's real cap in the remote build so a
+# cross-arch resource fallback cannot build mismatched SASS.
+TARGET_SM="${TARGET_SM:-}"
+
+# Forwarded to the remote build; setup.py compiles the Hopper (sm90) kernels only when "1".
+KERNEL_ALIGN_FORCE_SM90="${KERNEL_ALIGN_FORCE_SM90:-}"
 
 CI_IMAGE="${CI_IMAGE:-runpod/pytorch:2.4.0-py3.11-cuda12.4.1-devel-ubuntu22.04}"
 DISK_GB=40
@@ -127,17 +134,64 @@ if ! "$PY" -c "import torch" >/dev/null 2>&1; then
   done
 fi
 echo "[remote] Using interpreter: $PY"
-export TORCH_CUDA_ARCH_LIST=8.6
 export FORCE_CUDA=1
 export MAX_JOBS=8
+export KERNEL_ALIGN_FORCE_SM90="'"${KERNEL_ALIGN_FORCE_SM90}"'"
+
+# normalize_sm: compact (90) or dotted (9.0) compute cap -> torch dotted form, keeping +PTX.
+normalize_sm() {
+  sm_in="$1"; sm_ptx=""
+  case "$sm_in" in *+PTX) sm_ptx="+PTX"; sm_in="${sm_in%+PTX}";; esac
+  case "$sm_in" in
+    *.*) : ;;
+    [0-9][0-9]|[0-9][0-9][0-9]) sm_major="${sm_in%?}"; sm_in="${sm_major}.${sm_in#$sm_major}" ;;
+    *) return 1 ;;
+  esac
+  case "$sm_in" in [0-9]*.[0-9]|[0-9]*.[0-9][0-9]) echo "${sm_in}${sm_ptx}" ;; *) return 1 ;; esac
+}
+
+ACTUAL_SM=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null | head -1 | tr -d "[:space:]")
+[ -z "$ACTUAL_SM" ] && ACTUAL_SM=$("$PY" -c "import torch;a,b=torch.cuda.get_device_capability();print(f\"{a}.{b}\")" 2>/dev/null || true)
+[ -z "$ACTUAL_SM" ] && { echo "[remote] FATAL: cannot determine GPU compute capability"; exit 3; }
+
+REQUESTED_SM="'"${TARGET_SM}"'"
+if [ -n "$REQUESTED_SM" ]; then
+  NORM_REQ=$(normalize_sm "$REQUESTED_SM") || { echo "[remote] FATAL: unsupported TARGET_SM=$REQUESTED_SM"; exit 3; }
+  NORM_REQ_BASE="${NORM_REQ%+PTX}"
+  if [ "$NORM_REQ_BASE" != "$ACTUAL_SM" ]; then
+    echo "[remote] FATAL: requested TARGET_SM=$REQUESTED_SM (sm_$NORM_REQ_BASE) but provisioned GPU is sm_$ACTUAL_SM."
+    echo "[remote]        Refusing to build mismatched kernels (likely a cross-arch resource fallback)."
+    exit 3
+  fi
+  BUILD_SM="$NORM_REQ_BASE"
+else
+  BUILD_SM=$(normalize_sm "$ACTUAL_SM") || { echo "[remote] FATAL: unsupported detected arch $ACTUAL_SM"; exit 3; }
+fi
+# BUILD_SM is always bare here (both paths strip +PTX); +PTX gives forward-compat JIT.
+export TORCH_CUDA_ARCH_LIST="${BUILD_SM}+PTX"
+echo "[remote] Detected GPU sm_$ACTUAL_SM; building _C for TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST"
+
 cd /workspace
 git clone '"${PR_REPO_URL:-https://github.com/RL-Align/RL-Kernel.git}"' repo
 cd repo
 git fetch origin '"${PR_SHA}"'
 git checkout --detach '"${PR_SHA}"'
-"$PY" -m pip install -e .
-"$PY" -m pip install pytest
+"$PY" -c "import torch;print(f\"[remote] image torch {torch.__version__} cuda {torch.version.cuda}\")"
+# Pin torch (cu124, matching the CI image) so the extension is built against the exact
+# runtime torch, not a non-deterministic bare-install upgrade of the 2.4.0 in the image.
+TORCH_SPEC="${TORCH_SPEC:-torch==2.4.1}"
+TORCH_INDEX_URL="${TORCH_INDEX_URL:-https://download.pytorch.org/whl/cu124}"
+"$PY" -m pip install --no-cache-dir "$TORCH_SPEC" --index-url "$TORCH_INDEX_URL"
+"$PY" -c "import torch;print(f\"[remote] pinned torch {torch.__version__} cuda {torch.version.cuda}\")"
+# --no-build-isolation: torch must be visible to setup.py, else the extension is silently skipped.
+# --no-deps: keep the pinned torch; do not let the editable install re-resolve it.
+"$PY" -m pip install --no-build-isolation --no-deps -e .
+"$PY" -m pip install --no-cache-dir numpy tabulate accelerate transformers pytest
 nvidia-smi
+# Fail fast if _C did not build or cannot launch, instead of silently using native fallbacks.
+"$PY" scripts/ci_smoke.py
+# Enforce _C in the pytest suite too (test_extension_smoke.py skips unless this is set).
+export RL_KERNEL_REQUIRE_EXT=1
 '"${TEST_CMD}"
 
 echo "[ci] Launching remote test suite on GPU pod (Distributed Execution Mode: TP=${GPU_COUNT})..."

docs/getting_started/installation.md

@@ -5,24 +5,44 @@ CUDA toolchain; ROCm builds require a compatible ROCm environment.
 
 ## From Source
 
+For CUDA (or ROCm) source builds, install PyTorch first (matching your CUDA/ROCm
+runtime), then build with build isolation disabled so `setup.py` can access
+PyTorch's extension build utilities and compile the native kernels (`rl_engine._C`):
+
 ```bash
 git clone https://github.com/RL-Align/RL-Kernel.git
 cd RL-Kernel
-pip install -e .
+# Optional: pin the compile target. If unset, the build targets your GPU's arch.
+# export TORCH_CUDA_ARCH_LIST="9.0+PTX"   # e.g. Hopper; or "8.6+PTX", "12.0+PTX"
+pip install --no-build-isolation -e .
 ```
 
+Without `--no-build-isolation`, PyTorch is invisible to the isolated build
+environment, the extension is silently skipped, and the library falls back to the
+slower pure-PyTorch kernels. Confirm the compiled extension is present with:
+
+```bash
+python -c "from rl_engine import _C; print('compiled extension OK')"
+```
+
+A CPU-only install (plain `pip install -e .` on a machine with no GPU) remains
+supported and runs on the pure-PyTorch backends.
+
 ## Optional Backends
 
+The extras add optional dependencies on top of the compiled package, so they use
+the same `--no-build-isolation` flag as the source build above.
+
 ```bash
-pip install -e ".[cuda]"
+pip install --no-build-isolation -e ".[cuda]"
 ```
 
 ```bash
-pip install -e ".[rocm]"
+pip install --no-build-isolation -e ".[rocm]"
 ```
 
 ```bash
-pip install -e ".[vllm]"
+pip install --no-build-isolation -e ".[vllm]"
 ```
 
 Install the vLLM extra only on rollout or benchmark environments that need the

scripts/ci_smoke.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+"""CI fail-fast smoke check for the compiled CUDA extension (``rl_engine._C``).
+
+Exits non-zero with a clear message when either failure mode from issue #191
+occurs, so GPU CI cannot pass while silently running only native fallbacks:
+
+  * Bug A - the extension did not build at all (e.g. PEP 517 build isolation hid
+    torch from ``setup.py``), so ``from rl_engine import _C`` raises ImportError.
+  * Bug B - the extension built for the wrong GPU architecture; the import
+    succeeds (``dlopen`` does not check arch) but the first kernel launch raises
+    ``cudaErrorNoKernelImageForDevice`` once the stream is synchronized.
+
+Uses only ``fused_logp`` - the op registered unconditionally in ``csrc/ops.cpp`` -
+so it does not require ``KERNEL_ALIGN_FORCE_SM90=1`` / a Hopper build.
+"""
+import sys
+
+import torch
+
+
+def main() -> int:
+    if not torch.cuda.is_available():
+        print("[smoke] FATAL: CUDA is not available in this CI environment", file=sys.stderr)
+        return 2
+
+    print(f"[smoke] torch: {torch.__version__} (cuda {torch.version.cuda})")
+    print(f"[smoke] device: {torch.cuda.get_device_name()}")
+    cc = torch.cuda.get_device_capability()
+    print(f"[smoke] capability: sm_{cc[0]}{cc[1]}")
+
+    # (1) Import check -> catches Bug A (no .so was compiled at all).
+    try:
+        from rl_engine import _C
+    except ImportError as exc:
+        print(
+            "[smoke] FATAL: compiled extension rl_engine._C is missing - the CUDA "
+            "kernels were not built.\n"
+            "        Likely PEP 517 build isolation hid torch from setup.py; install "
+            "with `pip install --no-build-isolation -e .`.\n"
+            f"        Underlying error: {exc}",
+            file=sys.stderr,
+        )
+        return 1
+    print(f"[smoke] _C file: {getattr(_C, '__file__', None)}")
+
+    # (2) Real launch + synchronize -> catches Bug B (arch mismatch only surfaces
+    #     on launch, asynchronously, so the sync is required to raise it here).
+    try:
+        logits = torch.randn(4, 32, device="cuda", dtype=torch.float32)
+        token_ids = torch.randint(0, 32, (4,), device="cuda", dtype=torch.long)
+        out = _C.fused_logp(logits, token_ids)
+        torch.cuda.synchronize()
+    # Broad on purpose: an arch mismatch raises a CUDA RuntimeError (not ImportError),
+    # and any launch failure whatsoever must fail the smoke check loudly.
+    except Exception as exc:
+        print(
+            "[smoke] FATAL: rl_engine._C built but fused_logp failed to launch on "
+            f"sm_{cc[0]}{cc[1]}.\n"
+            "        The extension was likely compiled for a different architecture; "
+            "set TORCH_CUDA_ARCH_LIST / TARGET_SM to match this GPU.\n"
+            f"        Underlying error: {type(exc).__name__}: {exc}",
+            file=sys.stderr,
+        )
+        return 1
+
+    if tuple(out.shape) != (4,):
+        print(
+            f"[smoke] FATAL: unexpected fused_logp output shape {tuple(out.shape)}", file=sys.stderr
+        )
+        return 1
+
+    print(f"[smoke] OK: rl_engine._C built and fused_logp ran on sm_{cc[0]}{cc[1]}.")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tests/test_extension_smoke.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+"""Regression guard for issue #191.
+
+On a GPU host the compiled extension ``rl_engine._C`` MUST be present and
+launchable; a missing or arch-mismatched build must fail loudly here instead of
+silently degrading to the pure-PyTorch fallbacks. On a CPU host (no CUDA) the
+test skips, so the CPU CI job - which legitimately has no ``_C`` - stays green.
+
+Enforcement is opt-in via ``RL_KERNEL_REQUIRE_EXT=1``, which the GPU CI
+orchestrator sets right before pytest (after it has built the extension). That
+way the check does not fail in environments not expected to have ``_C`` compiled -
+a plain ``pytest`` run, or CI running an older orchestrator that has not built it
+yet (the compiled kernels are always enforced separately by scripts/ci_smoke.py).
+"""
+import os
+
+import pytest
+import torch
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA GPU")
+@pytest.mark.skipif(
+    os.environ.get("RL_KERNEL_REQUIRE_EXT") != "1",
+    reason="extension enforcement is CI-only (orchestrator sets RL_KERNEL_REQUIRE_EXT=1)",
+)
+def test_compiled_extension_present_and_launches():
+    # Import directly from the package, NOT from rl_engine.kernels.ops.base, which
+    # deliberately swallows the ImportError and falls back to _C = None.
+    try:
+        from rl_engine import _C
+    except ImportError as exc:  # Bug A: the extension was never built
+        pytest.fail(
+            "rl_engine._C is missing on a GPU host - the CUDA extension was not "
+            "built. Install with `pip install --no-build-isolation -e .`. "
+            f"Underlying error: {exc}"
+        )
+
+    logits = torch.randn(4, 32, device="cuda", dtype=torch.float32)
+    token_ids = torch.randint(0, 32, (4,), device="cuda", dtype=torch.long)
+    out = _C.fused_logp(logits, token_ids)
+    torch.cuda.synchronize()  # Bug B: an arch mismatch surfaces on synchronize
+    assert tuple(out.shape) == (4,)