Nano-vLLM 源碼學習

以下內容由Trae生成。我只管問和排版。

Nano-vLLM 是一個從零構建的輕量級vLLM實現項目，具備快速離線推理、代碼可讀性高和優化功能豐富等特點，性能表現和原版vLLM相媲美。以下是該項目各目錄和文件的功能說明：

1. 根目錄：

bench.py：用于存放基準測試代碼，可查看模型性能測試結果。
example.py：提供項目使用示例，展示API調用方法。
pyproject.toml：Python項目配置文件，管理項目依賴和元數據。

2. nanovllm目錄：

__init__.py：標識該目錄為Python包。
config.py：定義配置類 Config，用于管理模型運行的各種參數。
engine 子目錄：存放推理引擎相關代碼，包含 block_manager.py、llm_engine.py、 model_runner.py、scheduler.py 和 sequence.py。
layers 子目錄：存放模型各層的實現代碼，如激活層、注意力層等。
llm.py：可能實現LLM核心功能。
models 子目錄：存放特定模型的實現，如 qwen3.py。
sampling_params.py：可能用于定義采樣參數。
utils 子目錄：存放工具函數，如 context.py 和 loader.py。

一、engine 目錄

engine目錄負責推理引擎的實現，管理模型推理的整個流程，協調各組件工作。
Nano-vLLM engine推理引擎以模塊化設計，主要由 LLMEngine、ModelRunner、Scheduler、BlockManager 和 Sequence 幾個核心組件構成，各組件分工協作完成推理任務。

核心組件及實現原理

Sequence（序列管理）
- 用于管理單個推理序列的狀態和數據，包括 token ID、狀態（等待、運行、完成）、緩存信息等。
- 提供了獲取塊數據、追加 token 等方法，方便對序列進行操作。
BlockManager（塊管理）
- 負責管理緩存塊的分配和釋放，使用哈希算法來實現塊的緩存，提高推理效率。
- 通過 hash_to_block_id 字典記錄哈希值和塊 ID 的映射關系，減少重復計算。
- 提供了分配、釋放和追加塊的方法，確保塊資源的有效利用。
Scheduler（調度器）
- 負責管理待處理和正在運行的序列，根據資源限制和塊管理狀態進行序列調度。
- 分為預填充（prefill）和解碼（decode）兩個階段調度序列，在資源不足時會搶占序列。
- 對推理結果進行后處理，更新序列狀態并釋放完成序列的資源。
ModelRunner（模型運行器）
- 負責初始化模型、預熱模型、分配 KV 緩存等操作。
- 支持張量并行，使用共享內存進行進程間通信。
- 提供了準備輸入數據、運行模型和采樣等方法，支持 CUDA Graph 優化推理性能。
LLMEngine（推理引擎）
- 作為推理引擎的入口，負責初始化模型運行器、調度器和分詞器。
- 提供添加請求、單步推理和生成結果等方法，支持進度顯示和吞吐量統計。
- 使用多進程實現張量并行，提高推理速度。

推理流程

用戶通過 LLMEngine.generate 方法添加推理請求。
Scheduler 將請求加入等待隊列，并在資源充足時進行調度。
ModelRunner 準備輸入數據，調用模型進行推理，并使用 Sampler 進行采樣。
BlockManager 管理緩存塊，提高重復計算的效率。
Scheduler 對推理結果進行后處理，更新序列狀態，直到所有序列完成。

二、layers目錄

layers 目錄存放模型各基礎層的實現代碼，為模型提供基礎的計算單元，保證模型能完成復雜的計算任務。從已有代碼可知，該目錄主要有以下功能：

1、激活層（ activation.py ）

激活層的作用是向神經網絡中引入非線性因素，使得神經網絡可以擬合復雜的非線性函數。
項目中實現了 SiluAndMul 激活層，對輸入張量進行 SiLU 激活和元素乘法操作，引入非線性因素。代碼如下：

