Skip to content

vllm.v1.worker.gpu.spec_decode.rejection_sampler

Functions:

_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 
40
41
42
43
44
45
46
47
48
49
50
51
52
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

get_max_chunk_logits(vocab_size)

Largest number of logits rows one verification chunk may hold.

Source code in vllm/v1/worker/gpu/spec_decode/rejection_sampler.py 
35
36
37
def get_max_chunk_logits(vocab_size: int) -> int:
    """Largest number of logits rows one verification chunk may hold."""
    return max(1, MAX_CHUNK_BYTES // (vocab_size * _FP32_BYTES))