[https://nvbugs/6270671][fix] Replace the hardcoded multiBlock=1 with a call to… by tensorrt-cicd · Pull Request #15312 · NVIDIA/TensorRT-LLM

tensorrt-cicd · 2026-06-12T17:10:08Z

Summary

Root cause: BF16 (HMMA) spec-dec XQA kernel hardcodes multiBlock=1; PR [None][feat] Multi-block mode for Hopper spec dec XQA kernel #4416's multi-block enablement only covers the FP8/GMMA path
Fix: Replace the hardcoded multiBlock=1 with a call to computeMultiBlockCountSpecDecGMMA, passing nbTokenBlocksPerGrp as the spec-block count to match the HMMA gridDim layout
Automated fix generated by repair-bot

Test plan

Verify fix on the same GPU type as the original failure
Check for regressions in related tests

Links

Bug: https://nvbugs/6270671

Summary by CodeRabbit

Performance
- Optimized speculative decoding execution in the inference kernel for improved text generation performance during speculative decoding operations.

coderabbitai · 2026-06-12T17:14:33Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 94e7faa0-bd82-40d9-9222-8a078d72551c

📥 Commits

Reviewing files that changed from the base of the PR and between b03b78f and fdb0753.

📒 Files selected for processing (1)

cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.cpp

📝 Walkthrough

Walkthrough

The HMMA speculative-decoding code path now computes multi-block grid dimensions using computeMultiBlockCountSpecDecGMMA with token-block count input, replacing previous commented-out logic when multi-block mode is enabled.

Changes

HMMA Speculative-Decoding Multi-Block Grid Tuning

Layer / File(s)	Summary
Multi-block grid computation update `cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.cpp`	HMMA speculative-decoding branch computes `multiBlock` via `computeMultiBlockCountSpecDecGMMA` with `nbTokenBlocksPerGrp` when `xqaParams.multi_block_mode` is enabled, affecting the grid dimensions for kernel launches.

🎯 1 (Trivial) | ⏱️ ~3 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly indicates this is a fix that replaces hardcoded multiBlock=1 with a function call for spec-dec XQA kernel computation, directly corresponding to the main change in the changeset.
Description check	✅ Passed	The PR description includes a clear summary, root cause, fix explanation, test plan with verification steps, and bug link, addressing the key sections of the template.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

jhaotingc · 2026-06-12T18:06:09Z

Hi @pengbowang-nv, is mha.cu spec-dec + multi-block mode supported? (accuracy validated?)

Also, one problem with computeMultiBlockCountSpecDecGMMA was, it took max_seq_len (history kv) as a tunable heuristic, but actual behavior is the variable is fixed during cuda graph capture (a longest ISL prompt was sent during capture), and no matter the actual ISL is during actual serving, it's always selecting the config captured by longest ISL.

This can and should be addressed for both GMMA and HMMA for perf.

Thanks! cc @mikeiovine

pengbowang-nv · 2026-06-15T06:44:00Z

Hi @jhaotingc , I think we haven't done such experiment before. I think the kernel should support it but I'll need to run some test to be sure.

Also, one problem with computeMultiBlockCountSpecDecGMMA was, it took max_seq_len (history kv) as a tunable heuristic, but actual behavior is the variable is fixed during cuda graph capture (a longest ISL prompt was sent during capture), and no matter the actual ISL is during actual serving, it's always selecting the config captured by longest ISL.

Regarding this, it is hard to do so because cuda graph would keep kernel launch parameter and execute no CPU code, which means it just "remembers" a config during capture. We may have some solutions, like having a cudagraph dedicated for short sequence length. Or maybe we can try change the kernel to decide multiblock num in kernel, so that the kernel would work as a persistent kernel, it may launch maximum number of CTAs but will only use some of them at runtime and let other early exit.

pengbowang-nv · 2026-06-15T08:15:18Z

I have verified in unit test that the kernel would work with multiblock spec-dec.

jhaotingc · 2026-06-15T21:31:15Z

Sounds good, thank you! @pengbowang-nv
@mikeiovine the kernel accuracy is verified, so this PR can be a temporary fix until a better heuristic is written.

…ernel PR NVIDIA#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>

Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com>

pengbowang-nv · 2026-06-16T06:22:40Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-16T06:28:32Z

PR_Github #54499 [ run ] triggered by Bot. Commit: 0474212 Link to invocation

tensorrt-cicd · 2026-06-16T11:03:31Z

PR_Github #54499 [ run ] completed with state SUCCESS. Commit: 0474212
/LLM/main/L0_MergeRequest_PR pipeline #43560 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

pengbowang-nv · 2026-06-16T12:02:41Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-16T12:09:53Z

PR_Github #54579 [ run ] triggered by Bot. Commit: 0474212 Link to invocation

tensorrt-cicd · 2026-06-16T15:28:29Z

PR_Github #54579 [ run ] completed with state SUCCESS. Commit: 0474212
/LLM/main/L0_MergeRequest_PR pipeline #43622 completed with status: 'SUCCESS'

CI Report

Link to invocation

pengbowang-nv · 2026-06-17T05:31:05Z

/bot run --stage-list "DGX_H200-8_GPUs-PyTorch-PerfSanity-Post-Merge-1"

tensorrt-cicd · 2026-06-17T05:37:07Z

PR_Github #54769 [ run ] triggered by Bot. Commit: 0474212 Link to invocation

tensorrt-cicd · 2026-06-17T06:55:05Z

PR_Github #54769 [ run ] completed with state FAILURE. Commit: 0474212
/LLM/main/L0_MergeRequest_PR pipeline #43788 (Partly Tested) completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

pengbowang-nv · 2026-06-17T07:08:56Z

Previous run is just a perf instability. Not blocking nor caused by this PR.

pengbowang-nv · 2026-06-17T07:09:02Z

/bot run

tensorrt-cicd · 2026-06-17T07:15:17Z

PR_Github #54785 [ run ] triggered by Bot. Commit: 0474212 Link to invocation

tensorrt-cicd · 2026-06-17T07:39:47Z

PR_Github #54785 [ run ] completed with state SUCCESS. Commit: 0474212
/LLM/main/L0_MergeRequest_PR pipeline #43802 completed with status: 'SUCCESS'

CI Report

Link to invocation

…ec-dec (NVIDIA#15312) Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com> Co-authored-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com> Signed-off-by: GitLab CI Bot <gitlab-ci@nvidia.com>

tensorrt-cicd assigned mikeiovine Jun 12, 2026

github-actions Bot assigned tensorrt-cicd Jun 12, 2026

pengbowang-nv reviewed Jun 15, 2026

View reviewed changes

Comment thread cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.cpp Outdated

tensorrt-cicd and others added 2 commits June 16, 2026 14:22

rename multiblock calc functions

0474212

Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com>

pengbowang-nv force-pushed the repair-bot-bug6270671 branch from 8b52322 to 0474212 Compare June 16, 2026 06:22

pengbowang-nv approved these changes Jun 17, 2026

View reviewed changes

pengbowang-nv merged commit 071c287 into NVIDIA:main Jun 17, 2026
7 checks passed

Uh oh!

Conversation

tensorrt-cicd commented Jun 12, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Links

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Jun 12, 2026

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

jhaotingc commented Jun 12, 2026

Uh oh!

pengbowang-nv commented Jun 15, 2026

Uh oh!

Uh oh!

pengbowang-nv commented Jun 15, 2026

Uh oh!

jhaotingc commented Jun 15, 2026

Uh oh!

pengbowang-nv commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

pengbowang-nv commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

pengbowang-nv commented Jun 17, 2026

Uh oh!

tensorrt-cicd commented Jun 17, 2026

Uh oh!

tensorrt-cicd commented Jun 17, 2026

Uh oh!

pengbowang-nv commented Jun 17, 2026

Uh oh!

pengbowang-nv commented Jun 17, 2026

Uh oh!

tensorrt-cicd commented Jun 17, 2026

Uh oh!

tensorrt-cicd commented Jun 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

tensorrt-cicd commented Jun 12, 2026 •

edited by coderabbitai Bot

Loading