Skip to content

vllm.v1.worker.ubatch_utils

UBatchSlices module-attribute 

UBatchSlices: TypeAlias = list[UBatchSlice]

UBatchSlice dataclass

Source code in vllm/v1/worker/ubatch_utils.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
@dataclass
class UBatchSlice:
    request_slice: slice
    token_slice: slice

    def is_empty(self) -> bool:
        return self.request_slice.start == self.request_slice.stop \
            or self.token_slice.start == self.token_slice.stop

    @property
    def num_tokens(self) -> int:
        return self.token_slice.stop - self.token_slice.start

num_tokens property 

num_tokens: int

request_slice instance-attribute 

request_slice: slice

token_slice instance-attribute 

token_slice: slice

__init__ 

__init__(request_slice: slice, token_slice: slice) -> None

is_empty 

is_empty() -> bool
Source code in vllm/v1/worker/ubatch_utils.py 
13
14
15
def is_empty(self) -> bool:
    return self.request_slice.start == self.request_slice.stop \
        or self.token_slice.start == self.token_slice.stop

is_second_ubatch_empty 

is_second_ubatch_empty(
    orig_num_tokens_per_ubatch: int,
    padded_num_tokens_per_ubatch: int,
) -> bool
Source code in vllm/v1/worker/ubatch_utils.py 
25
26
27
def is_second_ubatch_empty(orig_num_tokens_per_ubatch: int,
                           padded_num_tokens_per_ubatch: int) -> bool:
    return padded_num_tokens_per_ubatch >= 2 * orig_num_tokens_per_ubatch