add Chinese annotations to all source files for learning purposes

Annotated 16 source files covering the full architecture:
engine (scheduler, block manager, model runner), layers (attention,
linear, sampler, etc.), model (qwen3), and utils.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

ANNOTATIONS.md

@@ -0,0 +1,22 @@
+# Nano-vLLM 注释说明
+
+**已添加注释的文件（16个）：**
+
+| 模块 | 文件 | 注释要点 |
+|------|------|---------|
+| 入口 | `__init__.py` | 项目架构和数据流概览 |
+| 配置 | `config.py`, `sampling_params.py` | 每个参数的含义和作用 |
+| 引擎 | `sequence.py` | 序列状态机、block_table、序列化机制 |
+| 引擎 | `block_manager.py` | 前缀缓存原理、哈希链式计算、引用计数 |
+| 引擎 | `scheduler.py` | prefill/decode调度策略、chunked prefill、抢占机制 |
+| 引擎 | `model_runner.py` | KV cache分配、CUDA Graph捕获、TP共享内存通信 |
+| 引擎 | `llm_engine.py` | 引擎初始化流程、step循环、吞吐量统计 |
+| 模型 | `qwen3.py` | Qwen3架构（GQA、Q/K Norm）、融合模块映射 |
+| 层 | `attention.py` | Triton kernel写KV cache、Flash Attention两阶段 |
+| 层 | `linear.py` | 5种并行线性层（列切/行切/融合QKV/融合gate_up） |
+| 层 | `sampler.py` | Gumbel-like采样方法 |
+| 层 | `activation.py` | SwiGLU (SiLU * up) |
+| 层 | `layernorm.py` | 残差融合RMSNorm |
+| 层 | `embed_head.py` | 词表并行Embedding、LM Head前缀优化 |
+| 层 | `rotary_embedding.py` | RoPE原理和预计算缓存 |
+| 工具 | `context.py`, `loader.py` | 全局上下文机制、safetensors权重加载 |

nanovllm/__init__.py

@@ -1,2 +1,12 @@
+# nano-vllm: 一个轻量级的 vLLM 实现
+#
+# 核心架构:
+# LLM (入口) → LLMEngine (引擎) → Scheduler (调度器) + ModelRunner (模型运行器)
+#                                      ↓                      ↓
+#                               BlockManager (KV缓存管理)    Qwen3ForCausalLM (模型)
+#
+# 数据流:
+# 用户请求 → Sequence → Scheduler 调度 → ModelRunner 准备输入 → 模型前向 → 采样 → 返回结果
+
 from nanovllm.llm import LLM
 from nanovllm.sampling_params import SamplingParams

nanovllm/config.py

@@ -5,6 +5,22 @@ from transformers import AutoConfig
 
 @dataclass(slots=True)
 class Config:
+    """推理引擎的全局配置。
+
+    Args:
+        model: HuggingFace 模型的本地路径（必须是一个目录，包含 safetensors 权重文件和 config.json）。
+        max_num_batched_tokens: 单次调度（schedule）中允许的最大 token 总数，控制 prefill 阶段的批处理粒度。
+        max_num_seqs: 单次调度中允许的最大序列数，限制同时处理的请求数量。
+        max_model_len: 模型支持的最大序列长度（prompt + 生成的总长度），会被 HF config 中的 max_position_embeddings 截断。
+        gpu_memory_utilization: GPU 显存使用率（0~1），决定 KV cache 可用显存大小。
+        tensor_parallel_size: 张量并行度（TP），即使用多少张 GPU 来并行运行模型。
+        enforce_eager: 是否强制使用 eager 模式（不使用 CUDA Graph），调试时设为 True。
+        hf_config: 从模型目录自动加载的 HuggingFace 配置对象，由 __post_init__ 自动设置。
+        eos: End-of-Sequence token 的 ID，由引擎初始化时从 tokenizer 获取，用于判断生成是否结束。
+        kvcache_block_size: KV cache 的块大小（token 数），每个 block 存储这么多 token 的 KV 向量。前缀缓存和 KV cache 管理的最小单位。
+        num_kvcache_blocks: KV cache 的总块数，-1 表示尚未分配，会在 ModelRunner.allocate_kv_cache() 中根据可用显存自动计算。
+    """
+
     model: str
     max_num_batched_tokens: int = 16384
     max_num_seqs: int = 512
@@ -21,5 +37,7 @@ class Config:
         assert os.path.isdir(self.model)
         assert self.kvcache_block_size % 256 == 0
         assert 1 <= self.tensor_parallel_size <= 8
+        # 从模型目录加载 HuggingFace 配置（如 num_layers, hidden_size, num_heads 等）
         self.hf_config = AutoConfig.from_pretrained(self.model)
+        # 确保最大序列长度不超过模型支持的位置编码上限
         self.max_model_len = min(self.max_model_len, self.hf_config.max_position_embeddings)

nanovllm/engine/block_manager.py

@@ -6,6 +6,16 @@ from nanovllm.engine.sequence import Sequence
 
 
 class Block:
+    """KV cache 的物理块。
+
+    每个 Block 对应 GPU 显存中一块固定大小的 KV cache 存储，可以存放 block_size 个 token 的 KV 向量。
+
+    Attributes:
+        block_id: 物理块的唯一 ID。
+        ref_count: 引用计数，有多少个序列正在共享这个块（前缀缓存复用）。
+        hash: 该块内容的哈希值，用于前缀缓存查找。初始为 -1 表示未计算哈希。
+        token_ids: 该块对应的 token ID 列表，用于验证缓存命中时内容是否一致。
+    """
 
     def __init__(self, block_id):
         self.block_id = block_id
@@ -14,16 +24,33 @@ class Block:
         self.token_ids = []
 
     def update(self, hash: int, token_ids: list[int]):
+        """更新块的内容哈希和 token_ids。"""
         self.hash = hash
         self.token_ids = token_ids
 
     def reset(self):
+        """重置块状态，准备被重新分配。"""
         self.ref_count = 1
         self.hash = -1
         self.token_ids = []
 
 
 class BlockManager:
+    """管理 KV cache 物理块的分配、释放和前缀缓存。
+
+    核心设计:
+    - 将 KV cache 分成固定大小的物理块（block_size 个 token/块）。
+    - 使用哈希表实现前缀缓存（prefix caching）：不同序列如果拥有相同的 prompt 前缀，
+      可以共享同一组 KV cache 块，避免重复计算。
+    - 块通过引用计数（ref_count）管理生命周期，所有引用释放后块变为空闲可复用。
+
+    Attributes:
+        block_size: 每个块存储的 token 数量。
+        blocks: 所有物理块的列表，索引即 block_id。
+        hash_to_block_id: 内容哈希 → 块 ID 的映射，用于前缀缓存查找。
+        free_block_ids: 空闲块的 ID 队列。
+        used_block_ids: 正在使用的块 ID 集合。
+    """
 
     def __init__(self, num_blocks: int, block_size: int):
         self.block_size = block_size
@@ -34,6 +61,15 @@ class BlockManager:
 
     @classmethod
     def compute_hash(cls, token_ids: list[int], prefix: int = -1):
+        """计算一个块的内容哈希值。
+
+        哈希是链式的：每个块的哈希值依赖于前一个块的哈希值（prefix 参数），
+        这确保了前缀缓存的一致性——只有完全相同的 token 序列前缀才会产生相同的哈希链。
+
+        Args:
+            token_ids: 该块对应的 token ID 列表。
+            prefix: 前一个块的哈希值，-1 表示第一个块（无前缀）。
+        """
         h = xxhash.xxh64()
         if prefix != -1:
             h.update(prefix.to_bytes(8, "little"))
@@ -41,9 +77,11 @@ class BlockManager:
         return h.intdigest()
 
     def _allocate_block(self) -> int:
+        """从空闲池中分配一个物理块。"""
         block_id = self.free_block_ids.popleft()
         block = self.blocks[block_id]
         assert block.ref_count == 0
+        # 如果该块之前有哈希记录，先从哈希表中移除
         if block.hash != -1 and self.hash_to_block_id.get(block.hash) == block_id:
             del self.hash_to_block_id[block.hash]
         block.reset()
@@ -51,11 +89,22 @@ class BlockManager:
         return block_id
 
     def _deallocate_block(self, block_id: int):
+        """释放一个物理块回空闲池。"""
         assert self.blocks[block_id].ref_count == 0
         self.used_block_ids.remove(block_id)
         self.free_block_ids.append(block_id)
 
     def can_allocate(self, seq: Sequence) -> int:
+        """检查是否有足够的空闲块来分配给序列，同时计算前缀缓存命中数。
+
+        遍历序列的所有逻辑块，逐个计算哈希并与已有缓存比对：
+        - 如果哈希匹配且 token 内容一致，则该块可以复用（前缀缓存命中）。
+        - 一旦遇到不匹配的块就停止，因为前缀缓存要求连续匹配。
+        - 最后检查剩余空闲块是否足够分配未命中部分。
+
+        Returns:
+            前缀缓存命中的块数，如果没有足够的空闲块则返回 -1。
+        """
         h = -1
         num_cached_blocks = 0
         num_new_blocks = seq.num_blocks
@@ -67,14 +116,25 @@ class BlockManager:
                 break
             num_cached_blocks += 1
             if block_id in self.used_block_ids:
+                # 块已在使用中（被其他序列共享），不需要额外分配
                 num_new_blocks -= 1
         if len(self.free_block_ids) < num_new_blocks:
             return -1
         return num_cached_blocks
 
     def allocate(self, seq: Sequence, num_cached_blocks: int):
