[CUDA] Fix XQA GroupQueryAttention cudaErrorInvalidValue on Blackwell (sm_120)

onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc

@@ -20,6 +20,7 @@
 #include "contrib_ops/cuda/bert/unfused_attention.h"
 #include "contrib_ops/cuda/utils/dump_cuda_tensor.h"
 #include "contrib_ops/cpu/utils/debug_macros.h"
+#include "contrib_ops/cuda/llm/common/cuda_runtime_utils.h"
 
 using namespace onnxruntime::cuda;
 using namespace ::onnxruntime::common;
@@ -491,6 +492,37 @@ Status GroupQueryAttention<T, U>::ComputeInternal(OpKernelContext* context) cons
 
     data.use_xqa = (is_non_quantized_supported || is_int8_quantized_supported || is_fp8_quantized_supported);
 
+    if (data.use_xqa) {
+      // Consumer Blackwell (sm_120) and some other devices expose a smaller per-block opt-in
+      // shared-memory limit than the SM the XQA kernel was compiled / JIT-compiled for (sm_80 and
+      // sm_90 use a larger K/V-tile layout, ~140 KB for head_size 128/256). Launching XQA there
+      // fails with cudaErrorInvalidValue because cudaFuncSetAttribute requests more shared memory
+      // than the device allows. Query the actual kernel requirement (read from the device symbol,
+      // so it is accurate even for a PTX kernel JIT-compiled for this SM) and fall back to cuDNN
+      // SDPA / Flash when it does not fit. Cached per node because head_size and group size are
+      // constant for a given GQA node (avoids a device->host copy on every decode step).
+      int xqa_smem_ok = xqa_shared_memory_ok_.load(std::memory_order_relaxed);
+      if (xqa_smem_ok < 0) {
+        // GetXQARequiredSharedMemoryBytes issues a cudaMemcpyFromSymbol, which synchronizes and is
+        // therefore illegal during CUDA graph capture (it would invalidate the capture). The result
+        // is invariant for a node (fixed head_size / group size / device), and ORT performs at least
+        // one non-captured warm-up run before capturing, so this normally resolves and caches during
+        // warm-up. Guard defensively: only issue the synchronizing query when the compute stream is
+        // not capturing. If we somehow reach here for the first time while capturing, conservatively
+        // skip XQA for this run instead of breaking the capture, and leave the cache unresolved so a
+        // later non-capturing run can determine it.
+        if (!onnxruntime::llm::common::isCapturing(Stream(context))) {
+          const size_t required_smem = onnxruntime::contrib::cuda::GetXQARequiredSharedMemoryBytes(
+              device_prop, parameters.head_size, parameters.num_heads, parameters.kv_num_heads);
+          xqa_smem_ok = (required_smem == 0 || required_smem <= device_prop.sharedMemPerBlockOptin) ? 1 : 0;
+          xqa_shared_memory_ok_.store(xqa_smem_ok, std::memory_order_relaxed);
+        } else {
+          xqa_smem_ok = 0;  // capturing (or status query failed) and not yet resolved -> fall back
+        }
+      }
+      data.use_xqa = (xqa_smem_ok != 0);
+    }
+
     if (data.use_xqa) {
       size_t xqa_internal_bytes = onnxruntime::contrib::cuda::GetXQAScratchSize(
           GetDeviceProp(),

onnxruntime/contrib_ops/cuda/bert/group_query_attention.h

@@ -3,6 +3,7 @@
 
 #pragma once
 
+#include <atomic>
 #include <memory>
 #include "core/providers/cuda/cuda_kernel.h"
 #include "contrib_ops/cuda/bert/group_query_attention_impl.h"
@@ -52,6 +53,10 @@ class GroupQueryAttention final : public CudaKernel {
   // FP32 head_sink cached in PrePack for the XQA path (empty when head_sink is not a constant initializer).
   IAllocatorUniquePtr<float> xqa_head_sink_;
   int xqa_head_sink_count_ = 0;  // Number of elements in xqa_head_sink_ (0 when not prepacked).
+  // Cached result of the XQA shared-memory fit check for this node (-1 unknown, 0 does not fit,
+  // 1 fits). head_size and group size are constant for a given GQA node, so the (device-symbol)
+  // query is done once and reused to avoid a per-decode-step device->host copy.
+  mutable std::atomic<int> xqa_shared_memory_ok_{-1};
   const AttentionKernelOptions* kernel_options_;
 };
 

onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh

@@ -30,6 +30,7 @@
 #ifndef GENERATE_CUBIN
 #include "hostUtils.h"
 #include <cuda_runtime.h>
+#include <string>
 #ifndef NDEBUG
 #include <cstdio>
 #endif
@@ -2491,6 +2492,20 @@ void launchMHA(cudaDeviceProp const& prop, uint32_t nbKHeads,
   static uint32_t const hostSmemSize = [&]() {
     uint32_t size;
     checkCuda(cudaMemcpyFromSymbol(&size, smemSize, sizeof(smemSize)));
+    // Defensive backstop: the kernel's shared-memory footprint is fixed at compile time for its
+    // target SM (sm_80/sm_90 use a larger K/V-tile layout than sm_86/sm_89/sm_120). When such a
+    // kernel is JIT-compiled from PTX onto a device with a smaller per-block opt-in limit (e.g.
+    // consumer Blackwell sm_120, ~99 KB), cudaFuncSetAttribute below returns cudaErrorInvalidValue.
+    // The GQA dispatcher already skips XQA in that case (see GetXQARequiredSharedMemoryBytes); this
+    // guard turns any remaining mismatch into an actionable message instead of an opaque CUDA error.
+    if (size > prop.sharedMemPerBlockOptin) {
+      throw std::runtime_error(
+          "XQA kernel requires " + std::to_string(size) +
+          " bytes of shared memory per block but this GPU allows only " +
+          std::to_string(prop.sharedMemPerBlockOptin) +
+          " bytes. Build ONNX Runtime with the device's native CUDA architecture (e.g. add 120 to "
+          "CMAKE_CUDA_ARCHITECTURES for sm_120 / RTX 50-series) or disable XQA (ORT_ENABLE_XQA=0).");
+    }
     checkCuda(cudaFuncSetAttribute(kernel_mha, cudaFuncAttributeMaxDynamicSharedMemorySize, size));
     return size;
   }();

onnxruntime/contrib_ops/cuda/bert/xqa/xqa_impl_gen.cuh

@@ -134,5 +134,30 @@ inline Status Launch(
   return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "XQA is only supported on Ampere (SM80) or newer GPUs.");
 #endif
 }
+
+#ifndef GENERATE_CUBIN
+// Host helper: dynamic shared-memory size (bytes) this kernel instantiation requests at launch,
+// read from the device `smemSize` symbol. Because it is read from the loaded module, the value
+// reflects the ACTUAL kernel that will run -- including a kernel JIT-compiled from PTX for a
+// different SM, whose shared-memory layout (and therefore smemSize) was fixed at PTX-generation
+// time. The GQA dispatcher uses this to avoid selecting XQA on devices whose per-block opt-in
+// shared memory is smaller than the kernel needs, which would otherwise fail at launch with
+// cudaErrorInvalidValue (e.g. consumer Blackwell sm_120 running a kernel JIT'd from sm_90 PTX,
+// where the Hopper layout needs ~140 KB but sm_120 allows only ~99 KB). Returns 0 if the value
+// cannot be queried.
+inline size_t GetSmemSize() {
+#ifdef XQA_HAS_SM80_TARGET
+  uint32_t size = 0;
+  if (cudaMemcpyFromSymbol(&size, smemSize, sizeof(smemSize)) != cudaSuccess) {
+    (void)cudaGetLastError();  // clear any sticky error so it does not leak to the next CUDA call
+    return 0;
+  }
+  return static_cast<size_t>(size);
+#else
+  return 0;
+#endif
+}
+#endif  // GENERATE_CUBIN
+
 #undef XQA_HAS_SM80_TARGET
 }  // namespace NAMESPACE_NAME

onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader.h

@@ -54,6 +54,18 @@ size_t GetXQAScratchSize(
     XqaQuantType kv_quant_type,
     bool is_bf16 = false);
 
+// Dynamic shared memory (bytes) the XQA decode kernel requires for the given head size and
+// query/KV group size on this device. The value is read from the loaded kernel module, so it is
+// accurate even when the kernel is a PTX image JIT-compiled for the running SM. Returns 0 when it
+// cannot be determined (e.g. unsupported head size / group size). Callers should skip XQA when the
+// returned value exceeds device_prop.sharedMemPerBlockOptin (otherwise the launch fails with
+// cudaErrorInvalidValue, e.g. consumer Blackwell sm_120 running a kernel built only for sm_90).
+size_t GetXQARequiredSharedMemoryBytes(
+    const cudaDeviceProp& device_prop,
+    int head_size,
+    int num_heads,
+    int kv_num_heads);
+
 }  // namespace cuda
 }  // namespace contrib
 }  // namespace onnxruntime

onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16.cu

@@ -35,6 +35,8 @@ Status LaunchXQAKernelImpl(
     void* workspace,
     size_t workspace_size);
 
+size_t GetXQAKernelSmemBytes(int group_size);
+
 }  // namespace H128
 
 namespace H64 {
@@ -61,6 +63,8 @@ Status LaunchXQAKernelImpl(
     void* workspace,
     size_t workspace_size);
 
+size_t GetXQAKernelSmemBytes(int group_size);
+
 }  // namespace H64
 
 namespace H256 {
@@ -87,6 +91,8 @@ Status LaunchXQAKernelImpl(
     void* workspace,
     size_t workspace_size);
 
+size_t GetXQAKernelSmemBytes(int group_size);
+
 }  // namespace H256
 
 // Dispatcher Implementation
@@ -178,6 +184,28 @@ size_t GetXQAScratchSize(
   return roundUp<size_t>(semaphore_size, 128) + scratch_size;
 }
 
+size_t GetXQARequiredSharedMemoryBytes(
+    const cudaDeviceProp& device_prop,
+    int head_size,
+    int num_heads,
+    int kv_num_heads) {
+  if (device_prop.major < 8 || kv_num_heads <= 0) {
+    return 0;
+  }
+  const int group_size = num_heads / kv_num_heads;
+  // The dtype (fp16 vs bf16) does not change the shared-memory footprint (both are 2-byte
+  // elements), so the fp16 kernels are queried for both. The non-quantized kernel is an upper
+  // bound for the int8/fp8 variants, so a single query covers all XQA paths.
+  if (head_size == 256) {
+    return H256::GetXQAKernelSmemBytes(group_size);
+  } else if (head_size == 128) {
+    return H128::GetXQAKernelSmemBytes(group_size);
+  } else if (head_size == 64) {
+    return H64::GetXQAKernelSmemBytes(group_size);
+  }
+  return 0;
+}
+
 // Instantiate template for half
 template Status LaunchXQAKernel<half>(
     const cudaDeviceProp& device_prop,

onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_impl.cuh

@@ -226,6 +226,37 @@ Status LaunchXQAKernelImpl(
   }
 }
 
