Rain-Bus
ffd2defdfc
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>
44 lines
2.3 KiB
Python
44 lines
2.3 KiB
Python
import os
|
|
from dataclasses import dataclass
|
|
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
|
|
max_model_len: int = 4096
|
|
gpu_memory_utilization: float = 0.9
|
|
tensor_parallel_size: int = 1
|
|
enforce_eager: bool = False
|
|
hf_config: AutoConfig | None = None
|
|
eos: int = -1
|
|
kvcache_block_size: int = 256
|
|
num_kvcache_blocks: int = -1
|
|
|
|
def __post_init__(self):
|
|
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)
|