+        """为序列分配 KV cache 物理块。
+
+        前缀缓存命中部分：增加已有块的引用计数。
+        未命中部分：从空闲池分配新块。
+
+        Args:
+            seq: 要分配块的序列。
+            num_cached_blocks: 前缀缓存命中的块数（由 can_allocate 计算）。
+        """
         assert not seq.block_table
         h = -1
+        # 处理缓存命中部分：复用已有块
         for i in range(num_cached_blocks):
             token_ids = seq.block(i)
             h = self.compute_hash(token_ids, h)
@@ -83,15 +143,22 @@ class BlockManager:
             if block_id in self.used_block_ids:
                 block.ref_count += 1
             else:
+                # 块在空闲池中但哈希匹配（之前被释放但内容未覆盖），重新激活
                 block.ref_count = 1
                 self.free_block_ids.remove(block_id)
                 self.used_block_ids.add(block_id)
             seq.block_table.append(block_id)
+        # 分配新块用于未命中部分
         for i in range(num_cached_blocks, seq.num_blocks):
             seq.block_table.append(self._allocate_block())
         seq.num_cached_tokens = num_cached_blocks * self.block_size
 
     def deallocate(self, seq: Sequence):
+        """释放序列占用的所有 KV cache 块。
+
+        递减每个块的引用计数，引用计数归零的块被回收进空闲池。
+        注意：遍历顺序为逆序，这是为了优先释放最新分配的块（它们在空闲队列尾部）。
+        """
         for block_id in reversed(seq.block_table):
             block = self.blocks[block_id]
             block.ref_count -= 1
@@ -101,16 +168,30 @@ class BlockManager:
         seq.block_table.clear()
 
     def can_append(self, seq: Sequence) -> bool:
+        """检查 decode 阶段是否能追加一个 token。
+
+        当序列当前最后一个块已满时（len(seq) % block_size == 1，即新增 token 是新块的第一个），
+        需要分配一个新块。否则只需要写入已有块，不需要额外分配。
+        """
         return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)
 
     def may_append(self, seq: Sequence):
+        """在 decode 阶段，如果需要则追加一个新块。"""
         if len(seq) % self.block_size == 1:
             seq.block_table.append(self._allocate_block())
 
     def hash_blocks(self, seq: Sequence):
+        """更新序列中已完成计算的块的哈希值。
+
+        在每次调度步骤完成后调用，将新计算完毕的块的 token 内容和哈希值
+        注册到哈希表中，以便后续序列复用（前缀缓存）。
+
+        只处理从 num_cached_tokens 到当前进度的块（即本次新完成的块）。
+        """
         start = seq.num_cached_tokens // self.block_size
         end = (seq.num_cached_tokens + seq.num_scheduled_tokens) // self.block_size
         if start == end: return
+        # 继承前一个块的哈希值作为前缀
         h = self.blocks[seq.block_table[start - 1]].hash if start > 0 else -1
         for i in range(start, end):
             block = self.blocks[seq.block_table[i]]

nanovllm/engine/llm_engine.py

@@ -13,40 +13,77 @@ from nanovllm.engine.model_runner import ModelRunner
 
 
 class LLMEngine:
+    """LLM 推理引擎：整个 nano-vllm 的入口和协调者。
+
+    职责:
+    1. 初始化配置、启动张量并行的子进程。
+    2. 管理请求（添加/调度/生成）。
+    3. 协调调度器和模型运行器之间的交互。
+
+    使用方式:
+    >>> engine = LLMEngine("/path/to/model", enforce_eager=True)
+    >>> outputs = engine.generate(["Hello"], SamplingParams(max_tokens=64))
+    """
 
     def __init__(self, model, **kwargs):
+        # 过滤出 Config 支持的参数，忽略其他参数
         config_fields = {field.name for field in fields(Config)}
         config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields}
         config = Config(model, **config_kwargs)
         Sequence.block_size = config.kvcache_block_size
+
+        # 启动张量并行的工作进程（rank 1, 2, ...）
+        # 使用 "spawn" 方式创建子进程，确保 CUDA context 独立
         self.ps = []
         self.events = []
         ctx = mp.get_context("spawn")
         for i in range(1, config.tensor_parallel_size):
-            event = ctx.Event()
+            event = ctx.Event()  # 用于通知工作进程有新任务
             process = ctx.Process(target=ModelRunner, args=(config, i, event))
             process.start()
             self.ps.append(process)
             self.events.append(event)
+
+        # 在主进程中初始化模型运行器（rank 0）
         self.model_runner = ModelRunner(config, 0, self.events)
+        # 加载 tokenizer（用于文本 ↔ token ID 转换）
         self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True)
         config.eos = self.tokenizer.eos_token_id
+        # 初始化调度器（此时 num_kvcache_blocks 已由 ModelRunner 分配完毕）
         self.scheduler = Scheduler(config)
+        # 注册退出清理函数
         atexit.register(self.exit)
 
     def exit(self):
+        """通知所有进程退出并等待。"""
         self.model_runner.call("exit")
         del self.model_runner
         for p in self.ps:
             p.join()
 
     def add_request(self, prompt: str | list[int], sampling_params: SamplingParams):
+        """添加一个生成请求到调度器的等待队列。
+
+        Args:
+            prompt: 可以是字符串（会被 tokenizer 编码为 token ID）或已编码的 token ID 列表。
+            sampling_params: 采样参数。
+        """
         if isinstance(prompt, str):
             prompt = self.tokenizer.encode(prompt)
         seq = Sequence(prompt, sampling_params)
         self.scheduler.add(seq)
 
     def step(self):
+        """执行一个推理步骤。
+
+        1. 调度器选择要处理的序列（prefill 或 decode）。
+        2. 模型运行器执行前向推理和采样。
+        3. 调度器后处理结果（更新缓存、检查终止条件）。
+
+        Returns:
+            outputs: 已完成序列的 (seq_id, completion_token_ids) 列表。
+            num_tokens: prefill 时为处理的 token 总数（正数），decode 时为序列数的负数。
+        """
         seqs, is_prefill = self.scheduler.schedule()
         num_tokens = sum(seq.num_scheduled_tokens for seq in seqs) if is_prefill else -len(seqs)
         token_ids = self.model_runner.call("run", seqs, is_prefill)
@@ -55,6 +92,7 @@ class LLMEngine:
         return outputs, num_tokens
 
     def is_finished(self):
+        """检查是否所有请求都已完成。"""
         return self.scheduler.is_finished()
 
     def generate(
@@ -63,6 +101,21 @@ class LLMEngine:
         sampling_params: SamplingParams | list[SamplingParams],
         use_tqdm: bool = True,
     ) -> list[str]:
+        """批量生成文本的入口方法。
+
+        工作流程:
+        1. 将所有 prompt 添加为请求。
+        2. 循环执行 step() 直到所有序列完成。
+        3. 收集结果并解码为文本。
+
+        Args:
+            prompts: prompt 列表，每个元素可以是字符串或 token ID 列表。
+            sampling_params: 单个采样参数（应用于所有 prompt）或每个 prompt 对应的采样参数列表。
+            use_tqdm: 是否显示进度条。
+
+        Returns:
+            列表，每个元素是 {"text": 解码文本, "token_ids": token ID 列表}。
+        """
         pbar = tqdm(total=len(prompts), desc="Generating", dynamic_ncols=True, disable=not use_tqdm)
         if not isinstance(sampling_params, list):
             sampling_params = [sampling_params] * len(prompts)
@@ -73,6 +126,7 @@ class LLMEngine:
         while not self.is_finished():
             t = perf_counter()
             output, num_tokens = self.step()
+            # num_tokens > 0 表示 prefill，< 0 表示 decode（取负得到序列数）
             if num_tokens > 0:
                 prefill_throughput = num_tokens / (perf_counter() - t)
             else:
@@ -85,6 +139,7 @@ class LLMEngine:
                 outputs[seq_id] = token_ids
                 pbar.update(1)
         pbar.close()
+        # 按 seq_id 排序输出（保证与输入顺序一致）
         outputs = [outputs[seq_id] for seq_id in sorted(outputs.keys())]
         outputs = [{"text": self.tokenizer.decode(token_ids), "token_ids": token_ids} for token_ids in outputs]
         return outputs

nanovllm/engine/model_runner.py

@@ -13,6 +13,18 @@ from nanovllm.utils.loader import load_model
 
 
 class ModelRunner:
+    """模型运行器：负责模型推理、KV cache 管理、CUDA Graph 捕获和张量并行通信。
+
+    在张量并行（TP）模式下:
+    - Rank 0 是主进程，负责采样和与引擎通信。
+    - Rank > 0 是工作进程，通过共享内存（SharedMemory）接收指令。
+    - 所有进程共享同一个模型和 KV cache 的分片。
+
+    生命周期:
+    1. 初始化: 加载模型 → warmup → 分配 KV cache → (可选)捕获 CUDA Graph
+    2. 推理: 接收序列 → 准备输入 → 运行模型 → 采样 token
+    3. 退出: 释放资源
+    """
 
     def __init__(self, config: Config, rank: int, event: Event | list[Event]):
         self.config = config
@@ -23,42 +35,54 @@ class ModelRunner:
         self.rank = rank
         self.event = event
 