+#ifndef GENERATE_CUBIN
+// Returns the dynamic shared-memory (bytes) the non-quantized XQA kernel for this head dim and
+// group size requests at launch (read from the device symbol, so it is accurate even for a PTX
+// kernel JIT-compiled for the running SM). Used by the GQA dispatcher to skip XQA when the
+// device's per-block opt-in shared memory is too small. The non-quantized kernel is an upper
+// bound for the int8/fp8 variants (smaller cache element), so this single query also guards the
+// quantized paths. Not marked inline: it is defined in exactly one TU per head-dim namespace
+// (xqa_loader_fp16_{64,128,256}.cu) and called from the head-size dispatcher TU, so it needs an
+// externally linkable definition.
+size_t GetXQAKernelSmemBytes(int group_size) {
+  switch (group_size) {
+    case 1:
+      return grp1_fp16::GetSmemSize();
+    case 2:
+      return grp2_fp16::GetSmemSize();
+    case 4:
+      return grp4_fp16::GetSmemSize();
+    case 5:
+      return grp5_fp16::GetSmemSize();
+    case 8:
+      return grp8_fp16::GetSmemSize();
+    case 16:
+      return grp16_fp16::GetSmemSize();
+    case 32:
+      return grp32_fp16::GetSmemSize();
+    default:
+      return 0;
+  }
+}
+#endif  // GENERATE_CUBIN
+
 // Instantiate template for half
 
 }  // namespace HEAD_DIM_NAMESPACE

onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h

@@ -22,6 +22,7 @@
 #include <cuda_fp8.h>
 #endif
 #include "core/providers/cuda/shared_inc/cuda_call.h"
+#include "core/platform/env_var_utils.h"
 
 namespace onnxruntime::llm::common {
 inline int getDevice() {
@@ -61,8 +62,7 @@ inline std::optional<bool> isCudaLaunchBlocking() {
   thread_local bool firstCall = true;
   thread_local std::optional<bool> result = std::nullopt;
   if (!firstCall) {
-    char const* env = std::getenv("CUDA_LAUNCH_BLOCKING");
-    if (env != nullptr && std::string(env) == "1") {
+    if (ParseEnvironmentVariableWithDefault<int>("CUDA_LAUNCH_BLOCKING", 0) == 1) {
       result = true;
     } else {
       result = false;