[https://nvbugs/6242591][fix] Fix bugs in Beam Search kernels (NVIDIA#15621)

Signed-off-by: wili <12345678+wili@users.noreply.github.com>

Author: wili-65535

cpp/tensorrt_llm/kernels/beamSearchKernels.cu

@@ -16,6 +16,7 @@
 
 #include "tensorrt_llm/common/config.h"
 #include "tensorrt_llm/common/cudaUtils.h"
+#include "tensorrt_llm/common/reduceKernelUtils.cuh"
 #include "tensorrt_llm/kernels/beamSearchKernels.h"
 
 using namespace tensorrt_llm::common;
@@ -135,21 +136,26 @@ void invokeUpdateCacheIndirection(int* tgtCI, int const* srcCI, BeamHypotheses&
     sync_check_cuda_error(stream);
 }
 
-__global__ void addCumLogProbs(float* __restrict pStage1LogProbs, float const* __restrict cumLogProbs,
-    FinishedState const* finished, int const* endIds, float const* diversityRates,
+__global__ void addCumLogProbs(float* __restrict pStage1LogProbs, int const* __restrict pStage1Ids,
+    float const* __restrict cumLogProbs, FinishedState const* finished, int const* endIds, float const* diversityRates,
     runtime::SizeType32 const* batchSlots, size_t const nBS, size_t const nBMIn, size_t const nBMOut, size_t const nBM)
 {
     int const bid = blockIdx.x; // Index of request in batch
     runtime::SizeType32 const slot = batchSlots[bid];
     float const diversityRate{diversityRates[slot]};
     float* pLocalLogProbs = pStage1LogProbs + bid * nBMIn * nBMOut * 2;
+    int const* pLocalIds = pStage1Ids + bid * nBMIn * nBMOut * 2;
 
     for (int i = threadIdx.x; i < nBMIn * nBMOut * 2; i += blockDim.x)
     {
         int const iBMIn = i / (nBMOut * 2);
-        if (finished[slot * nBMIn + iBMIn].isFinished())
+        if (finished[slot * nBM + iBMIn].isFinished())
         {
-            pLocalLogProbs[i] += (i == endIds[slot]) ? 1.0f : 0.0f;
+            // In V2 path, i is a candidate-slot index (0..nBMIn*nBMOut*2-1), NOT a vocab token id.
+            // Use pStage1Ids to look up the actual token id for the EOS comparison.
+            bool const isEOS = (pLocalIds[i] == endIds[slot]);
+            // Keep only the EOS candidate with its proper cumulative score; suppress all others.
+            pLocalLogProbs[i] = isEOS ? (pLocalLogProbs[i] + cumLogProbs[slot * nBM + iBMIn]) : -FLT_MAX;
         }
         else
         {
@@ -160,21 +166,27 @@ __global__ void addCumLogProbs(float* __restrict pStage1LogProbs, float const* _
     return;
 }
 
-__global__ void addCumLogProbs(half* __restrict pStage1LogProbs, float const* __restrict cumLogProbs,
-    FinishedState const* finished, int const* endIds, float const* diversityRates,
+__global__ void addCumLogProbs(half* __restrict pStage1LogProbs, int const* __restrict pStage1Ids,
+    float const* __restrict cumLogProbs, FinishedState const* finished, int const* endIds, float const* diversityRates,
     runtime::SizeType32 const* batchSlots, size_t const nBS, size_t const nBMIn, size_t const nBMOut, size_t const nBM)
 {
     int const bid = blockIdx.x; // Index of request in batch
     runtime::SizeType32 const slot = batchSlots[bid];
     float const diversityRate{diversityRates[slot]};
     half* pLocalLogProbs = pStage1LogProbs + bid * nBMIn * nBMOut * 2;
+    int const* pLocalIds = pStage1Ids + bid * nBMIn * nBMOut * 2;
 
     for (int i = threadIdx.x; i < nBMIn * nBMOut * 2; i += blockDim.x)
     {
         int const iBMIn = i / (nBMOut * 2);
-        if (finished[slot * nBMIn + iBMIn].isFinished())
+        if (finished[slot * nBM + iBMIn].isFinished())
         {
-            pLocalLogProbs[i] += (i == endIds[slot]) ? 1.0f : 0.0f;
+            // In V2 path, i is a candidate-slot index (0..nBMIn*nBMOut*2-1), NOT a vocab token id.
+            // Use pStage1Ids to look up the actual token id for the EOS comparison.
+            bool const isEOS = (pLocalIds[i] == endIds[slot]);
+            // Keep only the EOS candidate with its proper cumulative score; suppress all others.
+            pLocalLogProbs[i]
+                = isEOS ? (half) (float(pLocalLogProbs[i]) + cumLogProbs[slot * nBM + iBMIn]) : (half) -HALF_FLT_MAX;
         }
         else
         {

cpp/tensorrt_llm/kernels/beamSearchKernels.h

@@ -131,13 +131,15 @@ void invokeTopkBeamSearch(T const* logProbs, T const* bias, void* workspace, Bea
 void invokeUpdateCacheIndirection(int* tgtCI, int const* srcCI, BeamHypotheses& bh,
     runtime::SizeType32 const maxAttentionWindow, runtime::SizeType32 sinkTokenLength, cudaStream_t stream);
 
-__global__ void addCumLogProbs(float* __restrict pStage1LogProbs, float const* __restrict cumLogProbs,
-    ::tensorrt_llm::kernels::FinishedState const* finished, int const* endIds, float const* diversityRates,
-    runtime::SizeType32 const* batchSlots, size_t const nBS, size_t const nBMIn, size_t const nBMOut, size_t const nBM);
-
-__global__ void addCumLogProbs(half* __restrict pStage1LogProbs, float const* __restrict cumLogProbs,
-    ::tensorrt_llm::kernels::FinishedState const* finished, int const* endIds, float const* diversityRates,
-    runtime::SizeType32 const* batchSlots, size_t const nBS, size_t const nBMIn, size_t const nBMOut, size_t const nBM);
+__global__ void addCumLogProbs(float* __restrict pStage1LogProbs, int const* __restrict pStage1Ids,
+    float const* __restrict cumLogProbs, ::tensorrt_llm::kernels::FinishedState const* finished, int const* endIds,
+    float const* diversityRates, runtime::SizeType32 const* batchSlots, size_t const nBS, size_t const nBMIn,
+    size_t const nBMOut, size_t const nBM);
+
+__global__ void addCumLogProbs(half* __restrict pStage1LogProbs, int const* __restrict pStage1Ids,
+    float const* __restrict cumLogProbs, ::tensorrt_llm::kernels::FinishedState const* finished, int const* endIds,
+    float const* diversityRates, runtime::SizeType32 const* batchSlots, size_t const nBS, size_t const nBMIn,
+    size_t const nBMOut, size_t const nBM);
 
 __global__ void gatherId(int const* __restrict pStage1Id, int* __restrict pStage2Id, size_t const nBS,
     size_t const nBMIn, size_t const nBMOut, size_t const nV);

cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernelsTemplate.h

@@ -208,6 +208,7 @@ __launch_bounds__(BLOCK_SIZE) __global__ void beamStage3Kernel(
     __shared__ float smemCumLogProbs[PBM];
     __shared__ int smemSeqLen[PBM];
     __shared__ KVPair smemTopKV[(IS_V2) ? 1 : PBM * 2]; // Just a placeholder in V2 workflow
+    __shared__ int smemNBeamForNextStep;
 
     if (bh.numBeamsCBA != nullptr)
     {
@@ -217,14 +218,11 @@ __launch_bounds__(BLOCK_SIZE) __global__ void beamStage3Kernel(
             // Initialize worst score in the first call
             bh.minNormedScoresCBA[slot] = 0.0f; // logProbs is in range (-inf, 0]
         }
-        else if (earlyStopping == 1 && bh.numBeamsCBA[slot] == nBM
-            || earlyStopping != 1 && bh.finished[slot * nBM].isFinished())
+        else if (earlyStopping == 1 && bh.numBeamsCBA[slot] >= nBM || earlyStopping != 1 && bh.batchDones[slot])
         {
             // Condition of early return:
             // 1. In EarlyStopping mode, and we have got enough beams
             // 2. In NonEarlyStopping mode, and this batch has been marked as done
-            // TODO: improve the condition like below
-            // earlyStopping == 1 && bh.numBeamsCBA[slot] == nBM || earlyStopping != 1 && bh.batchDones[slot]
             return;
         }
     }
@@ -296,6 +294,11 @@ __launch_bounds__(BLOCK_SIZE) __global__ void beamStage3Kernel(
     }
     __syncthreads();
 
+    // Timestep at which each next-step destination beam's token is stored in the work tree.
+    // This is the parent beam's sequence length (the true generation step), which may differ
+    // from the destination slot's own (possibly stale) length when a finished slot is reused.
+    __shared__ int smemWriteStep[PBM];
+
     if (tid == 0)
     {
         int nBeamForNextStep{0};
@@ -324,10 +327,14 @@ __launch_bounds__(BLOCK_SIZE) __global__ void beamStage3Kernel(
             {
                 // Condition of this branch:
                 // This token is end-token and belongs to top nBM range in Beam search mode
-                int const nSeqLen = bh.sequenceLengths[slot * nBM + i] + 1 - bh.inputLengths[slot * nBM + i];
+                // Use the actual parent beam index (topId / nV) % nBM, not the candidate rank i,
+                // to look up the correct sequenceLength and inputLength for length-penalty scoring.
+                int const parentBeam = (topId / nV) % nBM;
+                int const nSeqLen
+                    = bh.sequenceLengths[slot * nBM + parentBeam] + 1 - bh.inputLengths[slot * nBM + parentBeam];
                 float const score = applyLengthPenalty(topLogProb, nSeqLen, lengthPenalty);
                 int nCBA = bh.numBeamsCBA[slot];
-                if (nCBA == nBM)
+                if (nCBA >= nBM)
                 {
                     // There are already nBM beams
                     if (score < bh.minNormedScoresCBA[slot])
@@ -407,13 +414,18 @@ __launch_bounds__(BLOCK_SIZE) __global__ void beamStage3Kernel(
                 // 1. bh.numBeamsCBA == nullptr && i <  nBM, i.e., beam search is disable
                 // 2. bh.numBeamsCBA != nullptr && i <  nBM && isEndToken == false, i.e., add token at the end
                 // 3. bh.numBeamsCBA != nullptr && i >= nBM && isEndToken == false, i.e., add token at the end
-                int const step = bh.sequenceLengths[slot * nBM + nBeamForNextStep];
+                // Write at the parent beam's sequence length (the actual generation step),
+                // not the destination slot's length, which can be stale if the slot was
+                // previously finished and is now being reused for a new continuation.
+                int const parentBeam = topId / nV % nBM;
+                int const step = bh.sequenceLengths[slot * nBM + parentBeam];
+                smemWriteStep[nBeamForNextStep] = step;
                 // Copy the selected token to work tree
                 bh.outputIdsPtr[slot][nBeamForNextStep * nMSL + step] = topId;
                 if (bh.logProbsTiled != nullptr)
                 {
                     int const index = step * nMBS * nBM + slot * nBM + nBeamForNextStep;
-                    int const indexBeam = topId / nV % nBM;
+                    int const indexBeam = parentBeam;
                     bh.logProbsTiled[index] = (float) topLogProb - smemCumLogProbs[indexBeam];
                 }
                 bh.cumLogProbs[slot * nBM + nBeamForNextStep] = (float) topLogProb;
@@ -437,6 +449,7 @@ __launch_bounds__(BLOCK_SIZE) __global__ void beamStage3Kernel(
                 break;
             }
         }
+        smemNBeamForNextStep = nBeamForNextStep;
     }
 
     // Update bh.batchDones
@@ -481,23 +494,40 @@ __launch_bounds__(BLOCK_SIZE) __global__ void beamStage3Kernel(
     if (tid < nBMOut)
     {
         int const indexBatchBeam = slot * nBM + tid;
-        int const step = smemSeqLen[tid];
-        if (!bh.finished[indexBatchBeam].isFinished())
+        if (tid < smemNBeamForNextStep)
         {
-            smemSeqLen[tid]++;
+            // This slot received a valid next-step token from the selection phase.
+            // Use the timestep recorded by the selection phase (the parent beam's length),
+            // which matches the position where the encoded token was stored.
+            int const step = smemWriteStep[tid];
+            int const newId = bh.outputIdsPtr[slot][tid * nMSL + step];
+            int const newBeamId = (newId / nV) % nBM;
+            int const newTokenId = newId % nV;
+            int const indexParentBeam = slot * nBM + newBeamId;
+            int const parentSeqLen = smemSeqLen[newBeamId];
+            bh.sequenceLengths[indexBatchBeam] = parentSeqLen + (!bh.finished[indexParentBeam].isFinished() ? 1 : 0);
+            if (newTokenId == bh.endIds[slot])
+            {
+                bh.finished[indexBatchBeam].setFinishedEOS();
+            }
+            else
+            {
+                // Reset any stale finished state: this slot may have been marked finished in a
+                // previous step and is now reused for a valid non-EOS beam; otherwise it would be
+                // wrongly skipped downstream.
+                bh.finished[indexBatchBeam] = FinishedState::empty();
+            }
+            bh.parentIdsPtr[slot][tid * nMSL + step] = newBeamId;
+            bh.outputIdsPtr[slot][tid * nMSL + step] = newTokenId;
         }
-        int const newId = bh.outputIdsPtr[slot][tid * nMSL + step];
-        int const newBeamId = (newId / nV) % nBM;
-        int const newTokenId = newId % nV;
-        bh.sequenceLengths[indexBatchBeam] = smemSeqLen[newBeamId];
-        if (newTokenId == bh.endIds[slot])
+        else
         {
-            bh.finished[indexBatchBeam].setFinishedEOS();
+            // No valid next-step token for this slot: all top candidates went to CBA.
+            // Mark as finished so downstream stages (cache indirection, next decode) skip it.
+            bh.finished[indexBatchBeam].setFinished();
         }
-        bh.parentIdsPtr[slot][tid * nMSL + step] = newBeamId;
-        bh.outputIdsPtr[slot][tid * nMSL + step] = newTokenId;
 
-        if ((earlyStopping == 1) && (bh.numBeamsCBA != nullptr && bh.numBeamsCBA[slot] == nBM)
+        if ((earlyStopping == 1) && (bh.numBeamsCBA != nullptr && bh.numBeamsCBA[slot] >= nBM)
             || (earlyStopping != 1) && bh.batchDones[slot])
         {
             bh.batchDones[slot] = true;
@@ -631,7 +661,7 @@ void beamSearchKernelLauncher(
         sync_check_cuda_error(stream);
 
         int nThread = min(roundUp(nBMIn * nBMOut * 2, 32), 1024);
-        addCumLogProbs<<<nBS, nThread, 0, stream>>>(pStage1LogProbs, bh.cumLogProbs, bh.finished, bh.endIds,
+        addCumLogProbs<<<nBS, nThread, 0, stream>>>(pStage1LogProbs, pStage1Ids, bh.cumLogProbs, bh.finished, bh.endIds,
             bh.diversityRates, bh.batchSlots, nBS, nBMIn, nBMOut, nBM);
         sync_check_cuda_error(stream);
 

cpp/tensorrt_llm/kernels/decodingKernels.cu

@@ -122,7 +122,7 @@ __global__ void gatherTree(gatherTreeParam param)
         {
             int const levelBeamIx = batch * param.beamWidth * param.maxSeqLen + beam * param.maxSeqLen + level;
             int const levelParentIx = batch * param.beamWidth * param.maxSeqLen + parent * param.maxSeqLen + level;
-            if (parent < 0 || parent > param.beamWidth)
+            if (parent < 0 || parent >= param.beamWidth)
             {
                 param.outputIds[levelBeamIx] = param.endTokens[batch];
                 parent = -1;
@@ -702,10 +702,11 @@ __global__ void transposeLogProbs(float* outputLogProbs, float* outputLogProbsTi
     }
 
     auto const batchSlot = batchSlots[batchIdx];
-    if (pos < sequenceLengths[batchSlot])
+    auto const batchBeamIdx = batchSlot * beamWidth + beamIdx;
+    if (pos < sequenceLengths[batchBeamIdx])
     {
-        auto const batchBeamIdx = batchSlot * beamWidth * maxSeqLen + beamIdx * maxSeqLen + pos;
-        outputLogProbs[batchBeamIdx]
+        auto const outputIndex = batchSlot * beamWidth * maxSeqLen + beamIdx * maxSeqLen + pos;
+        outputLogProbs[outputIndex]
             = outputLogProbsTiled[pos * maxBatchSize * beamWidth + batchSlot * beamWidth + beamIdx];
     }
 }

tests/unittest/_torch/sampler/test_beam_search.py

@@ -23,8 +23,9 @@
 
 import pytest
 import torch
-from test_beam_search_util import (BeamSearchTestOutput, DummyConfigLoader,
-                                   DummyWeightLoader, get_expected_outputs)
+from test_beam_search_util import (BeamSearchTestOutput, DummyConfig,
+                                   DummyConfigLoader, DummyWeightLoader,
+                                   get_expected_outputs)
 from utils.llm_data import llm_models_root
 from utils.util import assert_no_cuda_sync, force_ampere, run_test_with_warmup
 
@@ -543,6 +544,76 @@ def test_beam_search_disagg_e2e(
         gen_llm.shutdown()
 
 
+@pytest.mark.parametrize("beam_width", [10])
+@pytest.mark.threadleak(enabled=False)
+def test_beam_search_large_beam_width_regression(
+    beam_width: int,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """
+    Fix https://nvbugs/6242591
+    """
+    num_prompts = 2
+    input_prompts = [[1, 2, 3], [4, 5, 6]]
+    vocab_size = DummyConfig().vocab_size
+    # The dummy model pushes each beam towards (token + 3) mod vocab_size, so an
+    # end id a few multiples of 3 ahead of the prompt is reachable and different
+    # beams hit it at different steps -> exercises the finished-slot reuse path.
+    end_id = 24
+
+    checkpoint_loader = HfCheckpointLoader(
+        weight_loader=DummyWeightLoader(),
+        config_loader=DummyConfigLoader(),
+    )
+
+    gc.collect(2)  # force destruction of any other LLM instances
+    with _single_process_context():
+        llm = LLM(
+            model=_pl.Path("dummy_path"),
+            checkpoint_loader=checkpoint_loader,
+            sampler_type="TRTLLMSampler",
+            max_beam_width=beam_width,
+            max_batch_size=beam_width * num_prompts,
+            max_seq_len=64,
+            kv_cache_config=KvCacheConfig(max_tokens=10000),
+            disable_overlap_scheduler=True,
+            cuda_graph_config=None,
+        )
+        with llm:
+            sampling_params = SamplingParams(
+                max_tokens=16,
+                n=beam_width,
+                best_of=beam_width,
+                use_beam_search=True,
+                beam_search_diversity_rate=0.5,
+                early_stopping=1,
+                length_penalty=0.5,
+                end_id=end_id,
+            )
+            outputs = llm.generate(deepcopy(input_prompts),
+                                   sampling_params=deepcopy(sampling_params))
+
+    assert isinstance(outputs, list)
+    assert len(outputs) == num_prompts
+    for output in outputs:
+        beams = output.outputs
+        assert len(beams) == beam_width, (
+            f"expected {beam_width} beams, but got {len(beams)}")
+        beam_sequences = []
+        for beam_idx, beam in enumerate(beams):
+            token_ids = beam.token_ids
+            assert token_ids is not None, f"beam {beam_idx} has no token_ids"
+            assert len(token_ids) > 0, f"beam {beam_idx} is empty"
+            assert all(0 <= t < vocab_size for t in token_ids), (
+                f"beam {beam_idx} has out-of-vocab tokens: {token_ids}")
+            beam_sequences.append(tuple(token_ids))
+        # Beams must not all collapse to the same sequence: corrupted state
+        # produced duplicated / garbled beams. With a non-zero diversity rate
+        # the beams are expected to differ.
+        assert len(set(beam_sequences)) > 1, (
+            f"all {beam_width} beams are identical: {beam_sequences[0]}")
+
+
 ###########################################################################
 # Unit tests
 ###########################################################################