[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
This commit is contained in:
Satyam Srivastava
2026-10-05 15:42:48 -07:00
parent e71d01648c
commit 0ef47caf86
4 changed files with 73 additions and 17 deletions
+3 -2
View File
@@ -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