Skip to content

vllm.v1.worker.block_table

logger module-attribute 

logger = init_logger(__name__)

BlockTable

Source code in vllm/v1/worker/block_table.py 
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
class BlockTable:

    def __init__(
        self,
        max_num_reqs: int,
        max_num_blocks_per_req: int,
        max_num_batched_tokens: int,
        pin_memory: bool,
        device: torch.device,
    ):
        self.max_num_reqs = max_num_reqs
        self.max_num_blocks_per_req = max_num_blocks_per_req
        self.max_num_batched_tokens = max_num_batched_tokens
        self.pin_memory = pin_memory
        self.device = device

        self.block_table = torch.zeros(
            (max_num_reqs, max_num_blocks_per_req),
            device=self.device,
            dtype=torch.int32,
        )
        self.block_table_cpu = torch.zeros(
            (max_num_reqs, max_num_blocks_per_req),
            device="cpu",
            dtype=torch.int32,
            pin_memory=pin_memory,
        )
        self.block_table_np = self.block_table_cpu.numpy()
        self.num_blocks_per_row = np.zeros(max_num_reqs, dtype=np.int32)

        self.slot_mapping_cpu = torch.zeros(self.max_num_batched_tokens,
                                            dtype=torch.int64,
                                            device="cpu",
                                            pin_memory=self.pin_memory)
        self.slot_mapping_np = self.slot_mapping_cpu.numpy()
        self.slot_mapping = torch.zeros(self.max_num_batched_tokens,
                                        dtype=torch.int64,
                                        device=self.device)

    def append_row(
        self,
        block_ids: list[int],
        row_idx: int,
    ) -> None:
        if not block_ids:
            return
        num_blocks = len(block_ids)
        start = self.num_blocks_per_row[row_idx]
        self.num_blocks_per_row[row_idx] += num_blocks
        self.block_table_np[row_idx, start:start + num_blocks] = block_ids

    def add_row(self, block_ids: list[int], row_idx: int) -> None:
        self.num_blocks_per_row[row_idx] = 0
        self.append_row(block_ids, row_idx)

    def move_row(self, src: int, tgt: int) -> None:
        num_blocks = self.num_blocks_per_row[src]
        self.block_table_np[tgt, :num_blocks] = self.block_table_np[
            src, :num_blocks]
        self.num_blocks_per_row[tgt] = num_blocks

    def swap_row(self, src: int, tgt: int) -> None:
        num_blocks_src = self.num_blocks_per_row[src]
        num_blocks_tgt = self.num_blocks_per_row[tgt]
        self.num_blocks_per_row[src] = num_blocks_tgt
        self.num_blocks_per_row[tgt] = num_blocks_src

        self.block_table_np[[src, tgt]] = self.block_table_np[[tgt, src]]

    def commit(self, num_reqs: int) -> None:
        self.block_table[:num_reqs].copy_(self.block_table_cpu[:num_reqs],
                                          non_blocking=True)

    def clear(self) -> None:
        self.block_table.fill_(0)
        self.block_table_cpu.fill_(0)

    def get_device_tensor(self) -> torch.Tensor:
        """Ruturns the device tensor of the block table."""
        return self.block_table

    def get_cpu_tensor(self) -> torch.Tensor:
        """Returns the CPU tensor of the block table."""
        return self.block_table_cpu

    def get_numpy_array(self) -> np.ndarray:
        """Returns the numpy array of the block table."""
        return self.block_table_np

block_table instance-attribute 

block_table = zeros(
    (max_num_reqs, max_num_blocks_per_req),
    device=device,
    dtype=int32,
)

block_table_cpu instance-attribute 

block_table_cpu = zeros(
    (max_num_reqs, max_num_blocks_per_req),
    device="cpu",
    dtype=int32,
    pin_memory=pin_memory,
)

block_table_np instance-attribute 

block_table_np = numpy()

device instance-attribute 

device = device

max_num_batched_tokens instance-attribute 

max_num_batched_tokens = max_num_batched_tokens

max_num_blocks_per_req instance-attribute 

max_num_blocks_per_req = max_num_blocks_per_req

max_num_reqs instance-attribute 

max_num_reqs = max_num_reqs

num_blocks_per_row instance-attribute 

num_blocks_per_row = zeros(max_num_reqs, dtype=int32)

pin_memory instance-attribute 

pin_memory = pin_memory

slot_mapping instance-attribute 

slot_mapping = zeros(
    max_num_batched_tokens, dtype=int64, device=device
)

slot_mapping_cpu instance-attribute 

slot_mapping_cpu = zeros(
    max_num_batched_tokens,
    dtype=int64,
    device="cpu",
    pin_memory=pin_memory,
)

