[Performance]: updateCacheIndirectionKernel over-launches ~32× empty blocks due to `roundUp` vs `divUp` mismatch

Proposal to improve performance

Summary

In cpp/tensorrt_llm/kernels/beamSearchKernels.cu, the launch grid for updateCacheIndirectionKernel uses common::roundUp(nMaxSeqLen, 32) as the x-dimension block count. Since roundUp returns a value on the order of nMaxSeqLen (not nMaxSeqLen/32), this launches roughly 32× more blocks than necessary along the step dimension. ~31/32 of those blocks enter the kernel and immediately early-return, wasting SM scheduling and launch overhead on every beam-search decode step.

The kernel is correct — the extra blocks are guarded by step >= nMSL and write nothing — but the launch is unnecessarily expensive, especially since this runs once per decode step.

Location

cpp/tensorrt_llm/kernels/beamSearchKernels.cu, in invokeUpdateCacheIndirection:

https://github.com/NVIDIA/TensorRT-LLM/blob/3c68ae6ac79c48c6bad5816adcfcefc7d9897d55/cpp/tensorrt_llm/kernels/beamSearchKernels.cu#L131-L137

And the kernel's step indexing (updateCacheIndirectionKernel):

https://github.com/NVIDIA/TensorRT-LLM/blob/3c68ae6ac79c48c6bad5816adcfcefc7d9897d55/cpp/tensorrt_llm/kernels/beamSearchKernels.cu#L94-L112

Why this is a 32× over-launch

divUp and roundUp are defined in cpp/include/tensorrt_llm/common/cudaUtils.h:

https://github.com/NVIDIA/TensorRT-LLM/blob/3c68ae6ac79c48c6bad5816adcfcefc7d9897d55/cpp/include/tensorrt_llm/common/cudaUtils.h#L440-L451

So roundUp(nMSL, 32) == divUp(nMSL, 32) * 32.

With blockDim = 32:

• Intended block count (one block per 32-step tile): divUp(nMSL, 32) → nMSL/32 magnitude
• Current grid.x: roundUp(nMSL, 32) → nMSL magnitude

The current value is exactly 32× the intended one. Example with nMaxSeqLen = 2048:

	grid.x	blocks doing real work	empty blocks (early-return)
current (roundUp)	2048	64	1984 (96.9%)
fixed (divUp)	64	64	0

The step mapping step = blockIdx.x * blockDim.x + threadIdx.x is unchanged by the fix — the change only drops the 31/32 of blocks that currently enter and immediately return.

Suggested fix

- dim3 const grid(common::roundUp(bh.nMaxSeqLen, 32), bh.nBatchSize, bh.nBeamWidthOut);
+ dim3 const grid(common::divUp(bh.nMaxSeqLen, 32), bh.nBatchSize, bh.nBeamWidthOut);
  updateCacheIndirectionKernel<<<grid, 32, 0, stream>>>(tgtCI, srcCI, bh, maxAttentionWindow, sinkTokenLength);

Impact

• Correctness: unaffected — the extra blocks are already guarded by step >= nMSL.
• Performance: removes ~32× empty-block over-launch on the step dimension, which runs once per beam-search decode step. Small per-call saving, but it compounds over long generations.
• Risk: low. divUp(nMSL, 32) * 32 >= nMSL, so every valid step still gets a thread; no coverage is lost.

Report of performance regression

No response

Misc discussion on performance

No response

Your current environment (if you think it is necessary)

System Information:

• OS: Ubuntu
• Python version: 3.12
• CUDA version: 13.0
• GPU model(s):
• Driver version: 580.95.05
• TensorRT version:
• PyTorch version:
• TensorRT-LLM version: main branch

Detailed output:

Paste the output of the above commands here

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.