Skip to content

vllm.v1.worker.gpu.spec_decode.rejection_sampler

_iter_request_chunks(cu_num_logits, max_chunk_logits)

Yield maximally packed request ranges without splitting requests.

Source code in vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
def _iter_request_chunks(
    cu_num_logits: np.ndarray, max_chunk_logits: int
) -> Iterator[tuple[int, int]]:
    """Yield maximally packed request ranges without splitting requests."""
    assert max_chunk_logits > 0
    num_reqs = cu_num_logits.size - 1
    start = 0
    while start < num_reqs:
        max_logit = int(cu_num_logits[start]) + max_chunk_logits
        end = int(np.searchsorted(cu_num_logits, max_logit, side="right") - 1)
        end = min(num_reqs, max(start + 1, end))
        yield start, end
        start = end