From 886388cf9d2ce271fa2ec3e0ca261aeb45baad23 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 | 2 +- 2 files changed, 4 insertions(+), 2 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 72b4e6128331..d472e202b874 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp @@ -400,7 +400,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;