Initialize large-head FMHA shared memory per kernel variant

onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/fmha_launch_template.h

@@ -233,23 +233,36 @@ void LaunchCutlassFmha(const MemoryEfficientAttentionParams& params) {
     p.window_size = params.local_window_size;
   }
 
-  auto kernel_fn = attention_kernel_batched_impl<Attention>;
-
-  if (params.has_custom_right_padding) {
-    kernel_fn = attention_kernel_batched_impl_right_padding<Attention, queries_per_block>;
-  }
-
   int smem_bytes = sizeof(typename Attention::SharedStorage);
   if (smem_bytes > 0xc000) {
     ORT_ENFORCE(params.sm >= 70, "This kernel requires too much shared memory on this machine!");
-    static bool once = [&]() {
-      cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
-      return true;
-    }();
+    if (params.has_custom_right_padding) {
+      static bool right_padding_once = [&]() {
+        CUDA_CALL_THROW(cudaFuncSetAttribute(
+            attention_kernel_batched_impl_right_padding<Attention, queries_per_block>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes));
+        return true;
+      }();
+      ORT_UNUSED_PARAMETER(right_padding_once);
+    } else {
+      static bool default_once = [&]() {
+        CUDA_CALL_THROW(cudaFuncSetAttribute(
+            attention_kernel_batched_impl<Attention>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes));
+        return true;
+      }();
+      ORT_UNUSED_PARAMETER(default_once);
+    }
   }
 
   ORT_ENFORCE(Attention::check_supported(p));
-  kernel_fn<<<p.getBlocksGrid(), p.getThreadsGrid(), smem_bytes, params.stream>>>(p);
+  if (params.has_custom_right_padding) {
+    attention_kernel_batched_impl_right_padding<Attention, queries_per_block>
+        <<<p.getBlocksGrid(), p.getThreadsGrid(), smem_bytes, params.stream>>>(p);
+  } else {
+    attention_kernel_batched_impl<Attention>
+        <<<p.getBlocksGrid(), p.getThreadsGrid(), smem_bytes, params.stream>>>(p);
+  }
 }
 
 template <typename T, typename ArchTag, int queries_per_block, int keys_per_block, int max_head_size>

onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py

@@ -1535,6 +1535,60 @@ def test_gqa_large_head_unfused_softcap_additive_mask_poison_fp16(self):
         self.assertLess(out.float().max().item(), 1.0)
 
 
+@unittest.skipIf(not has_cuda_device(53), "CUDA device not available, skipping large head MEA tests.")
+@patch.dict(os.environ, {"ORT_DISABLE_FLASH_ATTENTION": "1", "ORT_DISABLE_MEMORY_EFFICIENT_ATTENTION": "0"})
+class TestONNXAttentionGQALargeHeadNonpadMEA(unittest.TestCase):
+    """Large-head normal and right-padding MEA kernels initialize independently."""
+
+    def test_gqa_large_head_nonpad_seqlen_after_default_mea_fp16(self):
+        # Keep this order: the regression requires the default kernel to opt in
+        # first, then verifies that the distinct right-padding kernel does too.
+        default_config = AttentionConfig(
+            batch_size=2,
+            q_sequence_length=32,
+            kv_sequence_length=32,
+            past_kv_sequence_length=0,
+            q_num_heads=8,
+            kv_num_heads=4,
+            head_size=512,
+            is_causal=0,
+        )
+        parity_check_gqa_prompt(
+            config=default_config,
+            ep="CUDAExecutionProvider",
+            device="cuda",
+            torch_type=torch.float16,
+            ort_type=TensorProto.FLOAT16,
+            causal=False,
+            rtol=rtol["fp16"],
+            atol=atol["fp16"],
+        )
+
+        nonpad_config = AttentionConfig(
+            batch_size=2,
+            q_sequence_length=32,
+            kv_sequence_length=32,
+            past_kv_sequence_length=0,
+            q_num_heads=8,
+            kv_num_heads=4,
+            head_size=512,
+            is_causal=0,
+            has_nonpad_kv_seqlen=True,
+        )
+        nonpad_seqlens = torch.tensor([32, 19], dtype=torch.int64, device="cuda")
+
+        parity_check_gqa_prompt_with_nonpad_kv_seqlen(
+            config=nonpad_config,
+            nonpad_seqlens=nonpad_seqlens,
+            ep="CUDAExecutionProvider",
+            device="cuda",
+            torch_type=torch.float16,
+            ort_type=TensorProto.FLOAT16,
+            rtol=rtol["fp16"],
+            atol=atol["fp16"],
+        )
+
+
 class TestONNXAttentionGQALargeHeadUnfusedCPU(unittest.TestCase):
     """CPU twin of TestONNXAttentionGQALargeHeadUnfused.