[bugfix]: select architecture-specific Blackwell kernel targets

- Detect GB200 and GB300 as 10.0a and 10.3a
- Build future GB200 CI images with 10.0a
- Add regression tests for detection, cache keys, and build arguments

.github/workflows/infra-build-image.yml

@@ -203,7 +203,8 @@ jobs:
           docker buildx imagetools create "${TAG_ARGS[@]}" "${IMAGE_REFS[@]}"
           docker buildx imagetools inspect "${TAGS[0]}"
 
-  # The CI runner is ARM64 like DGX Spark, but targets sm_100 rather than sm_121.
+  # The CI runner is ARM64 like DGX Spark, but targets sm_100a rather than sm_121.
+  # The architecture-specific target includes the GB200 VSA CUDA extensions.
   # Publish a single-architecture variant so the self-hosted CI runner can reuse
   # the exact prebuilt kernel instead of compiling it in every job.
   build-ci-runner-image:
@@ -219,7 +220,7 @@ jobs:
         PYTHON_VERSION=3.12
         CUDA_VERSION=13.0.0
         UV_TORCH_BACKEND=cu130
-        TORCH_CUDA_ARCH_LIST=10.0
+        TORCH_CUDA_ARCH_LIST=10.0a
         CMAKE_BUILD_PARALLEL_LEVEL=1
         FLASH_ATTN_WHEEL_TAG=cu130torch2.12
         FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.22

docs/contributing/ci_architecture.md

@@ -440,8 +440,11 @@ before later jobs consume the updated image pin.
 The same workflow publishes a single-architecture ARM64, CUDA 13, SM100 image
 for the self-hosted CI runner under the
 `py3.12-cuda13.0.0-sm100-{latest,sha-*}` tags. It carries the matching prebuilt
-kernel wheel so runner jobs can validate and install the exact source and ABI
-match instead of recompiling it in every lane.
+kernel wheel compiled with `TORCH_CUDA_ARCH_LIST=10.0a` to include the GB200
+VSA CUDA extensions. Runtime kernel detection uses the same target, so runner
+jobs can validate and install the exact source and ABI match instead of
+recompiling it in every lane. Older artifacts built for `10.0` have a different
+cache key and trigger a local rebuild when the worker detects `10.0a`.
 
 The optional Dreamverse matrix builds backend and UI images for CUDA 12.6 and
 CUDA 13 on `amd64`. Dreamverse remains `amd64`-only because its FA4 dependency

fastvideo/tests/modal/kernel_build_cache.py

@@ -147,6 +147,10 @@ def _detect_arch_from_torch() -> str:
         major, minor = torch.cuda.get_device_capability(0)
         if major == 9 and minor == 0:
             return "9.0a"
+        if major == 10 and minor in (0, 3):
+            # Match build.sh: data-center Blackwell VSA kernels require the
+            # architecture-specific target, not the generic compute capability.
+            return f"{major}.{minor}a"
         if major == 12 and minor == 0:
             return "12.0a"
         return f"{major}.{minor}"

fastvideo/tests/modal/test_kernel_build_cache.py

@@ -5,9 +5,12 @@ from __future__ import annotations
 import errno
 import importlib.util
 import json
+import shutil
 import subprocess
+import sys
 import zipfile
 from pathlib import Path
+from types import SimpleNamespace
 
 import pytest
 import yaml
@@ -168,7 +171,7 @@ def test_ci_runner_image_targets_arm64_sm100_with_opencv_runtime() -> None:
     assert job["with"]["tag_suffix"] == "py3.12-cuda13.0.0-sm100"
     assert "CUDA_VERSION=13.0.0" in job["with"]["build_args"]
     assert "UV_TORCH_BACKEND=cu130" in job["with"]["build_args"]
-    assert "TORCH_CUDA_ARCH_LIST=10.0" in job["with"]["build_args"]
+    assert "TORCH_CUDA_ARCH_LIST=10.0a" in job["with"]["build_args"].splitlines()
 
     dockerfile = (REPO_ROOT / "docker/Dockerfile").read_text(encoding="utf-8")
     assert "    ffmpeg \\\n" in dockerfile
