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

[tianleiwu]

Summary

GroupQueryAttention's XQA decode kernel failed on consumer Blackwell GPUs (RTX 50-series, sm_120) with cudaErrorInvalidValue, while working fine on A100 (sm_80) and H200 (sm_90). This adds a runtime shared-memory capability check so XQA is only selected when the device can actually satisfy the kernel's dynamic shared-memory request, and otherwise falls back to cuDNN SDPA / Flash. A100/H200 behavior is unchanged.

Root cause

XQA bakes its shared-memory layout at compile time from __CUDA_ARCH__:

• sm_80 / sm_87 / sm_90 use the large K/V-tile layout (preferedKHeadPartBytes=128, cacheVTileSeqLen=64) → up to ~140 KB of dynamic shared memory for head_size 128/256.
• sm_86 / sm_89 / sm_120 use the small layout (64 / 32) → ~78–96 KB.

Release/packaging binaries are built with a maximum arch of 90-virtual (compute_90 PTX only, no native sm_120 SASS). On sm_120 the driver JIT-compiles that sm_90 PTX, so the kernel's smemSize carries the Hopper value (~140 KB). launchMHA then calls cudaFuncSetAttribute(..., cudaFuncAttributeMaxDynamicSharedMemorySize, size), which exceeds sm_120's ~99 KB per-block opt-in limit (sharedMemPerBlockOptin) and returns cudaErrorInvalidValue. A100 (163 KB) and H200 (227 KB) have enough room, so they were unaffected.

Key changes

File	Change
xqa/xqa_impl_gen.cuh	Add GetSmemSize() host helper that reads the per-kernel smemSize device symbol (accurate even for a PTX kernel JIT-compiled for the running SM).
xqa/xqa_loader_fp16_impl.cuh	Add GetXQAKernelSmemBytes(group_size) head-dim dispatcher. The non-quantized fp16 footprint is an upper bound for the int8/fp8/bf16 variants (smaller cache element), so one query covers all XQA paths.
xqa/xqa_loader_fp16.cu, xqa/xqa_loader.h	Expose GetXQARequiredSharedMemoryBytes(device_prop, head_size, num_heads, kv_num_heads); a single non-templated entry point used by both the fp16 and bf16 GQA kernels.
group_query_attention.cc, group_query_attention.h	Gate XQA selection on required_smem <= device_prop.sharedMemPerBlockOptin; fall back to cuDNN SDPA / Flash when it does not fit. Result is cached per node (xqa_shared_memory_ok_) since head_size/group are constant.
xqa/mha_impl.cuh	Defensive backstop in launchMHA: if the requested shared memory still exceeds the device limit, throw an actionable message (which SM to build for / how to disable XQA) instead of the opaque cudaErrorInvalidValue.

CUDA graph safety

GetXQARequiredSharedMemoryBytes uses cudaMemcpyFromSymbol, which synchronizes and is illegal during CUDA graph capture. The query is:

• cached per node, so it runs at most once;
• guarded with onnxruntime::llm::common::isCapturing(Stream(context)) so the synchronizing call is only issued when the compute stream is not capturing;
• resolved during ORT's non-captured warm-up run(s) before capture begins.

If the value is somehow still unresolved while capturing, XQA is conservatively skipped for that run (safe fallback) without caching, so a later non-capturing run can resolve it. Warm-up and capture therefore make the same XQA/fallback decision, keeping the captured graph consistent with replay.

Testing notes

• Built the affected TUs (GQA dispatcher + fp16/bf16/int8/fp8 XQA loaders) with CMAKE_CUDA_ARCHITECTURES="80;90" (the configuration that reproduces the failure); all compile cleanly.
• To validate the fix end-to-end, run a fp16/bf16 GQA decode workload on an sm_120 GPU (e.g. RTX 5090): it should now run (via fallback) instead of returning cudaErrorInvalidValue. Set ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO=1 to confirm the selected backend.
• To actually run XQA (the fast path) on Blackwell, build with native arch 120 in CMAKE_CUDA_ARCHITECTURES (and 100 for datacenter Blackwell). With a native sm_120 cubin the layout is ~80 KB and fits, so XQA is selected.

[Wayne-Ch]

Approving from the Foundry Local customer validation perspective. This directly matches the HY-MT1.5/HY-MT2 RTX 5090 failure we reproduced: the reduced/no-cuDNN CUDA EP path fails in GroupQueryAttention with cudaErrorInvalidValue, while the full/cuDNN EP path passes on the same machine. Assuming CI completes successfully, this is needed to unblock the next validation package.