Optimize beam-search cache indirection for short generation
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 14.7k
- Forks
- 2.8k
- Avg merge
- 2d 23h
- Merged PRs (30d)
- 489
Description
Motivation
updateCacheIndirectionKernel currently launches one warp for every 32 positions of nMaxSeqLen, even though only the generated suffix between inputLength and sequenceLength - 1 can differ between beams. For short-output, large-maxSeqLen beam-search workloads, most launched warps return without useful work.
Proposed change
Launch one warp per output beam and let its lanes iterate only over the generated suffix. Preserve parent-beam selection, cyclic-cache addressing, sink tokens, finished-beam behavior, and the last-token mapping.
Direct-kernel evidence
On an uncontended H100 (SM90, CUDA 13.0), beam width 80, maxSeqLen=1024, inputLength=8, sequenceLength=13, nine repeats of 20,000 CUDA-event-timed launches after 1,000 warmups:
- current kernel median: 4.025690 us/call
- proposed kernel median: 2.297330 us/call
- delta: -42.9% (1.75x kernel speedup)
The baseline and candidate implementations were compiled into the same binary and their relevant output range was compared before timing. This is isolated kernel evidence, not a claim of equal end-to-end request speedup.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Search the codebase for updateCacheIndirectionKernel and read its launch configuration and indexing logic first. Implement the proposed one-warp-per-output-beam traversal while preserving parent-beam selection, cyclic-cache addressing, sink tokens, finished-beam behavior, and last-token mapping. Reproduce the direct-kernel comparison described in the issue and verify the relevant output range and timing.
Written by the indexing model from the issue text.
Assessment
- Domain
- performance
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100