slot_mapping_np instance-attribute 

slot_mapping_np = numpy()

__init__ 

__init__(
    max_num_reqs: int,
    max_num_blocks_per_req: int,
    max_num_batched_tokens: int,
    pin_memory: bool,
    device: device,
)
Source code in vllm/v1/worker/block_table.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
def __init__(
    self,
    max_num_reqs: int,
    max_num_blocks_per_req: int,
    max_num_batched_tokens: int,
    pin_memory: bool,
    device: torch.device,
):
    self.max_num_reqs = max_num_reqs
    self.max_num_blocks_per_req = max_num_blocks_per_req
    self.max_num_batched_tokens = max_num_batched_tokens
    self.pin_memory = pin_memory
    self.device = device

    self.block_table = torch.zeros(
        (max_num_reqs, max_num_blocks_per_req),
        device=self.device,
        dtype=torch.int32,
    )
    self.block_table_cpu = torch.zeros(
        (max_num_reqs, max_num_blocks_per_req),
        device="cpu",
        dtype=torch.int32,
        pin_memory=pin_memory,
    )
    self.block_table_np = self.block_table_cpu.numpy()
    self.num_blocks_per_row = np.zeros(max_num_reqs, dtype=np.int32)

    self.slot_mapping_cpu = torch.zeros(self.max_num_batched_tokens,
                                        dtype=torch.int64,
                                        device="cpu",
                                        pin_memory=self.pin_memory)
    self.slot_mapping_np = self.slot_mapping_cpu.numpy()
    self.slot_mapping = torch.zeros(self.max_num_batched_tokens,
                                    dtype=torch.int64,
                                    device=self.device)

add_row 

add_row(block_ids: list[int], row_idx: int) -> None
Source code in vllm/v1/worker/block_table.py 
64
65
66
def add_row(self, block_ids: list[int], row_idx: int) -> None:
    self.num_blocks_per_row[row_idx] = 0
    self.append_row(block_ids, row_idx)

append_row 

append_row(block_ids: list[int], row_idx: int) -> None
Source code in vllm/v1/worker/block_table.py 
52
53
54
55
56
57
58
59
60
61
62
def append_row(
    self,
    block_ids: list[int],
    row_idx: int,
) -> None:
    if not block_ids:
        return
    num_blocks = len(block_ids)
    start = self.num_blocks_per_row[row_idx]
    self.num_blocks_per_row[row_idx] += num_blocks
    self.block_table_np[row_idx, start:start + num_blocks] = block_ids

clear 

clear() -> None
Source code in vllm/v1/worker/block_table.py 
86
87
88
def clear(self) -> None:
    self.block_table.fill_(0)
    self.block_table_cpu.fill_(0)

commit 

commit(num_reqs: int) -> None
Source code in vllm/v1/worker/block_table.py 
82
83
84
def commit(self, num_reqs: int) -> None:
    self.block_table[:num_reqs].copy_(self.block_table_cpu[:num_reqs],
                                      non_blocking=True)

get_cpu_tensor 

get_cpu_tensor() -> Tensor

Returns the CPU tensor of the block table.

Source code in vllm/v1/worker/block_table.py 
94
95
96
def get_cpu_tensor(self) -> torch.Tensor:
    """Returns the CPU tensor of the block table."""
    return self.block_table_cpu

get_device_tensor 

get_device_tensor() -> Tensor

Ruturns the device tensor of the block table.

Source code in vllm/v1/worker/block_table.py 
90
91
92
def get_device_tensor(self) -> torch.Tensor:
    """Ruturns the device tensor of the block table."""
    return self.block_table

get_numpy_array 

get_numpy_array() -> ndarray

Returns the numpy array of the block table.

Source code in vllm/v1/worker/block_table.py 
 98
 99
100
def get_numpy_array(self) -> np.ndarray:
    """Returns the numpy array of the block table."""
    return self.block_table_np

move_row 

move_row(src: int, tgt: int) -> None
Source code in vllm/v1/worker/block_table.py 
68
69
70
71
72
def move_row(self, src: int, tgt: int) -> None:
    num_blocks = self.num_blocks_per_row[src]
    self.block_table_np[tgt, :num_blocks] = self.block_table_np[
        src, :num_blocks]
    self.num_blocks_per_row[tgt] = num_blocks

swap_row 

swap_row(src: int, tgt: int) -> None
Source code in vllm/v1/worker/block_table.py 
74
75
76
77
78
79
80
def swap_row(self, src: int, tgt: int) -> None:
    num_blocks_src = self.num_blocks_per_row[src]
    num_blocks_tgt = self.num_blocks_per_row[tgt]
    self.num_blocks_per_row[src] = num_blocks_tgt
    self.num_blocks_per_row[tgt] = num_blocks_src

    self.block_table_np[[src, tgt]] = self.block_table_np[[tgt, src]]