+        # 初始化分布式进程组（NCCL 后端），所有 GPU 通过 TCP 通信
         dist.init_process_group("nccl", "tcp://localhost:2333", world_size=self.world_size, rank=rank)
         torch.cuda.set_device(rank)
+
+        # 加载模型权重
         default_dtype = torch.get_default_dtype()
         torch.set_default_dtype(hf_config.dtype)
         torch.set_default_device("cuda")
         self.model = Qwen3ForCausalLM(hf_config)
         load_model(self.model, config.model)
         self.sampler = Sampler()
+
+        # Warmup: 运行一次前向传播以确定模型本身的显存占用
         self.warmup_model()
+        # 根据剩余显存分配 KV cache
         self.allocate_kv_cache()
+        # 捕获 CUDA Graph 以加速 decode 阶段的小批量推理
         if not self.enforce_eager:
             self.capture_cudagraph()
+
         torch.set_default_device("cpu")
         torch.set_default_dtype(default_dtype)
 
+        # 张量并行时，rank > 0 的工作进程进入消息循环
         if self.world_size > 1:
             if rank == 0:
+                # 主进程创建共享内存，工作进程打开它
                 self.shm = SharedMemory(name="nanovllm", create=True, size=2**20)
                 dist.barrier()
             else:
                 dist.barrier()
                 self.shm = SharedMemory(name="nanovllm")
-                self.loop()
+                self.loop()  # 工作进程在此循环，直到收到 exit 指令
 
     def exit(self):
+        """释放所有资源并退出。"""
         if self.world_size > 1:
             self.shm.close()
             dist.barrier()
             if self.rank == 0:
-                self.shm.unlink()
+                self.shm.unlink()  # 只有创建者需要 unlink
         if not self.enforce_eager:
             del self.graphs, self.graph_pool
         torch.cuda.synchronize()
         dist.destroy_process_group()
 
     def loop(self):
+        """工作进程的主循环：等待主进程指令，执行对应方法。"""
         while True:
             method_name, args = self.read_shm()
             self.call(method_name, *args)
@@ -66,14 +90,16 @@ class ModelRunner:
                 break
 
     def read_shm(self):
+        """从共享内存读取主进程发送的方法调用指令。"""
         assert self.world_size > 1 and self.rank > 0
-        self.event.wait()
+        self.event.wait()  # 等待主进程通知
         n = int.from_bytes(self.shm.buf[0:4], "little")
         method_name, *args = pickle.loads(self.shm.buf[4:n+4])
         self.event.clear()
         return method_name, args
 
     def write_shm(self, method_name, *args):
+        """将方法调用指令写入共享内存，通知工作进程。"""
         assert self.world_size > 1 and self.rank == 0
         data = pickle.dumps([method_name, *args])
         n = len(data)
@@ -83,12 +109,19 @@ class ModelRunner:
             event.set()
 
     def call(self, method_name, *args):
+        """调用指定方法。TP 模式下主进程先通知工作进程，再本地执行。"""
         if self.world_size > 1 and self.rank == 0:
             self.write_shm(method_name, *args)
         method = getattr(self, method_name, None)
         return method(*args)
 
     def warmup_model(self):
+        """预热模型：运行一次最大批量的前向传播。
+
+        目的是让 PyTorch 分配所有内部缓存（cuBLAS workspace 等），
+        然后通过 empty_cache 释放临时显存，这样后续的 peak memory 统计
+        就只包含模型权重，从而准确计算 KV cache 可用空间。
+        """
         torch.cuda.empty_cache()
         torch.cuda.reset_peak_memory_stats()
         max_num_batched_tokens, max_model_len = self.config.max_num_batched_tokens, self.config.max_model_len
@@ -101,6 +134,15 @@ class ModelRunner:
         torch.cuda.empty_cache()
 
     def allocate_kv_cache(self):
+        """根据剩余 GPU 显存分配 KV cache。
+
+        计算公式:
+        可用显存 = 总显存 × gpu_memory_utilization - 非模型占用
+        其中非模型占用 = 已用显存 - peak（模型权重）+ current（当前模型张量）
+
+        KV cache 形状: (2, num_layers, num_blocks, block_size, num_kv_heads/head_dim)
+        其中第一维 2 分别对应 K 和 V cache。
+        """
         config = self.config
         hf_config = config.hf_config
         free, total = torch.cuda.mem_get_info()
@@ -109,10 +151,13 @@ class ModelRunner:
         current = torch.cuda.memory_stats()["allocated_bytes.all.current"]
         num_kv_heads = hf_config.num_key_value_heads // self.world_size
         head_dim = getattr(hf_config, "head_dim", hf_config.hidden_size // hf_config.num_attention_heads)
+        # 每个块占用的字节数：2(K+V) × 层数 × block_size × KV头数 × head_dim × dtype字节数
         block_bytes = 2 * hf_config.num_hidden_layers * self.block_size * num_kv_heads * head_dim * hf_config.dtype.itemsize
         config.num_kvcache_blocks = int(total * config.gpu_memory_utilization - used - peak + current) // block_bytes
         assert config.num_kvcache_blocks > 0
+        # 分配 KV cache 张量，形状为 (2, num_layers, num_blocks, block_size, num_kv_heads, head_dim)
         self.kv_cache = torch.empty(2, hf_config.num_hidden_layers, config.num_kvcache_blocks, self.block_size, num_kv_heads, head_dim)
+        # 将 KV cache 的视图绑定到模型中每个 Attention 层
         layer_id = 0
         for module in self.model.modules():
             if hasattr(module, "k_cache") and hasattr(module, "v_cache"):
@@ -121,12 +166,26 @@ class ModelRunner:
                 layer_id += 1
 
     def prepare_block_tables(self, seqs: list[Sequence]):
+        """将序列的 block_table 列表填充为等长的二维张量，用于 GPU 计算。"""
         max_len = max(len(seq.block_table) for seq in seqs)
+        # 用 -1 填充短序列的 block_table
         block_tables = [seq.block_table + [-1] * (max_len - len(seq.block_table)) for seq in seqs]
         block_tables = torch.tensor(block_tables, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
         return block_tables
 
     def prepare_prefill(self, seqs: list[Sequence]):
+        """准备 prefill 阶段的模型输入张量。
+
+        Prefill 阶段需要处理多个序列的 prompt tokens，所有 token 被拼接成一个连续的输入。
+        使用 cu_seqlens（累积序列长度）来标记每个序列的边界，供 flash_attn_varlen 使用。
+
+        关键数据:
+        - input_ids: 所有序列的 token ID 拼接。
+        - positions: 每个 token 的位置 ID（考虑前缀缓存偏移）。
+        - cu_seqlens_q/k: 查询和键值的累积序列长度。
+        - slot_mapping: 将每个 token 映射到 KV cache 中的物理存储位置。
+        - block_tables: 前缀缓存命中时需要 block_table 来从 KV cache 读取已缓存的 K/V。
+        """
         input_ids = []
         positions = []
         cu_seqlens_q = [0]
@@ -136,18 +195,19 @@ class ModelRunner:
         slot_mapping = []
         block_tables = None
         for seq in seqs:
-            start = seq.num_cached_tokens
+            start = seq.num_cached_tokens  # 跳过已缓存的 token
             seqlen_q = seq.num_scheduled_tokens
             end = start + seqlen_q
-            seqlen_k = end
+            seqlen_k = end  # KV 的长度是从 0 到 end（包括缓存前缀）
             input_ids.extend(seq[start:end])
             positions.extend(range(start, end))
             cu_seqlens_q.append(cu_seqlens_q[-1] + seqlen_q)
             cu_seqlens_k.append(cu_seqlens_k[-1] + seqlen_k)
             max_seqlen_q = max(seqlen_q, max_seqlen_q)
             max_seqlen_k = max(seqlen_k, max_seqlen_k)
-            if not seq.block_table:    # warmup
+            if not seq.block_table:    # warmup 阶段没有 block_table
                 continue
+            # 计算 slot_mapping：每个 token 对应 KV cache 中的哪个 slot
             start_block = start // self.block_size
             end_block = (end + self.block_size - 1) // self.block_size
             for i in range(start_block, end_block):
@@ -159,7 +219,7 @@ class ModelRunner:
                 else:
                     slot_end = seq.block_table[i] * self.block_size + end - i * self.block_size
                 slot_mapping.extend(range(slot_start, slot_end))
-        if cu_seqlens_k[-1] > cu_seqlens_q[-1]:    # prefix cache
+        if cu_seqlens_k[-1] > cu_seqlens_q[-1]:    # 前缀缓存命中时，KV 长度 > Q 长度
             block_tables = self.prepare_block_tables(seqs)
         input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
         positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
@@ -170,6 +230,18 @@ class ModelRunner:
         return input_ids, positions
 
     def prepare_decode(self, seqs: list[Sequence]):
+        """准备 decode 阶段的模型输入张量。
+
+        Decode 阶段每个序列只处理 1 个 token（最新生成的 token）。
+        模型从 KV cache 中读取之前所有的 K/V 向量来做注意力计算。
+
+        关键数据:
+        - input_ids: 每个序列的最新 token ID。
+        - positions: 每个 token 的位置 ID（序列长度 - 1）。
+        - slot_mapping: 新 token 的 KV 写入位置。
+        - context_lens: 每个序列的上下文总长度。
+        - block_tables: KV cache 块映射表。
+        """
         input_ids = []
         positions = []
         slot_mapping = []
@@ -178,6 +250,7 @@ class ModelRunner:
             input_ids.append(seq.last_token)
             positions.append(len(seq) - 1)
             context_lens.append(len(seq))
+            # slot = 最后一个块的起始位置 + 该块内已有 token 数 - 1
             slot_mapping.append(seq.block_table[-1] * self.block_size + seq.last_block_num_tokens  - 1)
         input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
         positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
@@ -188,19 +261,30 @@ class ModelRunner:
         return input_ids, positions
 
     def prepare_sample(self, seqs: list[Sequence]):
+        """准备采样所需的温度参数张量。"""
         temperatures = [seq.temperature for seq in seqs]
         temperatures = torch.tensor(temperatures, dtype=torch.float32, pin_memory=True).cuda(non_blocking=True)
         return temperatures
 
     @torch.inference_mode()
     def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill: bool):
+        """运行模型前向传播。
+
+        对于 decode 阶段的小批量（<=512），使用 CUDA Graph 加速：
+        CUDA Graph 将整个计算图"录制"下来，后续只需回放即可，避免了
+        CPU 端的 kernel launch 开销，对 decode（每个 step 计算量很小）尤为有效。
+        """
         if is_prefill or self.enforce_eager or input_ids.size(0) > 512:
+            # 直接运行：prefill（批量动态）、eager 模式、或大批量 decode
             return self.model.compute_logits(self.model(input_ids, positions))
         else:
+            # 使用 CUDA Graph 回放加速小批量 decode
             bs = input_ids.size(0)
             context = get_context()
+            # 选择 >= bs 的最小预捕获图大小
             graph = self.graphs[next(x for x in self.graph_bs if x >= bs)]
             graph_vars = self.graph_vars
+            # 将实际输入拷贝到图预分配的固定大小缓冲区中
             graph_vars["input_ids"][:bs] = input_ids
             graph_vars["positions"][:bs] = positions
             graph_vars["slot_mapping"].fill_(-1)
@@ -208,10 +292,16 @@ class ModelRunner:
             graph_vars["context_lens"].zero_()
             graph_vars["context_lens"][:bs] = context.context_lens
             graph_vars["block_tables"][:bs, :context.block_tables.size(1)] = context.block_tables
+            # 回放图（比重新执行快，跳过了 Python/PyTorch 调度开销）
             graph.replay()
             return self.model.compute_logits(graph_vars["outputs"][:bs])
 
     def run(self, seqs: list[Sequence], is_prefill: bool) -> list[int]:
+        """执行一次完整的推理步骤：准备输入 → 模型前向 → 采样。
+
+        Returns:
+            采样得到的 token ID 列表（仅 rank 0 返回有效值）。
+        """
         input_ids, positions = self.prepare_prefill(seqs) if is_prefill else self.prepare_decode(seqs)
         temperatures = self.prepare_sample(seqs) if self.rank == 0 else None
         logits = self.run_model(input_ids, positions, is_prefill)
@@ -221,28 +311,39 @@ class ModelRunner:
 
     @torch.inference_mode()
     def capture_cudagraph(self):
+        """预捕获不同批量大小的 CUDA Graph。
+
+        CUDA Graph 要求输入张量的地址不变（同一个内存池），所以需要预分配
+        固定大小的输入缓冲区，并为每个 batch size 录制一个图。
+
+        预捕获的 batch size: [1, 2, 4, 8, 16, 32, ..., max_bs]
+        运行时选择 >= 实际 batch size 的最小预捕获图。
+        """
         config = self.config
         hf_config = config.hf_config
         max_bs = min(self.config.max_num_seqs, 512)
         max_num_blocks = (config.max_model_len + self.block_size - 1) // self.block_size
+        # 预分配固定地址的输入/输出缓冲区
         input_ids = torch.zeros(max_bs, dtype=torch.int64)
         positions = torch.zeros(max_bs, dtype=torch.int64)
         slot_mapping = torch.zeros(max_bs, dtype=torch.int32)
         context_lens = torch.zeros(max_bs, dtype=torch.int32)
         block_tables = torch.zeros(max_bs, max_num_blocks, dtype=torch.int32)
         outputs = torch.zeros(max_bs, hf_config.hidden_size)
+        # 要捕获的 batch size 列表
         self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
         self.graphs = {}
         self.graph_pool = None
 
+        # 逆序捕获：先捕获大的 batch size，共享同一个 graph pool
         for bs in reversed(self.graph_bs):
             graph = torch.cuda.CUDAGraph()
             set_context(False, slot_mapping=slot_mapping[:bs], context_lens=context_lens[:bs], block_tables=block_tables[:bs])
-            outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # warmup
+            outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # warmup 运行
             with torch.cuda.graph(graph, self.graph_pool):
-                outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # capture
+                outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # 捕获计算图
             if self.graph_pool is None:
-                self.graph_pool = graph.pool()
+                self.graph_pool = graph.pool()  # 所有图共享同一个内存池
             self.graphs[bs] = graph
             torch.cuda.synchronize()
             reset_context()

nanovllm/engine/scheduler.py

@@ -6,6 +6,22 @@ from nanovllm.engine.block_manager import BlockManager
 
 
 class Scheduler:
+    """调度器：决定每个步骤（step）中哪些序列被处理以及处理多少 token。
+
+    调度策略采用 vLLM 风格的 prefill-decode 分离调度：
+    1. **Prefill 优先**: 每次调度先尝试处理等待中的序列（计算其 prompt 的 KV cache）。
+    2. **Chunked prefill**: 如果一个序列的 prompt 太长，可以分多次调度处理。
+    3. **Decode**: 当没有等待中的序列时（或 prefill token 额度用完），处理正在解码的序列。
+    4. **抢占（Preemption）**: 当 KV cache 空间不足时，将最近的 running 序列抢占回 waiting 队列。
+
+    调度约束:
+    - 总 token 数不超过 max_num_batched_tokens（prefill 阶段）。
+    - 总序列数不超过 max_num_seqs。
+
+    Attributes:
+        waiting: 等待处理的序列队列（FIFO）。
+        running: 正在解码的序列队列。
+    """
 
     def __init__(self, config: Config):
         self.max_num_seqs = config.max_num_seqs
@@ -17,50 +33,79 @@ class Scheduler:
         self.running: deque[Sequence] = deque()
 
     def is_finished(self):
+        """检查是否所有序列都已完成。"""
         return not self.waiting and not self.running
 
     def add(self, seq: Sequence):
+        """将一个新序列加入等待队列。"""
         self.waiting.append(seq)
 
     def schedule(self) -> tuple[list[Sequence], bool]:
+        """执行一次调度，返回 (被调度的序列列表, 是否为 prefill 阶段)。
+
+        调度逻辑:
+        1. Prefill 阶段：从 waiting 队列中选取序列，检查前缀缓存命中情况，
+           为每个序列计算需要处理的 token 数量。支持 chunked prefill（长 prompt 分多次处理）。
+        2. Decode 阶段：从 running 队列中选取序列，每个序列处理 1 个 token。
+           如果 KV cache 空间不足，会抢占（preempt）最近加入 running 的序列。
+        """
         scheduled_seqs = []
         num_batched_tokens = 0
 
-        # prefill
+        # ========== Prefill 阶段 ==========
+        # 尝试从 waiting 队列中调度序列，计算它们的 prompt KV cache
         while self.waiting and len(scheduled_seqs) < self.max_num_seqs:
             seq = self.waiting[0]
             remaining = self.max_num_batched_tokens - num_batched_tokens
             if remaining == 0:
                 break
+
             if not seq.block_table:
+                # 序列尚未分配块，检查前缀缓存和空闲块
                 num_cached_blocks = self.block_manager.can_allocate(seq)
                 if num_cached_blocks == -1:
+                    # 空闲块不足，停止调度
                     break
+                # 需要实际处理的 token 数 = 总 prompt token 数 - 缓存命中的 token 数
                 num_tokens = seq.num_tokens - num_cached_blocks * self.block_size
             else:
+                # 序列已经有块（chunked prefill 的后续分片），只需处理未缓存的 token
                 num_tokens = seq.num_tokens - seq.num_cached_tokens
-            if remaining < num_tokens and scheduled_seqs:  # only allow chunked prefill for the first seq
+
+            if remaining < num_tokens and scheduled_seqs:
+                # token 预算不足以处理整个序列，且已有其他序列被调度
+                # 注意：第一个序列允许 chunked prefill（remaining < num_tokens 也可以）
                 break
+
             if not seq.block_table:
                 self.block_manager.allocate(seq, num_cached_blocks)
+
+            # 实际调度的 token 数取 min(num_tokens, remaining)，实现 chunked prefill
             seq.num_scheduled_tokens = min(num_tokens, remaining)
             num_batched_tokens += seq.num_scheduled_tokens
+
             if seq.num_cached_tokens + seq.num_scheduled_tokens == seq.num_tokens:
+                # 整个 prompt 已全部处理完毕，转移到 running 队列
                 seq.status = SequenceStatus.RUNNING
                 self.waiting.popleft()
                 self.running.append(seq)
+
             scheduled_seqs.append(seq)
 
         if scheduled_seqs:
-            return scheduled_seqs, True
+            return scheduled_seqs, True  # is_prefill = True
 
-        # decode
+        # ========== Decode 阶段 ==========
+        # 逐 token 解码，每个序列每次生成 1 个 token
         while self.running and len(scheduled_seqs) < self.max_num_seqs:
             seq = self.running.popleft()
+            # 检查是否有空闲块用于存储新的 KV cache
             while not self.block_manager.can_append(seq):
                 if self.running:
+                    # 空间不足，抢占最近加入 running 的序列
                     self.preempt(self.running.pop())
                 else:
+                    # 连当前序列都要被抢占
                     self.preempt(seq)
                     break
             else:
@@ -68,21 +113,37 @@ class Scheduler:
                 seq.is_prefill = False
                 self.block_manager.may_append(seq)
                 scheduled_seqs.append(seq)
-        assert scheduled_seqs
+
+        assert scheduled_seqs, "No sequences to schedule"
+        # 将调度过的序列放回 running 队列前端（保持顺序）
         self.running.extendleft(reversed(scheduled_seqs))
-        return scheduled_seqs, False
+        return scheduled_seqs, False  # is_prefill = False
 
     def preempt(self, seq: Sequence):
+        """抢占一个序列：释放其 KV cache 并放回等待队列头部。
+
+        抢占后序列需要重新做 prefill（重新计算 KV cache）。
+        这是一种牺牲吞吐量来换取 KV cache 空间的策略。
+        """
         seq.status = SequenceStatus.WAITING
         seq.is_prefill = True
         self.block_manager.deallocate(seq)
         self.waiting.appendleft(seq)
 
     def postprocess(self, seqs: list[Sequence], token_ids: list[int], is_prefill: bool):
+        """在模型推理完成后处理每个序列的结果。
+
+        1. 更新块的哈希值（用于前缀缓存）。
+        2. 更新已缓存的 token 计数。
+        3. 对于 prefill：如果整个 prompt 还没处理完，继续等待下一次调度。
+        4. 对于完成的序列（prefill 结束后或 decode 中）：追加生成的 token。
+        5. 检查终止条件（EOS 或达到 max_tokens），完成的序列被释放资源。
+        """
         for seq, token_id in zip(seqs, token_ids):
             self.block_manager.hash_blocks(seq)
             seq.num_cached_tokens += seq.num_scheduled_tokens
             seq.num_scheduled_tokens = 0
+            # prefill 阶段如果还没处理完整个 prompt，不采样 token，继续等待
             if is_prefill and seq.num_cached_tokens < seq.num_tokens:
                 continue
             seq.append_token(token_id)

nanovllm/engine/sequence.py

@@ -6,12 +6,35 @@ from nanovllm.sampling_params import SamplingParams
 
 
 class SequenceStatus(Enum):
+    """序列的生命周期状态。
+
+    WAITING:  等待被调度器选中进入 prefill 阶段。
+    RUNNING:  正在解码（decode）阶段，逐 token 生成。
+    FINISHED: 生成完成（遇到 EOS 或达到 max_tokens）。
+    """
     WAITING = auto()
     RUNNING = auto()
     FINISHED = auto()
 
 
 class Sequence:
+    """表示一个推理请求的序列。
+
+    每个 Sequence 封装了一个完整的生成请求，从 prompt tokens 到生成的 completion tokens。
+    它是调度器、块管理器和模型运行器之间传递的核心数据结构。
+
+    关键概念:
+    - prompt tokens: 用户输入的原始 token 序列。
+    - completion tokens: 模型生成的 token 序列。
+    - cached tokens: 已经计算过 KV cache 的 token 数量（用于前缀缓存和 chunked prefill）。
+    - scheduled tokens: 当前调度步骤中计划处理的 token 数量。
+    - block_table: 该序列在 KV cache 中占用的物理块 ID 列表。
+
+    类属性:
+        block_size: KV cache 块大小，由 Config 在引擎初始化时设置。
+        counter: 全局自增计数器，用于为每个序列分配唯一 ID。
+    """
+
     block_size = 256
     counter = count()
 
@@ -20,12 +43,12 @@ class Sequence:
         self.status = SequenceStatus.WAITING
         self.token_ids = copy(token_ids)
         self.last_token = token_ids[-1]
-        self.num_tokens = len(self.token_ids)
+        self.num_tokens = len(self.token_ids)  # 当前总 token 数（prompt + 已生成的）
-        self.num_prompt_tokens = len(token_ids)
+        self.num_prompt_tokens = len(token_ids)  # prompt 部分的 token 数，生成过程中不变
-        self.num_cached_tokens = 0
+        self.num_cached_tokens = 0  # 已计算 KV cache 的 token 数，用于前缀缓存命中判断
-        self.num_scheduled_tokens = 0
+        self.num_scheduled_tokens = 0  # 当前步骤中被调度处理的 token 数
-        self.is_prefill = True
+        self.is_prefill = True  # 是否处于 prefill 阶段（首次计算 prompt 的 KV cache）
-        self.block_table = []
+        self.block_table = []  # KV cache 物理块 ID 列表，索引为逻辑块号
         self.temperature = sampling_params.temperature
         self.max_tokens = sampling_params.max_tokens
         self.ignore_eos = sampling_params.ignore_eos
@@ -42,42 +65,56 @@ class Sequence:
 
     @property
     def num_completion_tokens(self):
+        """已生成的 completion token 数量。"""
         return self.num_tokens - self.num_prompt_tokens
 
     @property
     def prompt_token_ids(self):
+        """返回 prompt 部分的 token ID 列表。"""
         return self.token_ids[:self.num_prompt_tokens]
 
     @property
     def completion_token_ids(self):
+        """返回已生成的 completion token ID 列表。"""
         return self.token_ids[self.num_prompt_tokens:]
 
     @property
     def num_blocks(self):
+        """该序列需要的 KV cache 逻辑块数量（向上取整）。"""
         return (self.num_tokens + self.block_size - 1) // self.block_size
 
     @property
     def last_block_num_tokens(self):
+        """最后一个块中已使用的 token 数量。"""
         return self.num_tokens - (self.num_blocks - 1) * self.block_size
 
     def block(self, i):
+        """获取第 i 个逻辑块对应的 token ID 列表。"""
         assert 0 <= i < self.num_blocks
         return self.token_ids[i*self.block_size: (i+1)*self.block_size]
 
     def append_token(self, token_id: int):
+        """将一个新生成的 token 追加到序列末尾。"""
         self.token_ids.append(token_id)
         self.last_token = token_id
         self.num_tokens += 1
 
     def __getstate__(self):
+        """序列化时只保存必要的状态，用于多进程间传递序列数据。
+
+        prefill 阶段保存完整 token_ids（模型需要读取全部 prompt tokens），
+        decode 阶段只保存 last_token（模型只需要最新一个 token 的 ID）。
+        """
         last_state = self.last_token if not self.is_prefill else self.token_ids
         return (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.num_scheduled_tokens, self.block_table, last_state)
 
     def __setstate__(self, state):
+        """反序列化，恢复序列状态。"""
         self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.num_scheduled_tokens, self.block_table, last_state = state
         if isinstance(last_state, list):
             self.token_ids = last_state
             self.last_token = self.token_ids[-1]
         else:
+            # decode 阶段不需要完整的 token_ids，只保存了 last_token
             self.token_ids = []
             self.last_token = last_state

nanovllm/layers/activation.py

@@ -4,6 +4,12 @@ import torch.nn.functional as F
 
 
 class SiluAndMul(nn.Module):
+    """SwiGLU 激活函数：SiLU(gate) * up。
+
+    输入是 gate 和 up 拼接的张量，沿最后一维一分为二，
+    对前半部分应用 SiLU 激活后与后半部分逐元素相乘。
+    这是 LLaMA/Qwen 系列模型中 MLP 层的标准激活函数。
+    """
 
     @torch.compile
     def forward(self, x: torch.Tensor) -> torch.Tensor:

nanovllm/layers/attention.py

@@ -18,6 +18,12 @@ def store_kvcache_kernel(
     slot_mapping_ptr,
     D: tl.constexpr,
 ):
+    """Triton kernel：将新计算的 K/V 向量写入 KV cache。
+
+    每个 Triton program 处理一个 token 的 K/V 写入。
+    slot_mapping 指定了该 token 的 K/V 应该写入 cache 的哪个 slot。
+    slot == -1 表示该 token 不需要写入（如 warmup 阶段）。
+    """
     idx = tl.program_id(0)
     slot = tl.load(slot_mapping_ptr + idx)
     if slot == -1: return
@@ -31,6 +37,15 @@ def store_kvcache_kernel(
 
 
 def store_kvcache(key: torch.Tensor, value: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, slot_mapping: torch.Tensor):
+    """将新计算的 K/V 向量存储到 KV cache 中。
+
+    Args:
+        key: [N, num_heads, head_dim] 新计算的 K 向量。
+        value: [N, num_heads, head_dim] 新计算的 V 向量。
+        k_cache: KV cache 中 K 的存储区域。
+        v_cache: KV cache 中 V 的存储区域。
+        slot_mapping: [N] 每个 token 对应的 cache slot 索引。
+    """
     N, num_heads, head_dim = key.shape
     D = num_heads * head_dim
     assert key.stride(-1) == 1 and value.stride(-1) == 1
@@ -41,6 +56,18 @@ def store_kvcache(key: torch.Tensor, value: torch.Tensor, k_cache: torch.Tensor,
 
 
 class Attention(nn.Module):
+    """注意力层：封装了 Flash Attention 和 KV cache 的交互逻辑。
+
+    支持 two 阶段的注意力计算:
+    - Prefill: 使用 flash_attn_varlen_func（变长序列批量注意力），一次性处理整个 prompt。
+    - Decode: 使用 flash_attn_with_kvcache（带 KV cache 的注意力），逐 token 生成。
+
+    当存在前缀缓存时（block_tables 不为 None），prefill 阶段直接从 KV cache 读取
+    已缓存的 K/V，而不是用当前计算的 K/V。
+
+    Attributes:
+        k_cache, v_cache: 绑定到 ModelRunner 分配的全局 KV cache 的对应层视图。
+    """
 
     def __init__(
         self,
@@ -54,21 +81,23 @@ class Attention(nn.Module):
         self.head_dim = head_dim
         self.scale = scale
         self.num_kv_heads = num_kv_heads
-        self.k_cache = self.v_cache = torch.tensor([])
+        self.k_cache = self.v_cache = torch.tensor([])  # 占位，由 ModelRunner 分配后绑定
 
     def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
         context = get_context()
         k_cache, v_cache = self.k_cache, self.v_cache
         if k_cache.numel() and v_cache.numel():
+            # 将新计算的 K/V 写入 KV cache
             store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)
         if context.is_prefill:
-            if context.block_tables is not None:    # prefix cache
+            if context.block_tables is not None:    # 前缀缓存命中：从 KV cache 读取 K/V
                 k, v = k_cache, v_cache
+            # Flash Attention 变长版：支持不同长度的序列在同一批次中计算
             o = flash_attn_varlen_func(q, k, v,
                                        max_seqlen_q=context.max_seqlen_q, cu_seqlens_q=context.cu_seqlens_q,
                                        max_seqlen_k=context.max_seqlen_k, cu_seqlens_k=context.cu_seqlens_k,
                                        softmax_scale=self.scale, causal=True, block_table=context.block_tables)
-        else:    # decode
+        else:    # Decode 阶段：从 KV cache 中读取所有历史 K/V
             o = flash_attn_with_kvcache(q.unsqueeze(1), k_cache, v_cache,
                                         cache_seqlens=context.context_lens, block_table=context.block_tables,
                                         softmax_scale=self.scale, causal=True)

nanovllm/layers/embed_head.py

@@ -7,6 +7,12 @@ from nanovllm.utils.context import get_context
 
 
 class VocabParallelEmbedding(nn.Module):
+    """词表并行 Embedding：将词表按 TP rank 切分。
+
+    每个 rank 只存储词表中属于自己的分片。前向计算时：
+    1. 只查找属于自己的 token ID，其他位置输出零。
+    2. 通过 all-reduce 聚合所有 rank 的结果。
+    """
 
     def __init__(
         self,
@@ -25,6 +31,7 @@ class VocabParallelEmbedding(nn.Module):
         self.weight.weight_loader = self.weight_loader
 
     def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
+        """加载属于当前 rank 的词表分片。"""
         param_data = param.data
         shard_size = param_data.size(0)
         start_idx = self.tp_rank * shard_size
@@ -33,16 +40,27 @@ class VocabParallelEmbedding(nn.Module):
 
     def forward(self, x: torch.Tensor):
         if self.tp_size > 1:
+            # 构造 mask：标记哪些 token ID 属于当前 rank 的范围
             mask = (x >= self.vocab_start_idx) & (x < self.vocab_end_idx)
+            # 将全局 token ID 转为局部索引
             x = mask * (x - self.vocab_start_idx)
         y = F.embedding(x, self.weight)
         if self.tp_size > 1:
+            # 非当前 rank 范围的 token 输出清零，然后 all-reduce 求和
             y = mask.unsqueeze(1) * y
             dist.all_reduce(y)
         return y
 
 
 class ParallelLMHead(VocabParallelEmbedding):
+    """并行 LM Head：将隐藏状态映射为词表 logits。
+
+    与 Embedding 共享权重（如果模型配置了 tie_word_embeddings），
+    但前向逻辑不同：使用 F.linear（矩阵乘法）而非 F.embedding（查表）。
+
+    Prefill 阶段只需要每个序列最后一个 token 的 logits（因为只有最后一个 token 会用于采样），
+    所以先用 cu_seqlens 提取最后位置，再做矩阵乘法，减少计算量。
+    """
 
     def __init__(
         self,
@@ -56,10 +74,13 @@ class ParallelLMHead(VocabParallelEmbedding):
     def forward(self, x: torch.Tensor):
         context = get_context()
         if context.is_prefill:
+            # Prefill: 只取每个序列最后一个 token 的隐藏状态用于采样
             last_indices = context.cu_seqlens_q[1:] - 1
             x = x[last_indices].contiguous()
+        # 矩阵乘法得到 logits
         logits = F.linear(x, self.weight)
         if self.tp_size > 1:
+            # Gather 所有 rank 的 logits 分片到 rank 0
             all_logits = [torch.empty_like(logits) for _ in range(self.tp_size)] if self.tp_rank == 0 else None
             dist.gather(logits, all_logits, 0)
             logits = torch.cat(all_logits, -1) if self.tp_rank == 0 else None

nanovllm/layers/layernorm.py

@@ -3,6 +3,16 @@ from torch import nn
 
 
 class RMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization（RMSNorm）。
+
+    与 LayerNorm 相比，RMSNorm 不计算均值，只计算均方根，计算量更小。
+    公式: output = x / sqrt(mean(x^2) + eps) * weight
+
+    提供两个前向路径:
+    - rms_forward: 标准 RMSNorm。
+    - add_rms_forward: 将残差加法融合到 RMSNorm 中（x + residual → RMSNorm），
+      减少一次显存读写，是 vLLM 等推理框架的常见优化。
+    """
 
     def __init__(
         self,
@@ -31,6 +41,11 @@ class RMSNorm(nn.Module):
         x: torch.Tensor,
         residual: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        """融合残差加法和 RMSNorm。
+
+        Returns:
+            (normalized_output, updated_residual): 归一化后的输出和更新后的残差（= x + residual）。
+        """
         orig_dtype = x.dtype
         x = x.float().add_(residual.float())
         residual = x.to(orig_dtype)

nanovllm/layers/linear.py

@@ -5,11 +5,20 @@ import torch.distributed as dist
 
 
 def divide(numerator, denominator):
+    """整除断言，确保张量并行时维度能被均匀切分。"""
     assert numerator % denominator == 0
     return numerator // denominator
 
 
 class LinearBase(nn.Module):
+    """所有并行线性层的基类。
+
+    Attributes:
+        tp_dim: 张量并行切分的维度（0=列切分，1=行切分，None=不切分）。
+        tp_rank: 当前进程在 TP 组中的 rank。
+        tp_size: TP 组的总大小。
+        weight: 权重参数，带有 weight_loader 方法用于加载预训练权重。
+    """
 
     def __init__(
         self,
@@ -35,6 +44,7 @@ class LinearBase(nn.Module):
 
 
 class ReplicatedLinear(LinearBase):
+    """复制式线性层：所有 TP rank 持有完整的权重副本。用于不需要切分的层。"""
 
     def __init__(
         self,
@@ -52,6 +62,11 @@ class ReplicatedLinear(LinearBase):
 
 
 class ColumnParallelLinear(LinearBase):
+    """列并行线性层：将输出维度按 TP rank 切分。
+
+    每个 TP rank 持有输出维度的一个分片。例如输出维度为 4096，TP=2 时每个 rank 持有 2048。
+    常用于 QKV 投影和 FFN 的 gate/up 投影（这些层的输出可以独立计算）。
+    """
 
     def __init__(
         self,
@@ -63,6 +78,7 @@ class ColumnParallelLinear(LinearBase):
         super().__init__(input_size, divide(output_size, tp_size), bias, 0)
 
     def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
+        """加载权重时按 tp_rank 切取对应的列分片。"""
         param_data = param.data
         shard_size = param_data.size(self.tp_dim)
         start_idx = self.tp_rank * shard_size
@@ -74,6 +90,13 @@ class ColumnParallelLinear(LinearBase):
 
 
 class MergedColumnParallelLinear(ColumnParallelLinear):
+    """融合的列并行线性层：将多个线性层合并为一个矩阵乘法。
+
+    典型用途是将 gate_proj 和 up_proj 融合为 gate_up_proj，
+    减少 kernel launch 次数，提升计算效率。
+
+    权重加载时需要根据 shard_id（子层索引）定位到正确的权重分片位置。
+    """
 
     def __init__(
         self,
@@ -85,7 +108,9 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
         super().__init__(input_size, sum(output_sizes), bias)
 
     def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: int):
+        """根据 shard_id 将权重加载到融合矩阵的正确位置。"""
         param_data = param.data
+        # 计算该子层在融合矩阵中的偏移量
         shard_offset = sum(self.output_sizes[:loaded_shard_id]) // self.tp_size
         shard_size = self.output_sizes[loaded_shard_id] // self.tp_size
         param_data = param_data.narrow(self.tp_dim, shard_offset, shard_size)
@@ -94,6 +119,14 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
 
 
 class QKVParallelLinear(ColumnParallelLinear):
+    """QKV 融合的列并行线性层。
+
+    将 Q、K、V 三个投影合并为一个矩阵乘法。
+    权重按 [Q | K | V] 的顺序排列，加载时根据 shard_id（"q"/"k"/"v"）
+    定位到对应的位置。
+
+    支持 GQA：Q 的 head 数和 KV 的 head 数可以不同。
+    """
 
     def __init__(
         self,
@@ -112,6 +145,7 @@ class QKVParallelLinear(ColumnParallelLinear):
         super().__init__(hidden_size, output_size, bias)
 
     def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str):
+        """根据 shard_id ("q"/"k"/"v") 将权重加载到融合矩阵的正确位置。"""
         param_data = param.data
         assert loaded_shard_id in ["q", "k", "v"]
         if loaded_shard_id == "q":
@@ -120,7 +154,7 @@ class QKVParallelLinear(ColumnParallelLinear):
         elif loaded_shard_id == "k":
             shard_size = self.num_kv_heads * self.head_size
             shard_offset = self.num_heads * self.head_size
-        else:
+        else:  # "v"
             shard_size = self.num_kv_heads * self.head_size
             shard_offset = self.num_heads * self.head_size + self.num_kv_heads * self.head_size
         param_data = param_data.narrow(self.tp_dim, shard_offset, shard_size)
@@ -129,6 +163,14 @@ class QKVParallelLinear(ColumnParallelLinear):
 
 
 class RowParallelLinear(LinearBase):
+    """行并行线性层：将输入维度按 TP rank 切分。
+
+    每个 TP rank 持有输入维度的一个分片。前向计算后需要 all-reduce
+    将所有 rank 的结果求和，得到完整的输出。
+    常用于 O 投影和 FFN 的 down 投影（这些层的输出需要跨 rank 聚合）。
+
+    偏置项只在 rank 0 添加，避免 all-reduce 后重复加 bias。
+    """
 
     def __init__(
         self,
@@ -140,8 +182,10 @@ class RowParallelLinear(LinearBase):
         super().__init__(divide(input_size, tp_size), output_size, bias, 1)
 
     def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
+        """加载权重时按 tp_rank 切取对应的行分片。"""
         param_data = param.data
         if param_data.ndim == 1:
+            # bias 不切分，每个 rank 持有完整副本
             param_data.copy_(loaded_weight)
             return
         shard_size = param_data.size(self.tp_dim)
@@ -150,7 +194,8 @@ class RowParallelLinear(LinearBase):
         param_data.copy_(loaded_weight)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # 只有 rank 0 加 bias，避免 all-reduce 后重复
         y = F.linear(x, self.weight, self.bias if self.tp_rank == 0 else None)
         if self.tp_size > 1:
-            dist.all_reduce(y)
+            dist.all_reduce(y)  # 跨 rank 求和，得到完整输出
         return y

nanovllm/layers/rotary_embedding.py

@@ -8,6 +8,14 @@ def apply_rotary_emb(
     cos: torch.Tensor,
     sin: torch.Tensor,
 ) -> torch.Tensor:
+    """应用旋转位置编码（RoPE）。
+
+    将向量沿最后一维分成两半 (x1, x2)，然后做旋转变换:
+    y1 = x1 * cos - x2 * sin
+    y2 = x2 * cos + x1 * sin
+    这等价于在二维平面上将每个相邻的 (x1, x2) 对旋转 theta 角度，
+    其中 theta = position / (base^(2i/d))。
+    """
     x1, x2 = torch.chunk(x.float(), 2, dim=-1)
     y1 = x1 * cos - x2 * sin
     y2 = x2 * cos + x1 * sin
@@ -15,6 +23,15 @@ def apply_rotary_emb(
 
 
 class RotaryEmbedding(nn.Module):
+    """旋转位置编码（Rotary Position Embedding, RoPE）。
+
+    RoPE 通过旋转矩阵编码位置信息，使得注意力计算中
+    内积只依赖相对位置（q_i · k_j 只与 i-j 有关），
+    从而天然支持外推到更长序列。
+
+    预计算所有位置的 cos 和 sin 值并缓存，避免重复计算。
+    缓存形状: [max_position_embeddings, 1, head_dim]，中间维是 num_heads 的广播维度。
+    """
 
     def __init__(
         self,
@@ -26,11 +43,14 @@ class RotaryEmbedding(nn.Module):
         super().__init__()
         self.head_size = head_size
         assert rotary_dim == head_size
+        # 计算逆频率: 1 / (base^(2i/d)), i = 0, 1, ..., d/2-1
         inv_freq = 1.0 / (base**(torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim))
+        # 计算所有位置的频率: pos * inv_freq
         t = torch.arange(max_position_embeddings, dtype=torch.float)
         freqs = torch.einsum("i,j -> ij", t, inv_freq)
         cos = freqs.cos()
         sin = freqs.sin()
+        # 拼接 cos 和 sin，形状 [max_pos, 1, head_dim]
         cache = torch.cat((cos, sin), dim=-1).unsqueeze_(1)
         self.register_buffer("cos_sin_cache", cache, persistent=False)
 
@@ -41,6 +61,7 @@ class RotaryEmbedding(nn.Module):
         query: torch.Tensor,
         key: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        """根据位置索引查找 cos/sin 并应用到 Q 和 K。"""
         cos_sin = self.cos_sin_cache[positions]
         cos, sin = cos_sin.chunk(2, dim=-1)
         query = apply_rotary_emb(query, cos, sin)
@@ -55,5 +76,9 @@ def get_rope(
     max_position: int,
     base: float,
 ):
+    """获取 RotaryEmbedding 的单例实例。
+
+    使用 lru_cache 确保相同参数只创建一次。
+    """
     rotary_emb = RotaryEmbedding(head_size, rotary_dim, max_position, base)
     return rotary_emb

nanovllm/layers/sampler.py

@@ -3,10 +3,21 @@ from torch import nn
 
 
 class Sampler(nn.Module):
+    """采样器：将 logits 转换为 token ID。
+
+    使用 Gumbel-like 采样方法（而非标准的 top-k/top-p）：
+    1. 将 logits 除以温度（temperature）。
+    2. 计算 softmax 得到概率分布。
+    3. 用指数分布噪声扰动概率，取 argmax。
+
+    这种方法等价于从 softmax(logits/temperature) 分布中采样，
+    但避免了逐元素随机选择的低效操作，全部用张量运算实现。
+    """
 
     @torch.compile
     def forward(self, logits: torch.Tensor, temperatures: torch.Tensor):
         logits = logits.float().div_(temperatures.unsqueeze(dim=1))
         probs = torch.softmax(logits, dim=-1)
+        # 指数分布噪声采样：probs / Exp(1) 的 argmax 等价于按 probs 概率采样
         sample_tokens = probs.div_(torch.empty_like(probs).exponential_(1).clamp_min_(1e-10)).argmax(dim=-1)
         return sample_tokens

nanovllm/models/qwen3.py

@@ -12,6 +12,16 @@ from nanovllm.layers.embed_head import VocabParallelEmbedding, ParallelLMHead
 
 
 class Qwen3Attention(nn.Module):
+    """Qwen3 的注意力层。
+
+    支持 GQA（Grouped Query Attention）：num_kv_heads 可以小于 num_heads，
+    多个 query head 共享同一组 KV head，减少 KV cache 的显存占用。
+
+    当模型配置中没有 qkv_bias 时（Qwen3 默认无 bias），会在 Q 和 K 投影后
+    添加 RMSNorm（q_norm, k_norm），这是 Qwen3 的特殊设计。
+
+    支持张量并行（TP）：QKV 投影按列切分，O 投影按行切分。
+    """
 
     def __init__(
         self,
@@ -29,16 +39,17 @@ class Qwen3Attention(nn.Module):
         tp_size = dist.get_world_size()
         self.total_num_heads = num_heads
         assert self.total_num_heads % tp_size == 0
-        self.num_heads = self.total_num_heads // tp_size
+        self.num_heads = self.total_num_heads // tp_size  # 当前 TP rank 拥有的 query head 数
         self.total_num_kv_heads = num_kv_heads
         assert self.total_num_kv_heads % tp_size == 0
-        self.num_kv_heads = self.total_num_kv_heads // tp_size
+        self.num_kv_heads = self.total_num_kv_heads // tp_size  # 当前 TP rank 拥有的 KV head 数
         self.head_dim = head_dim or hidden_size // self.total_num_heads
         self.q_size = self.num_heads * self.head_dim
         self.kv_size = self.num_kv_heads * self.head_dim
-        self.scaling = self.head_dim ** -0.5
+        self.scaling = self.head_dim ** -0.5  # 注意力缩放因子: 1/sqrt(d_k)
         self.qkv_bias = qkv_bias
 
+        # QKV 融合投影：将 hidden_states 投影为 Q、K、V，按 TP 切分
         self.qkv_proj = QKVParallelLinear(
             hidden_size,
             self.head_dim,
@@ -46,6 +57,7 @@ class Qwen3Attention(nn.Module):
             self.total_num_kv_heads,
             bias=qkv_bias,
         )
+        # 输出投影：将注意力输出映射回 hidden_size，按 TP 行切分
         self.o_proj = RowParallelLinear(
             self.total_num_heads * self.head_dim,
             hidden_size,
@@ -53,6 +65,7 @@ class Qwen3Attention(nn.Module):
         )
         if isinstance(rope_scaling, dict):
             rope_theta = rope_scaling.get("rope_theta", rope_theta)
+        # 旋转位置编码（RoPE）
         self.rotary_emb = get_rope(
             self.head_dim,
             rotary_dim=self.head_dim,
@@ -65,6 +78,7 @@ class Qwen3Attention(nn.Module):
             self.scaling,
             self.num_kv_heads,
         )
+        # Qwen3 无 bias 时使用 Q/K norm（post-normalization）
         if not self.qkv_bias:
             self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
             self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
@@ -75,20 +89,30 @@ class Qwen3Attention(nn.Module):
         hidden_states: torch.Tensor,
     ) -> torch.Tensor:
         qkv = self.qkv_proj(hidden_states)
+        # 将 QKV 融合结果拆分为 Q、K、V
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
         q = q.view(-1, self.num_heads, self.head_dim)
         k = k.view(-1, self.num_kv_heads, self.head_dim)
         v = v.view(-1, self.num_kv_heads, self.head_dim)
         if not self.qkv_bias:
+            # Qwen3 特有：对 Q 和 K 做 RMSNorm
             q = self.q_norm(q)
             k = self.k_norm(k)
+        # 应用旋转位置编码
         q, k = self.rotary_emb(positions, q, k)
+        # 注意力计算（含 KV cache 读写）
         o = self.attn(q, k, v)
+        # 输出投影
         output = self.o_proj(o.flatten(1, -1))
         return output
 
 
 class Qwen3MLP(nn.Module):
+    """Qwen3 的前馈网络（MLP），使用 SwiGLU 激活函数。
+
+    结构: hidden_size → 2×intermediate_size (gate + up) → intermediate_size → hidden_size
+    其中 gate 和 up 投影融合为一个矩阵乘法，然后 SiLU(gate) * up。
+    """
 
     def __init__(
         self,
@@ -97,9 +121,10 @@ class Qwen3MLP(nn.Module):
         hidden_act: str,
     ) -> None:
         super().__init__()
+        # 融合 gate_proj 和 up_proj 为一个矩阵，减少一次 kernel launch
         self.gate_up_proj = MergedColumnParallelLinear(
             hidden_size,
-            [intermediate_size] * 2,
+            [intermediate_size] * 2,  # gate 和 up 各输出 intermediate_size
             bias=False,
         )
         self.down_proj = RowParallelLinear(
@@ -108,7 +133,7 @@ class Qwen3MLP(nn.Module):
             bias=False,
         )
         assert hidden_act == "silu"
-        self.act_fn = SiluAndMul()
+        self.act_fn = SiluAndMul()  # SiLU(gate) * up
 
     def forward(self, x):
         gate_up = self.gate_up_proj(x)
@@ -118,6 +143,11 @@ class Qwen3MLP(nn.Module):
 
 
 class Qwen3DecoderLayer(nn.Module):
+    """Qwen3 的单个 Transformer 解码层。
+
+    结构: Input RMSNorm → Self-Attention → Residual → Post-Attention RMSNorm → MLP → Residual
+    使用 Pre-Norm 架构（先归一化再进入子层），并将残差连接的计算融合到 RMSNorm 中以节省显存。
+    """
 
     def __init__(
         self,
@@ -149,6 +179,7 @@ class Qwen3DecoderLayer(nn.Module):
         hidden_states: torch.Tensor,
         residual: torch.Tensor | None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        # 残差连接融合到 RMSNorm 中：residual = hidden_states + residual, output = RMSNorm(residual)
         if residual is None:
             hidden_states, residual = self.input_layernorm(hidden_states), hidden_states
         else:
@@ -160,6 +191,7 @@ class Qwen3DecoderLayer(nn.Module):
 
 
 class Qwen3Model(nn.Module):
+    """Qwen3 的 Transformer 主体：Embedding → N × DecoderLayer → Final RMSNorm。"""
 
     def __init__(
         self,
@@ -179,11 +211,19 @@ class Qwen3Model(nn.Module):
         residual = None
         for layer in self.layers:
             hidden_states, residual = layer(positions, hidden_states, residual)
+        # 最终的 RMSNorm 也融合了残差加法
         hidden_states, _ = self.norm(hidden_states, residual)
         return hidden_states
 
 
 class Qwen3ForCausalLM(nn.Module):
+    """Qwen3 因果语言模型：用于文本生成。
+
+    包含映射表，将 HuggingFace 的独立权重名（q_proj, k_proj, v_proj, gate_proj, up_proj）
+    映射到本项目融合后的权重名（qkv_proj, gate_up_proj），以便正确加载权重。
+    """
+
+    # 融合模块的映射关系：HF 权重名 → (本项目模块名, shard_id)
     packed_modules_mapping = {
         "q_proj": ("qkv_proj", "q"),
         "k_proj": ("qkv_proj", "k"),
@@ -199,6 +239,7 @@ class Qwen3ForCausalLM(nn.Module):
         super().__init__()
         self.model = Qwen3Model(config)
         self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
+        # 如果模型配置了权重共享（tie），LM Head 和 Embedding 使用同一个权重矩阵
         if config.tie_word_embeddings:
             self.lm_head.weight.data = self.model.embed_tokens.weight.data
 
@@ -213,4 +254,5 @@ class Qwen3ForCausalLM(nn.Module):
         self,
         hidden_states: torch.Tensor,
     ) -> torch.Tensor:
+        """将最后一层隐藏状态通过 LM Head 转换为词表 logits。"""
         return self.lm_head(hidden_states)

nanovllm/sampling_params.py

@@ -3,6 +3,16 @@ from dataclasses import dataclass
 
 @dataclass(slots=True)
 class SamplingParams:
+    """生成采样的参数配置。
+
+    Args:
+        temperature: 采样温度，控制输出的随机性。值越大越随机，越接近 0 越确定。
+                     注意：本项目不支持 temperature=0（贪心解码），必须大于 1e-10。
+        max_tokens: 单个请求最大生成的 token 数量。
+        ignore_eos: 是否忽略 EOS token。设为 True 时即使遇到结束符也继续生成，直到 max_tokens 耗尽。
+                    基准测试中用于确保每个请求都生成固定数量的 token。
+    """
+
     temperature: float = 1.0
     max_tokens: int = 64
     ignore_eos: bool = False

nanovllm/utils/context.py

@@ -4,6 +4,21 @@ import torch
 
 @dataclass(slots=True)
 class Context:
+    """全局上下文：存储当前推理步骤的注意力相关元数据。
+
+    这个对象在每次推理步骤开始时被 ModelRunner 设置，在模型的前向传播中
+    被 Attention 层读取。它是一个全局单例，避免了通过函数参数层层传递。
+
+    Attributes:
+        is_prefill: 当前是否为 prefill 阶段。
+        cu_seqlens_q: 查询的累积序列长度（prefill 阶段使用），标记每个序列的边界。
+        cu_seqlens_k: 键值的累积序列长度（prefill 阶段使用），可能与 cu_seqlens_q 不同（前缀缓存）。
+        max_seqlen_q: 批次中最长的查询序列长度（prefill 使用）。
+        max_seqlen_k: 批次中最长的键值序列长度（prefill 使用）。
+        slot_mapping: 每个 token 在 KV cache 中的存储位置索引（用于写入新的 K/V）。
+        context_lens: 每个序列的上下文总长度（decode 阶段使用）。
+        block_tables: KV cache 块映射表，将逻辑块映射到物理块（decode 和前缀缓存使用）。
+    """
     is_prefill: bool = False
     cu_seqlens_q: torch.Tensor | None = None
     cu_seqlens_k: torch.Tensor | None = None
@@ -16,12 +31,15 @@ class Context:
 _CONTEXT = Context()
 
 def get_context():
+    """获取当前全局上下文。"""
     return _CONTEXT
 
 def set_context(is_prefill, cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=0, max_seqlen_k=0, slot_mapping=None, context_lens=None, block_tables=None):
+    """设置当前推理步骤的全局上下文。"""
     global _CONTEXT
     _CONTEXT = Context(is_prefill, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, slot_mapping, context_lens, block_tables)
 
 def reset_context():
+    """重置全局上下文（推理步骤结束后调用）。"""
     global _CONTEXT
     _CONTEXT = Context()

nanovllm/utils/loader.py

@@ -6,23 +6,38 @@ from safetensors import safe_open
 
 
 def default_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor):
+    """默认权重加载器：直接将加载的权重拷贝到参数中。"""
     param.data.copy_(loaded_weight)
 
 
 def load_model(model: nn.Module, path: str):
+    """从 HuggingFace safetensors 格式加载模型权重。
+
+    支持融合模块的权重加载：本项目将 Q/K/V 投影融合为 qkv_proj，
+    将 gate/up 投影融合为 gate_up_proj。加载时需要通过 packed_modules_mapping
+    将原始的独立权重名映射到融合后的模块，并使用自定义的 weight_loader
+    将权重放置到正确位置。
+
+    Args:
+        model: 要加载权重的模型。
+        path: 模型目录路径，包含 .safetensors 文件。
+    """
     packed_modules_mapping = getattr(model, "packed_modules_mapping", {})
     for file in glob(os.path.join(path, "*.safetensors")):
         with safe_open(file, "pt", "cpu") as f:
             for weight_name in f.keys():
+                # 检查是否为融合模块的子权重（如 q_proj, k_proj, gate_proj 等）
                 for k in packed_modules_mapping:
                     if k in weight_name:
                         v, shard_id = packed_modules_mapping[k]
+                        # 替换权重名：如 "model.layers.0.self_attn.q_proj.weight" → "...qkv_proj.weight"
                         param_name = weight_name.replace(k, v)
                         param = model.get_parameter(param_name)
                         weight_loader = getattr(param, "weight_loader")
                         weight_loader(param, f.get_tensor(weight_name), shard_id)
                         break
                 else:
+                    # 普通权重：直接加载
                     param = model.get_parameter(weight_name)
                     weight_loader = getattr(param, "weight_loader", default_weight_loader)
                     weight_loader(param, f.get_tensor(weight_name))