From 54f2dbc7ee522783112bea31291a134512858786 Mon Sep 17 00:00:00 2001
From: Anurag Mukkara <134339030+amukkara@users.noreply.github.com>
Date: Thu, 17 Sep 2026 18:21:27 +0000
Subject: [PATCH] Add larger tile size in fmha_v2

Signed-off-by: Anurag Mukkara <134339030+amukkara@users.noreply.github.com>
---
 cpp/kernels/fmha_v2/setup.py                                  | 4 +++-
 .../kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp     | 4 ++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/cpp/kernels/fmha_v2/setup.py b/cpp/kernels/fmha_v2/setup.py
index 82cd51ed345b..fddf77bf67ac 100644
--- a/cpp/kernels/fmha_v2/setup.py
+++ b/cpp/kernels/fmha_v2/setup.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2020-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2020-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -5610,6 +5610,8 @@ def enumerate_qmma_flash_kernels(specs,
         ((192, 128), (64, 64), 1),
         ((576, 512), (64, 64), 1),
     ]
+    if sm == 120:
+        params_q_kv_step.append((256, (64, 128), 1))
     input_layouts = [
         InputLayout.PACKED_QKV, InputLayout.CONTIGUOUS_Q_KV,
         InputLayout.Q_PAGED_KV, InputLayout.SEPARATE_Q_K_V
diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp
index 6aefa2a66bae..a9cdd99918db 100644
--- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp
+++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2020-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -405,7 +405,7 @@ void FusedMHARunnerV2::setupLaunchParams(MHARunnerParams runnerParams)
         mLaunchParams.kernel_s = 0;
         mLaunchParams.force_unroll = true;
         // enable tiled kernels on Ampere/Ada
-        if ((isSm89 || isSm120f) && mFixedParams.dataType == DATA_TYPE_E4M3)
+        if (isSm89 && mFixedParams.dataType == DATA_TYPE_E4M3)
         {
             // so far Ada QMMA only supports non-tiled kernels.
             mLaunchParams.granular_tiling = false;