About diverse beam search implementation

[shenlanyilang]

The diverse beam search in tensorrtllm seems so different with huggingface and the paper. Why using the parent beam's index as a additional score, is there any paper/work suggest that?
The huggingface's implementation will consider other beam's word, if having overlaps, will got a penalty score, which can avoid generating duplicates words.

# code in cpp/tensorrt_llm/kernels/beamSearchKernels.cu
template <typename T>
__global__ void addCumLogProbs(T* __restrict pStage1LogProbs, 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]};
    T* pLocalLogProbs = pStage1LogProbs + 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())
        {
            pLocalLogProbs[i] += (i == endIds[slot]) ? 1.0f : 0.0f;
        }
        else
        {
            // nBM is used in VBWS since `cumLogProbs` is initialized with kMaxBeamWidth earlier than BeamSearchLayer
            pLocalLogProbs[i] += cumLogProbs[slot * nBM + iBMIn] + diversityRate * iBMIn;
        }
    }
    return;
}

[Advait251206]

The implementation in TensorRT-LLM is indeed different from the original Diverse Beam Search (Vijayakumar et al., 2016) and from Hugging Face's implementation.

The key line is:

pLocalLogProbs[i] += cumLogProbs[slot * nBM + iBMIn] + diversityRate * iBMIn;

Here, the diversity term depends only on the parent beam index (iBMIn), not on the candidate token. This is fundamentally different from the DBS paper.

In the original Diverse Beam Search, diversity is introduced by penalizing candidate tokens that have already been selected by previous beam groups at the current decoding step. The diversity penalty is therefore token-dependent, encouraging different groups to choose different next tokens.

Similarly, Hugging Face's implementation applies a Hamming diversity penalty, where if another beam group has already selected the same token, that token receives an additional penalty. This directly discourages duplicate token choices across beam groups.

By contrast, TensorRT-LLM simply adds a constant offset based on the parent beam index. Since this value is constant for every candidate token originating from the same parent beam, it changes the relative preference between parent beams but does not discourage different beams from selecting the same token. In other words, it is not implementing the diversity objective described in the DBS paper.

To see why, suppose we have two parent beams, beam0 and beam1, and both assign the following log probabilities to the next tokens:

Parent Beam	Token A	Token B
beam0	-0.2	-0.5
beam1	-0.2	-0.5

With diversityRate = 0.5, TensorRT-LLM computes:

beam0:
A = -0.2 + 0.5 * 0 = -0.2
B = -0.5 + 0.5 * 0 = -0.5

beam1:
A = -0.2 + 0.5 * 1 = 0.3
B = -0.5 + 0.5 * 1 = 0.0

Notice that both tokens from beam1 receive exactly the same bonus. The ranking within a parent beam remains unchanged:

A > B   (before)
A > B   (after)

No token receives an additional penalty because another beam already selected it. Therefore, both beams may still choose exactly the same token.

In contrast, the DBS paper and Hugging Face's implementation would apply the diversity penalty to the token itself. If Token A had already been selected by an earlier beam group, only Token A would be penalized, making Token B more competitive and encouraging different beams to generate different continuations.

For this reason, TensorRT-LLM's implementation appears to be a beam-level bias rather than a token-level diversity penalty.

I am also not aware of any published paper that proposes using the parent beam index itself as a diversity score in this way. It may simply be a lightweight heuristic or an implementation shortcut to bias beam selection while avoiding the additional synchronization and bookkeeping required for token-level diversity penalties on the GPU.

So unless I'm missing another part of the beam search pipeline where token-level diversity penalties are applied later, this implementation does not appear to be equivalent to either:

• Diverse Beam Search (Vijayakumar et al., 2016), or
• Hugging Face's Hamming diversity implementation.

I'd be interested to know whether this was an intentional approximation for performance reasons, or whether there's another paper or internal design document motivating this implementation.

---

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206