@@ -191,23 +194,68 @@ def test_docker_image_bakes_modal_apt_and_rust_layer() -> None:
     assert "ENV PATH=/root/.cargo/bin:${PATH}" in dockerfile
 
 
-def test_cache_key_uses_resolved_arch_not_raw_env(monkeypatch, tmp_path) -> None:
-    _patch_stable_metadata(monkeypatch)
-    monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: "9.0a")
+def _patch_cuda_capability(monkeypatch, capability) -> None:
+    cuda = SimpleNamespace(is_available=lambda: True, get_device_capability=lambda device: capability)
+    monkeypatch.setitem(sys.modules, "torch", SimpleNamespace(cuda=cuda))
 
-    monkeypatch.setenv("TORCH_CUDA_ARCH_LIST", "9.0a")
-    explicit_hopper = kernel_build_cache._build_metadata(tmp_path)
+
+@pytest.mark.parametrize("capability, expected", [
+    ((8, 9), "8.9"),
+    ((9, 0), "9.0a"),
+    ((10, 0), "10.0a"),
+    ((10, 3), "10.3a"),
+    ((12, 0), "12.0a"),
+    ((12, 1), "12.1"),
+])
+def test_detected_arch_enables_device_specific_kernels(monkeypatch, capability, expected) -> None:
+    _patch_cuda_capability(monkeypatch, capability)
+
+    assert kernel_build_cache._detect_arch_from_torch() == expected
+
+
+@pytest.mark.parametrize("arch", ["9.0a", "10.0a", "10.3a"])
+def test_cache_key_uses_resolved_arch_not_raw_env(monkeypatch, tmp_path, arch) -> None:
+    _patch_stable_metadata(monkeypatch)
+    # An explicit target must win even when the detected GPU differs.
+    monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: "8.9")
+
+    monkeypatch.setenv("TORCH_CUDA_ARCH_LIST", arch)
+    explicit = kernel_build_cache._build_metadata(tmp_path)
 
     monkeypatch.delenv("TORCH_CUDA_ARCH_LIST", raising=False)
-    detected_hopper = kernel_build_cache._build_metadata(tmp_path)
+    monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: arch)
+    detected = kernel_build_cache._build_metadata(tmp_path)
 
-    monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: "8.9")
-    detected_l40s = kernel_build_cache._build_metadata(tmp_path)
+    monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: arch.removesuffix("a"))
+    generic = kernel_build_cache._build_metadata(tmp_path)
 
-    assert explicit_hopper["cache_key"] == detected_hopper["cache_key"]
-    assert explicit_hopper["build"]["torch_cuda_arch_list"] == "9.0a"
-    assert detected_hopper["build"]["torch_cuda_arch_list"] == ""
-    assert detected_l40s["cache_key"] != detected_hopper["cache_key"]
+    assert explicit["cache_key"] == detected["cache_key"]
+    assert explicit["build"]["torch_cuda_arch_list"] == arch
+    assert detected["build"]["torch_cuda_arch_list"] == ""
+    assert generic["cache_key"] != detected["cache_key"]
+
+
+@pytest.mark.parametrize("capability, expected", [((10, 0), "10.0a"), ((10, 3), "10.3a")])
+def test_blackwell_detection_reaches_kernel_build(monkeypatch, tmp_path, capability, expected) -> None:
+    _patch_stable_metadata(monkeypatch)
+    _patch_cuda_capability(monkeypatch, capability)
+    metadata = kernel_build_cache._build_metadata(tmp_path)
+    build_arches = []
+
+    def fake_run(command, *, cwd, env):
+        assert command[:2] == ["./build.sh", "--wheel-dir"]
+        assert cwd == tmp_path / "fastvideo-kernel"
+        build_arches.append(env["TORCH_CUDA_ARCH_LIST"])
+        _write_wheel(Path(command[2]) / WHEEL_NAME)
+        return ""
+
+    monkeypatch.setattr(kernel_build_cache, "_run", fake_run)
+    wheel = kernel_build_cache._build_wheel(tmp_path, metadata)
+    try:
+        assert wheel.is_file()
+        assert build_arches == [expected]
+    finally:
+        shutil.rmtree(wheel.parent)
 
 
 def test_cache_key_ignores_runtime_only_torch_config(monkeypatch, tmp_path) -> None: