1. Block 概覽

vLLM 的一個很大創新點是將物理層面的 GPU 和 CPU 可用內存切分成若干個 block,這樣可以有效降低內存碎片化問題。具體而言，vLLM 的 block 分爲邏輯層面（logical）和物理層面（physical），二者之間存在映射關係。下圖很好解釋了兩個層面 block 的關係。

假設每個 block 可以用來存 4 個 token 的kv cache數據。一個句子的 token在邏輯層面是緊鄰的，每次 decoding 生成新的 token 就往空閒的 block 裏放。但是對應到物理層面的 block，一個句子的 token 可能分佈在並不相鄰的 block內，不過沒關係，vLLM 會爲每個句子的每個 token記錄邏輯和物理block 的映射關係，方便查找和讀取。

接下來我們詳細介紹 block 大小的含義，以及 block 的數量是如何計算的，最後介紹 vLLM 是如何管理 block 的。

2. Block 大小如何計算

block 的大小可以自定義，上面定義爲 4，簡單理解就是每個 block 最多存儲 4 個 token 的 kv cache 數據。但是 block 設置爲 4 的時候對應到 GPU 內存到底是多大呢？其實這很好計算，

一個 block 佔用內存大小（Byte）= token 數量 (block_size) ✖️ 一個 token 的 kv cache 佔用內存大小。

所以，我們只需要計算出單個 token 的 kv cache 對應的大小即可。block 大小的計算方法由vllm/vllm/worker/cache_engine.py文件裏CacheEngine類的get_cache_block_size函數實現，代碼也很簡單，簡化後如下：

# vllm/vllm/worker/cache_engine.py
class CacheEngine:
    @staticmethod
    def get_cache_block_size(
        block_size: int,
        cache_dtype: str,
        model_config: ModelConfig,
        parallel_config: ParallelConfig,
    ) -> int:
        head_size = model_config.get_head_size()
        num_heads = model_config.get_num_kv_heads(parallel_config)
        num_layers = model_config.get_num_layers(parallel_config)

        key_cache_block = block_size * num_heads * head_size
        value_cache_block = key_cache_block
        total = num_layers * (key_cache_block + value_cache_block)
        if cache_dtype == "auto":
            dtype = model_config.dtype
        else:
            dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype]
        dtype_size = _get_dtype_size(dtype)
        return dtype_size * total

上面代碼中首先拿到 num_heads和head_size兩個變量的值，
num_heads * head_size就表示單個 token 在單層多頭注意力機制計算中所需要的參數量，不過這只是 key 或者 value cache 所佔用的參數量。

一個 block 佔用的內存 = token 數量（block_size）✖️ 層數 (num_layers) ✖️ 單層 kv cache 佔用內存（2✖️num_heads✖️head_size）✖️ 數據類型大小（如果是 fp16，則每個數據佔用 2 Bytes）

舉例來說，假設 block_size=4， num_layers=4, num_heads=8, heads_size=128，採用 fp16 存儲數據，那麼

一個 block 佔用內存大小 = 4 ✖️ 4 ✖️ 8 ✖️ 128 ✖️ 2 = 32,768 Bytes。

總結，一個 block 所佔用的內存大小就是 block_size 個 token kv cache 所佔內存的總和。不同模型的 block 各不相同。

2. Block 數量如何計算

block 數量計算由vllm/vllm/worker/worker.py文件中Worker類的profile_num_available_blocks函數實現，該函數很簡單，簡化代碼如下：

class Worker
    @torch.inference_mode()
    def profile_num_available_blocks(
        self,
        block_size: int,
        gpu_memory_utilization: float,
        cpu_swap_space: int,
        cache_dtype: str,
    ) -> Tuple[int, int]:
        torch.cuda.empty_cache()
		
		# 這一行其實就是用模擬數據跑一下forward 來統計GPU 的使用情況
        self.model_runner.profile_run()

        torch.cuda.synchronize()
        free_gpu_memory, total_gpu_memory = torch.cuda.mem_get_info()
        peak_memory = total_gpu_memory - free_gpu_memory

        cache_block_size = CacheEngine.get_cache_block_size(
            block_size, cache_dtype, self.model_config, self.parallel_config)
        num_gpu_blocks = int(
            (total_gpu_memory * gpu_memory_utilization - peak_memory) //
            cache_block_size)
        num_cpu_blocks = int(cpu_swap_space // cache_block_size)
        num_gpu_blocks = max(num_gpu_blocks, 0)
        num_cpu_blocks = max(num_cpu_blocks, 0)
        if self.model_runner.lora_manager:
            self.model_runner.remove_all_loras()
        gc.collect()
        torch.cuda.empty_cache()
        return num_gpu_blocks, num_cpu_blocks

整個函數的邏輯很清晰，簡單理解就是先用模擬數據跑一次 forward 記錄下 GPU 的使用情況，這樣可以知道 peak memory，然後計算每個 block 需要用到的 memory，接着就可以計算出 block 數量了。具體而言：

13 行：vllm 默認用 256 個句子來做 profile，每個句子長度爲 128
15 到 17 行：統計 GPU 內存使用情況，返回的是以字節（Byte）爲單位的數值，後面也都是基於 Byte 爲單位進行計算的
19 行：計算每個 block 的大小，這個在前面已經介紹。
20-23 行：計算可用的 GPU block 數量。num_gpu_blocks = int( (total_gpu_memory * gpu_memory_utilization - peak_memory) // cache_block_size)：gpu_memory_utilization: 默認值是 0.9，表示 GPU 內存利用率是 90%，這挺高的了。所以最終的可用 GPU block 數量等於剩餘 GPU 內存大小除以每個 block 的大小
24 行：計算可用的 CPU block 數量。 num_cpu_blocks = int(cpu_swap_space // cache_block_size)這裏的cpu_swap_space 代表每個 GPU 對應的 CPU swap 空間大小，單位是（GB），默認是是 4。也就是說每個 GPU 對應的 CPU swap 空間大小是 4 GB。

3. Block 如何管理？

3.1 邏輯 Block 定義和使用

邏輯 Block（LogicalTokenBlock）定義如下：

# vllm/vllm/block.py
class LogicalTokenBlock:
    def __init__(
        self,
        block_number: int,
        block_size: int,
    ) -> None:
        self.block_number = block_number
        self.block_size = block_size

        self.token_ids = [_BLANK_TOKEN_ID] * block_size
        self.num_tokens = 0

    def is_empty(self) -> bool:
        return self.num_tokens == 0

    def get_num_empty_slots(self) -> int:
        return self.block_size - self.num_tokens

    def is_full(self) -> bool:
        return self.num_tokens == self.block_size

    def append_tokens(self, token_ids: List[int]) -> None:
        assert len(token_ids) <= self.get_num_empty_slots()
        curr_idx = self.num_tokens
        self.token_ids[curr_idx:curr_idx + len(token_ids)] = token_ids
        self.num_tokens += len(token_ids)

    def get_token_ids(self) -> List[int]:
        return self.token_ids[:self.num_tokens]

    def get_last_token_id(self) -> int:
        assert self.num_tokens > 0
        return self.token_ids[self.num_tokens - 1]

block_number: int: 這個是 PhysicalTokenBlock 實例對象的索引，可以理解成是 flag，用於區分不同 block
block_size: int：表示一個 block 內存儲多少個 token 的 kv cache 數據。
__init__函數中self.token_ids初始化是一個長度爲 block_size 的全爲 -1 的list。後續可以通過append_tokens將新的 token添加到這個 list 中去。
self.num_tokens會統計已使用的 token 數量，當self.num_tokens==block_size時則表示這個 block 已經被使用完了。

邏輯 Block 的使用邏輯是根據需要實時實例化一個對象，如果當前的 LogicalBlock沒有剩餘空間了，就再實例化一個新的。

在 vLLm 的使用場景是在vllm/vllm/sequence.py裏的Sequence類中根據需要動態創建LogicalBlock。

Sequence類在之前介紹 vLLM 的文章【大模型推理框架 vLLM 源碼解析（一）】中已經有詳細介紹，這裏你只需要知道這個類記錄了每個輸入句子整個推理過程（prefilling 和 decoding）的所有信息。

我們結合代碼來看會更好理解，如下：

# vllm/vllm/sequence.py
class Sequence:
	def __init__(self, ...):
		...
		self.logical_token_blocks: List[LogicalTokenBlock] = []
		self._append_tokens_to_blocks(prompt_token_ids)
		...

    def _append_tokens_to_blocks(self, token_ids: List[int]) -> None:
        cursor = 0
        while cursor < len(token_ids):
            if not self.logical_token_blocks:
                self._append_logical_block()

            last_block = self.logical_token_blocks[-1]
            if last_block.is_full():
                self._append_logical_block()
                last_block = self.logical_token_blocks[-1]

            num_empty_slots = last_block.get_num_empty_slots()
            last_block.append_tokens(token_ids[cursor:cursor +
                                               num_empty_slots])
            cursor += num_empty_slots
			
    def _append_logical_block(self) -> None:
        block = LogicalTokenBlock(
            block_number=len(self.logical_token_blocks),
            block_size=self.block_size,
        )
        self.logical_token_blocks.append(block)

__init__函數中會初始化self.logical_token_blocks空數組，用來存LogicalBlock。可以看到會先將 prompt 的所有 token 通過_append_tokens_to_blocks存入到 block 中
_append_tokens_to_blocks函數會遍歷傳入的 token_ids 數組中的每個 token id，將該 token 信息存入到 LogicalBlock 中。
- 第 12 行：如果self.logical_token_blocks爲空，則會動態調用_append_logical_block來創建一個LogicalBlock，並存到self.logical_token_blocks變量中去
- 第 16 行：如果最新創建的LogicalBlock空間已經滿了，則同樣會動態調用_append_logical_block來創建一個新的LogicalBlock

3.2 物理Block 定義和管理

物理 Block (PhysicalTokenBlock)的代碼定義如下：

device: Device: 是一個 enum.Enum 實例對象，要麼是 CPU 要麼是 GPU。
self.ref_count 變量用來指示這個 block 被使用的次數，默認爲 0，代表沒有使用。可以大於等於1，表示這個 block 內 token的 cache 被重複利用，使用場景比如可以是 beam search，這樣可以重複利用cache，減少內存開銷。

# vllm/vllm/block.py
class PhysicalTokenBlock:
    def __init__(
        self,
        device: Device,
        block_number: int,
        block_size: int,
    ) -> None:
        self.device = device
        self.block_number = block_number
        self.block_size = block_size

        self.ref_count = 0

    def __repr__(self) -> str:
        return (f'PhysicalTokenBlock(device={self.device}, '
                f'block_number={self.block_number}, '
                f'ref_count={self.ref_count})')

PhysicalTokenBlock只是針對單個 block 的描述。vLLM 在vllm/vllm/core/block_manager.py文件下實現了BlockAllocator類用來初始化所有物理 block，並負責分配這些 block。

BlockAllocator這個類代碼很簡單，如下。主要作用有三個：

__init__: 初始化指定數量的物理層面 block，這個數量在前面一節已經介紹過如何計算。
allocate: 通過 list的 pop() 函數返回一個可用的 block，並將該 block 的ref_count設置爲 1
free：回收一個指定的 PhysicalBlock，但是回收的前提是這個 block 的ref_count變量值爲 0，表示這個 block 內的 token kv cache 數據不再需要了。

# vllm/vllm/core/block_manager.py
class BlockAllocator:
    def __init__(
        self,
        device: Device,
        block_size: int,
        num_blocks: int,
    ) -> None:
        self.device = device
        self.block_size = block_size
        self.num_blocks = num_blocks

        # Initialize the free blocks.
        self.free_blocks: BlockTable = []
        for i in range(num_blocks):
            block = PhysicalTokenBlock(device=device,
                                       block_number=i,
                                       block_size=block_size)
            self.free_blocks.append(block)

    def allocate(self) -> PhysicalTokenBlock:
        if not self.free_blocks:
            raise ValueError("Out of memory! No free blocks are available.")
        block = self.free_blocks.pop()
        block.ref_count = 1
        return block

    def free(self, block: PhysicalTokenBlock) -> None:
        if block.ref_count == 0:
            raise ValueError(f"Double free! {block} is already freed.")
        block.ref_count -= 1
        if block.ref_count == 0:
            self.free_blocks.append(block)

    def get_num_free_blocks(self) -> int:
        return len(self.free_blocks)

3.3 Block 管理和映射模塊

在介紹這個Block 管理模塊之前，我們先了解 vLLM 中設置的用來判斷句子是否能夠被分配物理 Block 的三種狀態，代碼如下：

# vllm/vllm/core/block_manager.py
class AllocStatus(enum.Enum):
    """Result for BlockSpaceManager.can_allocate
    """
    OK = enum.auto()
    LATER = enum.auto()
    NEVER = enum.auto()

三種狀態的含義如下：

OK: seq_group 可以現在被分配。
LATER: seq_group 不能被分配。分配器的容量大於 seq_group 所需。
NEVER: seq_group 永遠不能被分配。seq_group 太大，無法在 GPU 中分配。

vllm/vllm/core/block_manager.py下的BlockSpaceManager是一個高級內存管理器，它在內存密集型計算任務（尤其是在使用GPU和CPU進行大規模數據處理的情況下）中管理邏輯數據塊和物理內存塊之間的映射。

接下來，我們結合代碼介紹BlockSpaceManager一些重要的函數。

初始化函數__init__:
- watermark: 一種閾值機制，用來決定何時停止在GPU上分配新的塊，以避免內存不足
- watermark_blocks: 計算出在達到內存不足前，還能在GPU上分配多少個塊。
- sliding_window: 可選參數，用來限制在任意給定時間內活躍的邏輯塊的數量，有助於控制內存使用。
- 創建了 cpu 和 gpu 兩種 BlockAllocator,不過需要注意這裏都是物理層面的 Block
- 創建了一個字典 block_tables，用於存儲每個 sequence id 和它所使用的物理塊之間的映射。通過這個 sequence id ，我們就能找到對應的前面介紹的Sequence實例化對象，通過這個字典，就建立了邏輯 block 和物理 block 的映射關係。

# vllm/vllm/core/block_manager.py
class BlockSpaceManager:
    def __init__(
        self,
        block_size: int,
        num_gpu_blocks: int,
        num_cpu_blocks: int,
        watermark: float = 0.01,
        sliding_window: Optional[int] = None,
    ) -> None:
        self.block_size = block_size
        self.num_total_gpu_blocks = num_gpu_blocks
        self.num_total_cpu_blocks = num_cpu_blocks

        self.block_sliding_window = None
        if sliding_window is not None:
            assert sliding_window % block_size == 0, (sliding_window,
                                                      block_size)
            self.block_sliding_window = sliding_window // block_size

        self.watermark = watermark
        assert watermark >= 0.0

        self.watermark_blocks = int(watermark * num_gpu_blocks)
        self.gpu_allocator = BlockAllocator(Device.GPU, block_size,
                                            num_gpu_blocks)
        self.cpu_allocator = BlockAllocator(Device.CPU, block_size,
                                            num_cpu_blocks)
        # Mapping: seq_id -> BlockTable.
        self.block_tables: Dict[int, BlockTable] = {}

can_allocate

class BlockSpaceManager:
    def can_allocate(self, seq_group: SequenceGroup) -> AllocStatus:
        seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0]
        num_required_blocks = len(seq.logical_token_blocks)

        if seq_group.prefix is not None and seq_group.prefix.allocated:
            num_required_blocks -= seq_group.prefix.get_num_blocks()

        if self.block_sliding_window is not None:
            num_required_blocks = min(num_required_blocks,
                                      self.block_sliding_window)
        num_free_gpu_blocks = self.gpu_allocator.get_num_free_blocks()

        # Use watermark to avoid frequent cache eviction.
        if (self.num_total_gpu_blocks - num_required_blocks <
                self.watermark_blocks):
            return AllocStatus.NEVER
        if num_free_gpu_blocks - num_required_blocks >= self.watermark_blocks:
            return AllocStatus.OK
        else:
            return AllocStatus.LATER

can_allocate方法用於判斷一個序列組（seq_group）是否能被成功分配所需的內存塊。此方法首先計算該序列組基於當前任務的邏輯數據塊所需的總物理內存塊數量。接着，它會檢查GPU分配器中的空閒內存塊數量，以確認是否有足夠的資源滿足需求。

方法中引入了watermark_blocks概念，其主要目的是防止因頻繁進行內存塊的緩存淘汰而影響系統性能。在模型訓練或數據處理的動態環境中，內存需求持續變化，如果因缺乏足夠的空閒內存塊而不得不頻繁淘汰並重新分配內存塊，將會造成性能損耗。這是因爲被淘汰的內存塊很可能很快再次需要使用，其重新分配過程會消耗額外的時間和資源。

通過設置watermark_blocks閾值，當GPU上的空閒內存塊數量低於此閾值時，系統將避免分配新的內存塊，以留出緩衝區域，減少緩存淘汰的發生。只有當空閒內存塊數量高於此閾值時，系統纔會繼續進行新的內存塊分配。這種策略旨在平衡內存分配需求和系統性能，避免因頻繁的內存操作而降低效率。

如果根據當前的資源狀態，確定序列組所需的內存塊永遠無法被滿足，則返回AllocStatus.NEVER，意味着該序列組在當前條件下無法被分配。如果當前不可分配但未來有可能，返回AllocStatus.LATER，表明序列組暫時無法分配，但隨着系統狀態的改變，可能在將來能夠分配。如果有足夠的空閒內存塊滿足分配需求，則返回AllocStatus.OK，表示序列組可以立即被分配所需內存。

這種方式確保了watermark_blocks在滿足內存分配需求的同時，有效避免了頻繁的緩存淘汰問題，從而優化了整體的系統性能和資源利用效率。

allocate 代碼有簡化，但是不影響理解

class BlockSpaceManager:
   def allocate(self, seq_group: SequenceGroup) -> None:
        seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0]
        num_prompt_blocks = len(seq.logical_token_blocks)

        block_table: BlockTable = []
        for logical_idx in range(num_prompt_blocks):
			block = self.gpu_allocator.allocate()
            block.ref_count = seq_group.num_seqs()
            block_table.append(block)

        for seq in seq_group.get_seqs(status=SequenceStatus.WAITING):
            self.block_tables[seq.seq_id] = block_table.copy()

allocate 方法用於爲序列組分配內存塊。它會遍歷序列組中的每個序列，爲每個序列分配足夠的內存塊，並將這些塊添加到序列的塊表中。同時，它會更新序列的塊表，以便在後續的訓練過程中可以正確地訪問這些塊。

BlockSpaceManager還有很多其它的函數，爲了避免文章累贅，這裏不做詳細介紹。

後面會繼續寫一篇 vLLM 的調度Scheduler模塊的文章，對BlockSpaceManager更加詳細地介紹。相信通過本篇文章，你應該能夠對 vLLM 的 block 有一個清楚的瞭解了，如果還是不清楚，可以反覆閱讀直到清楚爲止。

大模型推理框架 vLLM 源碼解析（二）：Block 模塊分配和管理

1. Block 概覽

2. Block 大小如何計算

2. Block 數量如何計算

3. Block 如何管理？

3.1 邏輯 Block 定義和使用

3.2 物理Block 定義和管理

3.3 Block 管理和映射模塊

參考

NSCC集羣使用筆記

Huggingface Transformers實現張量並行的小坑 set/get_output_embeddings

Pytorch 如何使用 storage 實現參數 offload？

TACC 集羣使用筆記

圖解 vLLM 的推理調度策略

Mac下配置sublime實現LaTeX

https://yachay.unat.edu.pe/blog/index.php?comment_area=format_blog&comment_component=blog&comment_co

linux以太網驅動總結