Lilicorr Chain Walk#

void trt_edgellm::kernel::launchLiLiCorrChainWalk(
rt::Tensor const &candidateIds,
rt::Tensor const &transitionScores,
rt::Tensor const &uniforms,
rt::Tensor const &temperatures,
rt::Tensor const &greedyMask,
rt::Tensor &proposalTokenIds,
rt::Tensor &proposalProbabilities,
cudaStream_t stream
)#

Walks the LiLiCorr candidate lattice one slot at a time, starting from candidate index 0. transitionScores[b,s,prev,next] scores candidate next of slot s given the index chosen at slot s-1. Greedy rows (greedyMask[b] != 0 or temperatures[b] < 1e-3) take the argmax with ties to the lowest index and emit a one-hot q; other rows sample softmax(row / temperature) by inverse CDF on uniforms[b,s] and emit it as q.

Parameters:

  • candidateIds – [B,S,K] INT32, K <= kLiLiCorrMaxTopK

  • transitionScores – [B,S,K,K] FP32

  • uniforms – [B,S] FP32

  • temperatures – [B] FP32

  • greedyMask – [B] INT32

  • proposalTokenIds – [B,S] INT32 output

  • proposalProbabilities – [B,S,K] FP32 output