From a910d8609d6fd5b6b936a9e90d7c1716df911e65 Mon Sep 17 00:00:00 2001 From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> Date: Fri, 12 Jun 2026 09:27:13 -0700 Subject: [PATCH 1/2] [nvbugs/6270671][fix] Enable multi-block mode for BF16 spec-dec XQA kernel PR #4416 enabled multi-block mode for the FP8 (GMMA) spec-dec XQA kernel on Hopper, but the same treatment was never applied to the BF16 (HMMA) spec-dec XQA kernel. On Hopper SM=90 with BF16 input + BF16 KV cache (e.g. GPT-OSS-120B + Eagle3), the GMMA path is rejected by supportConfigQGMMA (requires E4M3 KV or skip-softmax), so the launch falls back to the HMMA kernel with multiBlock hardcoded to 1. At ISL=32k, concurrency=1 the resulting ~4-CTA grid cannot saturate the SMs, so per-token latency dominates and Eagle3 acceptance ~1.8 fails to translate into an end-to-end speedup. Wire computeMultiBlockCountSpecDecGMMA into the HMMA spec-dec branch. The HMMA gridDim is {multiBlock, nbKVHeads * nbTokenBlocksPerGrp, batchSize}, so passing nbTokenBlocksPerGrp where the GMMA path passes specDecBlocks gives the helper the correct singleBlockCount and reuses the existing wave-count / history-aware tuning. The mha.cu kernel already implements the SPEC_DEC multi-block reduction (gridDim.x splits, per-(head, q-token) sub-sequence indexing, scratch + semaphore handling sized via gridDim.y), so no kernel-side change is needed. Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> --- .../decoderXQARunner.cpp | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.cpp b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.cpp index a443e618f241..7579e9662285 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.cpp +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.cpp @@ -462,10 +462,14 @@ void DecoderXQARunner::runImpl(XQAParams const& xqaParams, KVCacheBuffer const& appendParam(&launchParams.scratch); uint32_t multiBlock = 1; - // if (xqaParams.multi_block_mode) - // { - // multiBlock = computeMultiBlockCount(xqaParams, xqaParams.batch_size, multiprocessorCount); - // } + if (xqaParams.multi_block_mode) + { + // HMMA spec-dec gridDim is {multiBlock, nbKVHeads * nbTokenBlocksPerGrp, batchSize}; + // mirror the GMMA spec-dec wave-count tuning by passing nbTokenBlocksPerGrp where the GMMA + // path passes specDecBlocks (singleBlockCount factor along the head/q-token axis). + multiBlock = computeMultiBlockCountSpecDecGMMA( + xqaParams, xqaParams.batch_size, multiprocessorCount, nbTokenBlocksPerGrp); + } auto const gridDim = (dim3{multiBlock, xqaParams.num_kv_heads * nbTokenBlocksPerGrp, xqaParams.batch_size}); dim3 const blockDim(128, 1, 2); From 04742120402067091d95f77a931714bf3d1b67fa Mon Sep 17 00:00:00 2001 From: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com> Date: Tue, 16 Jun 2026 06:07:35 +0000 Subject: [PATCH 2/2] rename multiblock calc functions Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com> --- .../decoderMaskedMultiheadAttention/decoderXQARunner.cpp | 4 ++-- .../decoderMaskedMultiheadAttention/decoderXQARunnerUtils.h | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.cpp b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.cpp index 7579e9662285..e49dcf849ed1 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.cpp +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.cpp @@ -467,7 +467,7 @@ void DecoderXQARunner::runImpl(XQAParams const& xqaParams, KVCacheBuffer const& // HMMA spec-dec gridDim is {multiBlock, nbKVHeads * nbTokenBlocksPerGrp, batchSize}; // mirror the GMMA spec-dec wave-count tuning by passing nbTokenBlocksPerGrp where the GMMA // path passes specDecBlocks (singleBlockCount factor along the head/q-token axis). - multiBlock = computeMultiBlockCountSpecDecGMMA( + multiBlock = computeMultiBlockCountSpecDec( xqaParams, xqaParams.batch_size, multiprocessorCount, nbTokenBlocksPerGrp); } auto const gridDim = (dim3{multiBlock, xqaParams.num_kv_heads * nbTokenBlocksPerGrp, xqaParams.batch_size}); @@ -540,7 +540,7 @@ void DecoderXQARunner::runImpl(XQAParams const& xqaParams, KVCacheBuffer const& { if (isSpecDec && isGMMAKernel) { - multiBlock = computeMultiBlockCountSpecDecGMMA( + multiBlock = computeMultiBlockCountSpecDec( xqaParams, xqaParams.batch_size, multiprocessorCount, specDecBlocks); } else if (!isSpecDec) diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunnerUtils.h b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunnerUtils.h index 0038a47e668b..b57abf025583 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunnerUtils.h +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunnerUtils.h @@ -442,7 +442,7 @@ inline int computeMultiBlockCountForMLA(XQAParams const& xqaParams, int multipro return 1; // disable multi-block for MLA kernel for now. } -inline int computeMultiBlockCountSpecDecGMMA( +inline int computeMultiBlockCountSpecDec( XQAParams const& xqaParams, int batchSize, int multiprocessorCount, int specDecBlocks) { auto const userSpecified = tensorrt_llm::common::getEnvXqaBlocksPerSequence();