class SiluAndMul(nn.Module):def __init__(self):super().__init__()@torch.compiledef forward(self, x: torch.Tensor) -> torch.Tensor:x, y = x.chunk(2, -1)return F.silu(x) * y

SiluAndMul 類繼承自 nn.Module ，在 forward 方法中，先將輸入張量 x 沿最后一個維度分成兩部分，對第一部分應用 SiLU（Sigmoid Linear Unit）激活函數，再與第二部分進行元素乘法運算。

2、注意力層（ attention.py ）：

注意力層是Transformer架構的核心組件，它允許模型在處理序列數據時，動態地關注序列中不同位置的信息。
項目中的 Attention 類實現了注意力機制，支持多頭注意力機制，利用 Triton 優化 KV 緩存存儲，結合 Flash Attention 提高計算效率。代碼如下：

class Attention(nn.Module):def __init__(self,num_heads,head_dim,scale,num_kv_heads,):super().__init__()# ...已有代碼...def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):# ...已有代碼...

Attention 類支持多頭注意力機制，在 forward 方法中，根據上下文 context 的狀態選擇不同的注意力計算方式：

prefill 階段：使用 flash_attn_varlen_func 計算注意力，支持前綴緩存（prefix cache）。
decode 階段：使用 flash_attn_with_kvcache 結合KV緩存計算注意力。
此外，項目還使用 Triton 實現了 store_kvcache_kernel 核函數，用于將鍵（key）和值（value）存儲到緩存中，提高計算效率。

3、采樣（sampler.py）

項目里， Sampler 類是一個繼承自 torch.nn.Module 的模塊，用于實現采樣操作。采樣是從模型輸出的概率分布中選取下一個 token 的過程，在自然語言處理里常用于生成文本。

從 sampler.py 文件可知， Sampler 類的 forward 方法接收 logits （模型未經過 softmax 的原始輸出）和 temperatures 作為輸入，實現了兩種采樣方式：

貪心采樣（Greedy Sampling）：當 temperatures 為 0 時，直接選取 logits 中概率最大的 token。
基于溫度調整的采樣：當 temperatures 不為 0 時，先對 logits 除以溫度值，再通過 torch.softmax 轉換為概率分布，最后結合指數分布對概率進行調整，選取調整后概率最大的 token。
這種采樣機制可以在確定性輸出（貪心采樣）和隨機性輸出（基于溫度調整的采樣）之間進行平衡，從而控制生成文本的多樣性。

4、嵌入層與語言模型頭（embed_head.py）

VocabParallelEmbedding 類：實現了詞表并行的嵌入層，支持在分布式環境下對詞表進行分片處理，不同進程處理詞表的不同部分，最后通過 all_reduce 操作匯總結果。
ParallelLMHead 類：繼承自 VocabParallelEmbedding ，實現了并行的語言模型頭。在 prefill 階段會提取最后一個時間步的特征，然后進行線性變換得到對數幾率（logits）。

5、歸一化層（layernorm.py）

RMSNorm 類：實現了 Root Mean Square (RMS) 歸一化。支持兩種前向傳播方式：一種是普通的 RMS 歸一化，另一種是帶有殘差連接的 RMS 歸一化。

6、線性層（linear.py）

LinearBase 類：線性層的基類，定義了線性層的基本屬性和方法。
ReplicatedLinear 類：實現了普通的線性層，權重在所有進程中保持一致。
ColumnParallelLinear 類：實現了列并行的線性層，輸出維度在不同進程間進行分片。
MergedColumnParallelLinear 類：繼承自 ColumnParallelLinear ，支持將多個輸出維度合并處理。
QKVParallelLinear 類：繼承自 ColumnParallelLinear ，專門用于處理查詢（Q）、鍵（K）、值（V）的并行線性變換。
RowParallelLinear 類：實現了行并行的線性層，輸入維度在不同進程間進行分片，計算結果通過 all_reduce 操作匯總。

7、旋轉位置編碼層（rotary_embedding.py）