MultiGroupBlockTable

The BlockTables for each KV cache group.

Source code in vllm/v1/worker/block_table.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
class MultiGroupBlockTable:
    """The BlockTables for each KV cache group."""

    def __init__(self, max_num_reqs: int, max_model_len: int,
                 max_num_batched_tokens: int, pin_memory: bool,
                 device: torch.device, block_sizes: list[int]) -> None:
        self.block_tables = [
            BlockTable(max_num_reqs, cdiv(max_model_len, block_size),
                       max_num_batched_tokens, pin_memory, device)
            for block_size in block_sizes
        ]

    def append_row(self, block_ids: tuple[list[int], ...],
                   row_idx: int) -> None:
        for i, block_table in enumerate(self.block_tables):
            block_table.append_row(block_ids[i], row_idx)

    def add_row(self, block_ids: tuple[list[int], ...], row_idx: int) -> None:
        for i, block_table in enumerate(self.block_tables):
            block_table.add_row(block_ids[i], row_idx)

    def move_row(self, src: int, tgt: int) -> None:
        for block_table in self.block_tables:
            block_table.move_row(src, tgt)

    def swap_row(self, src: int, tgt: int) -> None:
        for block_table in self.block_tables:
            block_table.swap_row(src, tgt)

    def commit(self, num_reqs: int) -> None:
        for block_table in self.block_tables:
            block_table.commit(num_reqs)

    def clear(self) -> None:
        for block_table in self.block_tables:
            block_table.clear()

    def __getitem__(self, idx: int) -> "BlockTable":
        """Returns the BlockTable for the i-th KV cache group."""
        return self.block_tables[idx]

block_tables instance-attribute 

block_tables = [
    BlockTable(
        max_num_reqs,
        cdiv(max_model_len, block_size),
        max_num_batched_tokens,
        pin_memory,
        device,
    )
    for block_size in block_sizes
]

__getitem__ 

__getitem__(idx: int) -> BlockTable

Returns the BlockTable for the i-th KV cache group.

Source code in vllm/v1/worker/block_table.py 
140
141
142
def __getitem__(self, idx: int) -> "BlockTable":
    """Returns the BlockTable for the i-th KV cache group."""
    return self.block_tables[idx]

__init__ 

__init__(
    max_num_reqs: int,
    max_model_len: int,
    max_num_batched_tokens: int,
    pin_memory: bool,
    device: device,
    block_sizes: list[int],
) -> None
Source code in vllm/v1/worker/block_table.py 
106
107
108
109
110
111
112
113
def __init__(self, max_num_reqs: int, max_model_len: int,
             max_num_batched_tokens: int, pin_memory: bool,
             device: torch.device, block_sizes: list[int]) -> None:
    self.block_tables = [
        BlockTable(max_num_reqs, cdiv(max_model_len, block_size),
                   max_num_batched_tokens, pin_memory, device)
        for block_size in block_sizes
    ]

add_row 

add_row(
    block_ids: tuple[list[int], ...], row_idx: int
) -> None
Source code in vllm/v1/worker/block_table.py 
120
121
122
def add_row(self, block_ids: tuple[list[int], ...], row_idx: int) -> None:
    for i, block_table in enumerate(self.block_tables):
        block_table.add_row(block_ids[i], row_idx)

append_row 

append_row(
    block_ids: tuple[list[int], ...], row_idx: int
) -> None
Source code in vllm/v1/worker/block_table.py 
115
116
117
118
def append_row(self, block_ids: tuple[list[int], ...],
               row_idx: int) -> None:
    for i, block_table in enumerate(self.block_tables):
        block_table.append_row(block_ids[i], row_idx)

clear 

clear() -> None
Source code in vllm/v1/worker/block_table.py 
136
137
138
def clear(self) -> None:
    for block_table in self.block_tables:
        block_table.clear()

commit 

commit(num_reqs: int) -> None
Source code in vllm/v1/worker/block_table.py 
132
133
134
def commit(self, num_reqs: int) -> None:
    for block_table in self.block_tables:
        block_table.commit(num_reqs)

move_row 

move_row(src: int, tgt: int) -> None
Source code in vllm/v1/worker/block_table.py 
124
125
126
def move_row(self, src: int, tgt: int) -> None:
    for block_table in self.block_tables:
        block_table.move_row(src, tgt)

swap_row 

swap_row(src: int, tgt: int) -> None
Source code in vllm/v1/worker/block_table.py 
128
129
130
def swap_row(self, src: int, tgt: int) -> None:
    for block_table in self.block_tables:
        block_table.swap_row(src, tgt)