fix(driver/cuda): gate FlashInfer Mamba SM90 SSU on min arch, not any arch (#438)

The SM90 SSU vertical/horizontal kernels use CCCL's experimental TMA
intrinsics (cuda::device::experimental::cp_async_bulk_tensor_* /
fence_proxy_async_shared_cta), which <cuda/barrier> exposes only behind
__cccl_lib_experimental_ctk12_cp_async_exposure — defined solely when
__CUDA_MINIMUM_ARCH__ >= 900. That macro is the *minimum* of the whole
CUDA_ARCHITECTURES list and is constant across device passes, so a fat
multi-arch build (e.g. the Docker image's 80;86;89;90) pins it to 800 and
#ifdef's the intrinsics out of every pass — even the sm_90 one — breaking
the build with "namespace cuda::device::experimental has no member ...".

flashinfer_mamba.cu built fine at v0.3.0 (no mamba kernel); the SSU kernels
arrived in v0.4.0 and the gate enabled FLASHINFER_MAMBA_ENABLE_SM90 whenever
*any* targeted arch was 90/100, which a mixed build can never satisfy.
Enable it only when *every* targeted arch is sm_90+ (a pure Hopper/Blackwell
build, e.g. the per-compute-capability release binaries); mixed builds fall
back to the simple SSU algorithm. flashinfer's kAuto dispatch already selects
kSimple on every device (incl. Hopper) when FLASHINFER_MAMBA_ENABLE_SM90 is
unset, so this is correct — just not the TMA-optimized path in a fat binary.

Hopper/FA3 attention gating is unchanged (it compiles fine in fat builds; it
uses CUTLASS TMA, not the CCCL experimental gate).

Unbreaks the slim CUDA Docker image build added in #363. Verified e2e on
yecl-gpu-02 (RTX 4090, CUDA 12.9): image builds (4.11 GB runtime), pie doctor
detects the GPU with cuda_native compiled in, and `pie run text-completion`
returns a coherent completion.

Author: shsym

driver/cuda/CMakeLists.txt

@@ -119,24 +119,28 @@ endif()
 
 set(PIE_CUDA_FLASHINFER_HOPPER_SOURCE src/ops/attention_flashinfer_hopper_stub.cpp)
 set(PIE_CUDA_FLASHINFER_MAMBA_SM90 OFF)
+set(_PIE_CUDA_HAS_SM90PLUS OFF)
+set(_PIE_CUDA_HAS_PRE_SM90 OFF)
 foreach(PIE_CUDA_ARCH IN LISTS CMAKE_CUDA_ARCHITECTURES)
   string(REGEX REPLACE "[^0-9].*$" "" PIE_CUDA_ARCH_NUMBER "${PIE_CUDA_ARCH}")
   if(PIE_CUDA_ARCH_NUMBER STREQUAL "90")
     set(PIE_CUDA_FLASHINFER_HOPPER_SOURCE
       src/ops/attention_flashinfer_hopper.cu
       src/ops/attention_xqa_gqa8_sm90.cu
     )
-    set(PIE_CUDA_FLASHINFER_MAMBA_SM90 ON)
   endif()
-  # FlashInfer's Mamba SSU vertical/horizontal kernels use TMA +
-  # memcpy_async_tx, which Blackwell also supports. The runtime
-  # algorithm auto-selector in selective_state_update prefers the
-  # horizontal kernel for sm >= 10 with bf16/fp16 state — so the same
-  # FLASHINFER_MAMBA_ENABLE_SM90 compile gate should be on for sm_100.
-  if(PIE_CUDA_ARCH_NUMBER STREQUAL "100")
-    set(PIE_CUDA_FLASHINFER_MAMBA_SM90 ON)
+  if(PIE_CUDA_ARCH_NUMBER GREATER_EQUAL 90)
+    set(_PIE_CUDA_HAS_SM90PLUS ON)
+  else()
+    set(_PIE_CUDA_HAS_PRE_SM90 ON)
   endif()
 endforeach()
+# SM90 SSU uses CCCL TMA intrinsics that compile only when __CUDA_MINIMUM_ARCH__
+# >= 900 — i.e. the whole arch list is sm_90+. Enable the gate only for pure
+# sm_90+ builds; mixed builds fall back to the simple SSU algo (kAuto picks kSimple).
+if(_PIE_CUDA_HAS_SM90PLUS AND NOT _PIE_CUDA_HAS_PRE_SM90)
+  set(PIE_CUDA_FLASHINFER_MAMBA_SM90 ON)
+endif()
 if(PIE_CUDA_FLASHINFER_HOPPER_SOURCE MATCHES "stub")
   message(STATUS "Disabling FlashInfer Hopper/FA3 attention for CUDA architecture ${CMAKE_CUDA_ARCHITECTURES}")
 else()