RotaryEmbedding 類：實現了旋轉位置編碼（Rotary Position Embedding, RoPE），通過正弦和余弦函數為查詢（query）和鍵（key）添加位置信息。
get_rope 函數：用于獲取 RotaryEmbedding 實例，使用 lru_cache 緩存實例，避免重復創建。

layers構建語言模型推理鏈路流程

輸入處理階段

嵌入層（ embed_head.py ）：將輸入token轉換為連續向量表示（詞嵌入），是模型處理的起點。
位置編碼（ rotary_embedding.py ）：通過旋轉位置編碼（RoPE）為嵌入向量添加位置信息，使模型能感知序列順序（如 apply_rotary_emb 函數實現）。

特征提取與變換階段

注意力層（ attention.py ）：利用帶位置信息的嵌入向量，通過多頭注意力機制計算上下文依賴（如 flash_attn_varlen_func 處理前綴緩存， flash_attn_with_kvcache 結合KV緩存優化推理）。
激活層（ activation.py ）：在注意力輸出或前饋網絡中應用SiLU激活函數（如 SiluAndMul 層將輸入分塊后執行 F.silu(x) * y ），增強模型非線性表達能力。
線性變換層（ linear.py ）：通過 ColumnParallelLinear / RowParallelLinear 等并行線性層，完成向量空間的線性變換（如QKV投影、前饋網絡映射）。
歸一化層（ layernorm.py ）：使用 RMSNorm 對各層輸出進行歸一化（如 rms_forward 方法），穩定訓練過程。

輸出生成階段

采樣器（ sampler.py ）：基于模型最終輸出的logits，通過貪心采樣或溫度調整采樣選擇下一個token（如 temperatures=0 時選最大概率token，非0時縮放logits后softmax）。

協作核心機制

各層通過 engine 模塊的 ModelRunner 和 LLMEngine 協調：

ModelRunner 初始化時為注意力層分配共享KV緩存（ k_cache / v_cache 屬性），并通過 context 對象傳遞 slot_mapping 等元數據（如 attention.py 中 context = get_context() 獲取狀態）。
LLMEngine 調度請求，根據 context.is_prefill 標志切換注意力層計算模式（前綴填充/單步解碼），確保各層在正確上下文（如內存管理、并行策略）中執行。

綜上， layers 的基礎層通過“輸入處理→特征變換→輸出生成”的流水線協作，配合 engine 的資源管理，共同實現高效的大語言模型推理。

三、models目錄

models 目錄實現具體的模型架構，會調用 layers 目錄中的基礎層來構建完整模型。以 qwen3.py 為例，實現了 Qwen3ForCausalLM 模型，其功能如下：

模型組件封裝：將 Attention 、 SiluAndMul 等基礎層封裝到 Qwen3Attention 、 Qwen3MLP 等模塊中。
模型架構搭建：通過 Qwen3DecoderLayer 和 Qwen3Model 類構建完整的 Qwen3 模型架構。

四、各目錄關聯工作方式

layers 、 models 和 engine 目錄的組件通過以下流程協同工作：

初始化階段：
- LLMEngine 初始化時創建 ModelRunner 實例， ModelRunner 初始化 Qwen3ForCausalLM 模型。
- Qwen3ForCausalLM 模型在構建過程中調用 layers 目錄下的基礎層，如 Attention 、 SiluAndMul 等。
- ModelRunner 分配 KV 緩存，并將緩存張量綁定到模型各層的 k_cache 和 v_cache 屬性上。
推理階段：
- 用戶請求通過 LLMEngine 的 add_request 方法添加到調度隊列。
- Scheduler 調度請求， LLMEngine 調用 ModelRunner 的 run 方法執行推理。
- ModelRunner 準備輸入數據，設置上下文信息，調用 Qwen3ForCausalLM 模型進行前向傳播。
- Qwen3ForCausalLM 模型在計算過程中調用 layers 目錄下的基礎層完成具體計算。
- ModelRunner 使用 Sampler 進行采樣，生成最終結果。
結果處理階段：
- LLMEngine 收集推理結果，通過 tokenizer 解碼生成文本返回給用戶。