[Performance]: Avoid full priority-queue scans for bounded admission peeks #21912

Proposal to improve performance

PriorityWaitingQueue.peek_requests(k) currently uses heapq.nsmallest, which scans every queued request and maintains a size-k heap. With priority scheduling and KV-balanced attention-DP admission enabled, this runs every admitting scheduler iteration; the configured admission window is normally 32 requests.

For a bounded peek, the queue's existing heap structure can enumerate the next requests through a best-first frontier. This visits only the relevant heap branches, reducing the work from O(n log k) to O(k log k) without copying, sorting, or mutating the full queue. Large result sets can retain the existing nsmallest path.

I have a focused patch and CPU-only equivalence tests ready.

Report of performance regression

Benchmark on current main (a545268cf3e6), using the real PriorityWaitingQueue implementation with identical randomized priority heaps. Each process used 5 warmups and 31 timed calls; the table reports the median of five process medians.

At the default 32-request KV-balance window, an 8,192-request waiting queue saves about 0.33 ms per admitting iteration.

Misc discussion on performance

The output order is identical to repeated priority-queue pops, and the source heap remains unchanged. I found no open pull request touching this function and no open or closed issue covering this overhead.

Your current environment

• OS: Linux x86_64
• CPU: Intel Xeon Platinum 8352Y
• Python: 3.11.18
• GPU/CUDA: not used
• TensorRT-LLM: main at a545268cf3e6

Before submitting a new issue...

• I searched relevant open and closed issues and pull requests.