From 501d7cd636fbc61ba61536344ece9973c67c131b Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Mon, 6 Jul 2026 09:57:34 +0000 Subject: [PATCH] fix(fmha): refresh JIT libraries Refresh x86_64 and aarch64 FMHA JIT libraries for skip-softmax generation and sm_103a NVRTC linking. Synchronize FmhaAutoTuner, FmhaOptions, and KernelTraits with the matching FMHA snapshot. Retain all other exported headers and cubins. Remove the related test waivers. Refs NVBug 6388157 Refs NVBug 6390307 Refs TRTLLM-12988 Refs TRTLLM-12990 Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- .../fmha/lib/aarch64-linux-gnu/libTrtLlmGen.a | 4 ++-- .../lib/aarch64-linux-gnu/libTrtLlmGenFmhaLib.a | 4 ++-- .../fmha/lib/x86_64-linux-gnu/libTrtLlmGen.a | 4 ++-- .../lib/x86_64-linux-gnu/libTrtLlmGenFmhaLib.a | 4 ++-- .../fmha/trtllmGen_fmha_export/FmhaAutoTuner.h | 3 +++ .../fmha/trtllmGen_fmha_export/FmhaOptions.h | 10 +++++++--- .../fmha/trtllmGen_fmha_export/KernelTraits.h | 14 ++++++++++---- tests/integration/test_lists/waives.txt | 4 ---- 8 files changed, 28 insertions(+), 19 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/lib/aarch64-linux-gnu/libTrtLlmGen.a b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/lib/aarch64-linux-gnu/libTrtLlmGen.a index d6e6baffa77e..b00b8dabaa81 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/lib/aarch64-linux-gnu/libTrtLlmGen.a +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/lib/aarch64-linux-gnu/libTrtLlmGen.a @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:2b0ff40af56e9dc2544827244101f22f03abf7942103f76ec628b47b53a52f90 -size 7293250 +oid sha256:7057231cff1dc01ca4249ca8da35869df4fc6403d1cd629f657c3dcf08c6440d +size 7299448 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/lib/aarch64-linux-gnu/libTrtLlmGenFmhaLib.a b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/lib/aarch64-linux-gnu/libTrtLlmGenFmhaLib.a index 21abec0014ed..f2bcb9d72da7 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/lib/aarch64-linux-gnu/libTrtLlmGenFmhaLib.a +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/lib/aarch64-linux-gnu/libTrtLlmGenFmhaLib.a @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:eb90c34477c6b9564417dcadcd05a582715fd6d5e86e222db10c512c5ec1559c -size 5074186 +oid sha256:32243e9163f943f6c8908a55cfe287255e5268a545d04b794b764e374f15ea0d +size 5075466 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/lib/x86_64-linux-gnu/libTrtLlmGen.a b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/lib/x86_64-linux-gnu/libTrtLlmGen.a index 9b421ae07d3f..0f82d902a00b 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/lib/x86_64-linux-gnu/libTrtLlmGen.a +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/lib/x86_64-linux-gnu/libTrtLlmGen.a @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:814ebc3b8b3abc1836879943b77c43d8baff0804dd784ea1792bd05910600469 -size 7345482 +oid sha256:37db982312fe782dd0ad8282e695b6c8a83428d306dbe2f2c94b6a578f5bc7f6 +size 7352462 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/lib/x86_64-linux-gnu/libTrtLlmGenFmhaLib.a b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/lib/x86_64-linux-gnu/libTrtLlmGenFmhaLib.a index 199dfbb235d7..bf41da145bdd 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/lib/x86_64-linux-gnu/libTrtLlmGenFmhaLib.a +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/lib/x86_64-linux-gnu/libTrtLlmGenFmhaLib.a @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:69e1110204c5222f97d165a706fc1dfd6dbba55a3a49ae85e2a92b4367cfe498 -size 5145196 +oid sha256:50cb4f541392084afda1d4c809ff387134b0cedebb35a9bcea683926b42afee0 +size 5146356 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/FmhaAutoTuner.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/FmhaAutoTuner.h index 76e1a19efc59..00d1270ffdf4 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/FmhaAutoTuner.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/FmhaAutoTuner.h @@ -66,6 +66,9 @@ class FmhaAutoTuner { // Select the MLA generation kernel. void selectMlaGenerationKernel(); + // Select the grouped Q64 MLA generation kernel when the capability predicate matches. + void selectGroupedMlaGenerationKernel(); + private: // computeNumCtas uses the same cluster occupancy helper as enableCgaReduction. friend std::tuple computeNumCtas(FmhaOptions& options, diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/FmhaOptions.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/FmhaOptions.h index 4e7557c3faf0..33f3a7304753 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/FmhaOptions.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/FmhaOptions.h @@ -50,6 +50,8 @@ struct FmhaOptions : public KernelConfigBase { int32_t mDsv4ScaleBufM{0}; // Enable the auto tuner. bool mEnablesAutoTuner{false}; + // Select the grouped MLA generation kernel in the auto tuner. + bool mSelectsGroupedMla{false}; // Enable the BF16Q+FP8KV K-only transform path. Disabled by default. bool mEnablesBf16QFp8KvKOnlyTransform{false}; // Whether is exporting cubin. @@ -133,6 +135,7 @@ struct FmhaOptions : public KernelConfigBase { TO_JSON(mDryRun); TO_JSON(mDsv4ScaleBufM); TO_JSON(mEnablesAutoTuner); + TO_JSON(mSelectsGroupedMla); TO_JSON(mEnablesBf16QFp8KvKOnlyTransform); TO_JSON(mIsExportingCubin); TO_JSON(mIsTracing); @@ -529,12 +532,13 @@ inline void checkFmhaOptions(FmhaOptions const& options, } } - // The mGroupsTokensHeadsQ only works with GQA generation kernels. + // groupsTokensHeadsQ is allowed for GQA gen and explicitly selected MLA gen kernels. if (options.mGroupsTokensHeadsQ) { TLLM_CHECK_ERROR(!isContextKernel(options.mFmhaKernelType), "mGroupsTokensHeadsQ should only be enabled for generation kernels."); - TLLM_CHECK_ERROR(!options.mIsMlaGen, - "MLA gen kernels haven't supported mGroupsTokensHeadsQ yet."); + TLLM_CHECK_ERROR(!options.mIsMlaGen + || options.mSelectsGroupedMla, + "MLA generation with mGroupsTokensHeadsQ requires mSelectsGroupedMla."); } diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/KernelTraits.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/KernelTraits.h index 3ea36fc90f96..4a9c801a7230 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/KernelTraits.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/KernelTraits.h @@ -207,8 +207,14 @@ struct KernelConfig : public KernelConfigBase { // Set numStagesQ for headDim > 128 kernels. if (mNumInstsQ * mNumInstsKv == 1) { - TLLM_CHECK_INFO(mTileSizeQ == 64 || (mHeadDimQk > 128 && mHeadDimV > 128), - "Consider using numInstsQ = 2 for better performance."); + // Whether the kernel is a generation kernel that skips softmax when possible. + bool isGenerationSkipsSoftmax = + options.mSkipsSoftmaxWhenPossible && !isContextKernel(options.mFmhaKernelType); + // Skip the check for skipsSoftmax generation kernels because this is intended. + if (!isGenerationSkipsSoftmax) { + TLLM_CHECK_INFO(mTileSizeQ == 64 || (mHeadDimQk > 128 && mHeadDimV > 128), + "Consider using numInstsQ = 2 for better performance."); + } // There is no enough shared memory for 2 stages when the headDim is not split into multiple // stages. if (mHeadDimPerStageKv == 0 && keepsMmaAbForDsMlaGen) { @@ -748,8 +754,8 @@ struct KernelTraits : public KernelConfig, public MmaTraits { mSeparateSmemKv = true; } - // The tile size for the correction step. - mCorrTileSize = std::min(mValidTilePvN, 64); + // Use the largest power-of-two tile up to 64 that exactly divides the correction width. + mCorrTileSize = std::gcd(mValidTilePvN, 64); // The number of keys per tile. mNumKeysPerTile = std::min(mNumTokensPerPage, mTileSizeKv); diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 709b74932fcb..46fd4aeb615d 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -111,7 +111,6 @@ accuracy/test_llm_api_pytorch.py::TestQwen3NextInstruct::test_bf16_4gpu[tep4] SK accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_fp8[latency] SKIP (https://nvbugs/6177390) accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_fp8[throughput_latency] SKIP (https://nvbugs/6177390) accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4[latency_moe_trtllm] SKIP (https://nvbugs/6402009) -accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention_4gpus[target_sparsity_0.9-fp8kv=False] SKIP (https://nvbugs/6388157) accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp1-TRTLLM] SKIP (https://nvbugs/6428008) accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp SKIP (https://nvbugs/6206179) accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8[enable_block_reuse=True] SKIP (https://nvbugs/6210714) @@ -223,9 +222,6 @@ full:B300/accuracy/test_llm_api_pytorch_multimodal.py::TestQwen2_5_VL_7B::test_a full:B300/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_genpp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6322073) full:B300/test_e2e.py::test_qwen_e2e_cpprunner_large_new_tokens[DeepSeek-R1-Distill-Qwen-1.5B-DeepSeek-R1-Distill-Qwen-1.5B] SKIP (https://nvbugs/6414760) full:B300/unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" SKIP (https://nvbugs/6165866) -full:DGX_B200/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention_4gpus[target_sparsity_0.5-fp8kv=False] SKIP (https://nvbugs/6390307) -full:DGX_B200/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention_4gpus[target_sparsity_0.5-fp8kv=True] SKIP (https://nvbugs/6390307) -full:DGX_B200/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention_4gpus[target_sparsity_0.9-fp8kv=True] SKIP (https://nvbugs/6390307) full:DGX_B200/unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" SKIP (https://nvbugs/6165866) full:DGX_B200/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py::test_allreduce_strategies[MIN_LATENCY] SKIP (https://nvbugs/6403920) full:DGX_B200/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py::test_allreduce_strategies[ONESHOT] SKIP (https://nvbugs/6403920)