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)