模型训练收工的那一刻通常是最开心的接下去要把模型送上线才是最磨人的。手里一个小众模型训练时用的自定义框架checkpoint 命名随心所欲架构和主流模型有那么点像又不完全一样。你想用 vLLM 上线启动脚本直接报错“Model architectures are not supported”换回原项目自带的推理脚本单请求低吞吐批处理调度、显存管理全得自己写。于是只好硬着头皮打开 vLLM 源码研究怎么把一个自定义模型塞进去。这篇文章就是完整记录这件事的做法。我会拿一个完全虚构、但结构上很典型的小众因果语言模型 MiniStream 作为例子从权重转换、模型类编写、注册机制到启动验证把“从零给 vLLM 添加自定义模型推理支持”的完整链路捋一遍。适合三类人被“不支持的架构”卡住的部署工程师想把自己训练的模型接入 vLLM 的算法同学以及想搞懂 vLLM 模型适配机制的框架学习者。1. 为什么“加模型支持”这件事绕不开1.1 小众模型上线的两条死路先说一个现实vLLM 的高吞吐能力本质上来自它对模型结构做了大量假设和预优化。连续批处理、paged KV cache、量化、算子融合这些都是按主流 decoder-only 架构的规律去设计的。你的模型越是“长得像主流模型”接入成本越低越是个性化越要自己动手补适配层。小众模型常见的困境是两条路都不好走。第一条路是用原项目自带推理脚本上线。论文代码也好、内部实验代码也好大部分只保证“能跑”而不是“跑得快”。单条请求推理时GPU 利用率经常只有个位数百分比同时来十个请求内存翻倍、速度骤降。想支持流式输出、多轮对话、并发排队都得自己搭。最要命的是显存里的 KV cache 每次请求都要重新算长对话根本扛不住。这条路对小模型或许能撑模型一上亿参数就非常难受。第二条路是用通用推理框架跑但挑一个“看起来支持”的。很多框架只支持固定名单里的模型要么你的模型压根不在列表里要么得走 ONNX/TensorRT 导出精度和算子覆盖又容易出问题。好不容易导出成功模型里一个自定义位置编码可能就让整条链路废掉。折腾半天最后还是回到原点。相比之下直接给 vLLM 写一个模型适配文件一次性投入之后就能吃到 vLLM 整套调度和显存管理红利这是性价比最高的路径。只是这活儿有个门槛你得先搞清楚 vLLM 到底是怎么“认识”一个模型的。1.2 vLLM 认识模型的最小链条很多人以为 vLLM 支持某个模型是靠“扫描模型文件自动识别”其实完全不是。它内部维护了一张注册表这张表把“模型架构名”映射到“模型类”。你调用 LLM(model...) 时引擎会走一条固定的识别链路读取模型目录下的 config.json拿到 architectures 字段比如 [MiniStreamForCausalLM]。拿这个架构名去模型注册表里查看有没有注册过对应的类。查到了就用配置文件实例化这个模型类然后调用它的 load_weights 方法把权重灌进去。查不到直接抛“Model architectures are not supported”。所以“给 vLLM 加一个模型”这个任务拆开其实就是四件事写一个模型类、把它注册到注册表、保证配置文件能被正确解析、实现权重加载逻辑。其中权重加载通常是最耗时的因为你的 checkpoint 命名几乎一定和 vLLM 期望的不一样。这个链路听起来简单实际做起来处处是细节。下面把四个关键点一个个说清楚。2. 动手前必须摸清的四个关键点2.1 模型类长什么样先看一个最小可用的模型类长什么样。以 vLLM 0.6.x 这个相对稳定的版本为例模型类一般长这样import torch import torch.nn as nn from transformers import PretrainedConfig from vllm.model_executor.models.registry import ModelRegistry from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.sampler import Sampler ModelRegistry.register_model(MiniStreamForCausalLM) class MiniStreamForCausalLM(nn.Module): def __init__(self, config: PretrainedConfig, cache_configNone): super().__init__() self.config config self.model MiniStreamModel(config) self.lm_head ColumnParallelLinear( config.hidden_size, config.vocab_size, biasFalse ) self.sampler Sampler() def forward(self, input_ids, positions, kv_caches, attn_metadata, **kwargs): hidden_states self.model(input_ids, positions, kv_caches, attn_metadata) return hidden_states def compute_logits(self, hidden_states, sampling_metadata): return self.lm_head(hidden_states) def sample(self, logits, sampling_metadata): return self.sampler(logits, sampling_metadata) def load_weights(self, weights): # 核心工作量都在这 pass几个容易被忽略的细节forward 不直接返回 logits而是返回最后一层 hidden states。后续的 logits 计算和采样由 compute_logits、sample 这两个方法负责这是 vLLM 为了做采样处理、logits processor 而设计的拆分。init里那个 cache_config 参数有的版本有有的版本没有。保险写法是给它默认值 None避免不同版本实例化时报参数数量不匹配。注册装饰器里的字符串必须和 config.json 里 architectures[0] 完全一致大小写、空格都得一样否则注册了也找不到。2.2 权重映射才是真正的重头戏模型类说白了是个空壳真正让一个模型跑起来的是权重映射。vLLM 加载权重时会把 checkpoint 里的每个参数按名字喂给 load_weights你要做的就是把 checkpoint 名字翻译成你的模型参数名再把数据拷进去。最朴素的写法是def load_weights(self, weights): params_dict dict(self.named_parameters()) for item in weights: name getattr(item, name, item[0]) loaded_weight getattr(item, loaded_weight, item[1]) if name in params_dict: params_dict[name].data.copy_(loaded_weight)为什么这里是重头戏因为绝大多数自定义 checkpoint 不会按照“model.layers.0.self_attn.q_proj.weight”这种规范命名。我见过太多情况QKV 是融合成一个 Wqkv.weight 矩阵需要拆成 q_proj、k_proj、v_proj 三份门控 MLP 的两个权重拼在一起需要拆成 gate_proj 和 up_projembedding 和 lm_head 是权重绑定的checkpoint 里只存一份层编号从 1 开始而不是从 0 开始还有直接把 bias、scale 之类参数塞进同一个文件里的。你每遇到一种情况就要在 load_weights 里加一个分支去处理和翻译。如果模型有 30 多层每一层都这样处理分支代码会很长但逻辑本身并不复杂。关键是先把 checkpoint 的键名和模型参数名两边都打印出来做一次差集心里有数再动手。2.3 config.json 别靠猜vLLM 读配置的方式是先尝试用 transformers 生态的标准配置读取再拿着里面的字段去实例化模型类。所以 config.json 里的字段质量直接决定后面是否顺利。一份能用的 config.json 至少得包含这些{ architectures: [MiniStreamForCausalLM], model_type: ministream, hidden_size: 768, num_hidden_layers: 8, num_attention_heads: 12, num_key_value_heads: 12, head_dim: 64, intermediate_size: 2048, vocab_size: 32000, max_position_embeddings: 2048, rms_norm_eps: 1e-6, rope_theta: 10000.0, bos_token_id: 1, eos_token_id: 2, tie_word_embeddings: true }我这里特别提醒三个坑。第一个坑是 head_dim。很多模型 config 里不写 head_dim而是靠 hidden_size 除以 num_attention_heads 去推算。如果你的模型里 head_dim 和这个除法结果不一致就得显式写上否则注意力计算会直接形状不匹配。第二个坑是 architectures 字段。有些实验代码导出的 config 里architectures 写的是训练脚本里的类名比如 “MiniStreamTrainerModel”跟推理模型类名完全对不上。vLLM 只会拿这个字段去注册表里查名字对不上就是 not supported。第三个坑是 tie_word_embeddings。如果你的模型 embedding 和 lm_head 是共享权重的而 vLLM 模型实现里两者是独立参数那么加载时要人为复制一份权重。这个处理不好会出现“生成时第一个词概率异常”这类极难排查的问题。2.4 注意力能不能复用决定你后面省不省心给 vLLM 写模型类最大的性能分水岭在于注意力层怎么实现。如果你的模型用的是标准 RoPE MHA/GQA直接复用 vLLM 的 Attention 层就好paged KV cache、前缀缓存、滑动窗口这些能力全部自动生效。这是最理想的情况。如果你的模型有特殊的位置编码、特殊 attention mask 或者需要往 attention score 里加额外的 bias就得斟酌了。vLLM 的 Attention 层封装了下层各种后端不是所有自定义行为都支持。有人选择回到 PyTorch 原生实现比如 F.scaled_dot_product_attention这样功能上没问题模型照样能跑但你会失去 paged KV cache 的优势KV cache 显存占用和推理延迟都会明显变差。我的建议是分两步走第一版先保证能跑通哪怕用原生注意力实现跑通之后再回过头来看能不能把特殊注意力行为改造到 vLLM Attention 支持的范围内。别一上来就追求完美性能先让整条链路活起来。3. 实操把一个非标准 checkpoint 接进 vLLM3.1 样例模型MiniStream 与原始权重结构现在进入实战。假设我们手里有个内部训练的小模型 MiniStream参数量约 1.3 亿结构不复杂standard decoder-only8 层 Transformerhidden size 76812 个注意力头SwiGLU 激活RMSNormRoPE 位置编码。结构上一点都不特殊但 checkpoint 格式非常“实验风”。原始 checkpoint 是单个 model.pt 文件键名长这样embedding.weight形状 [32000, 768]blocks.0.attn.Wqkv.weight形状 [2304, 768]Q、K、V 按行顺序拼在一起blocks.0.attn.out_proj.weight形状 [768, 768]blocks.0.mlp.w12.weight形状 [4096, 768]SwiGLU 的两个门控矩阵拼在一起blocks.0.mlp.w3.weight形状 [768, 2048]blocks.0.norm1.weight、blocks.0.norm2.weight形状 [768]没有独立的 lm_headembedding 与输出头共享权重也就是说训练框架导出时完全按“角色”命名不关心 transformers 生态的命名规范。我们要做的第一件事就是写个脚本把它翻译成 vLLM 能消化的 HF 风格权重。3.2 第一步写权重转换脚本转换脚本的逻辑很简单逐个键名做映射。需要特别注意的是 Wqkv 融合矩阵和 w12 融合矩阵的处理方式。import torch from safetensors.torch import save_file raw torch.load(mini-stream-132m/model.pt, map_locationcpu) raw raw[model] if model in raw else raw mapped {} mapped[model.embed_tokens.weight] raw[embedding.weight] mapped[lm_head.weight] raw[embedding.weight] # tied 权重存两份 hidden_size 768 mid hidden_size # 因为 MiniStream 是 MHAQ/K/V 各占 hidden_size intermediate 2048 for i in range(8): src fblocks.{i} dst fmodel.layers.{i} fused_qkv raw[f{src}.attn.Wqkv.weight] # [2304, 768] q, k, v torch.split(fused_qkv, [hidden_size, hidden_size, hidden_size], dim0) mapped[f{dst}.self_attn.q_proj.weight] q mapped[f{dst}.self_attn.k_proj.weight] k mapped[f{dst}.self_attn.v_proj.weight] v mapped[f{dst}.self_attn.o_proj.weight] raw[f{src}.attn.out_proj.weight] w12 raw[f{src}.mlp.w12.weight] # [4096, 768]对应 gate_up_proj mapped[f{dst}.mlp.gate_up_proj.weight] w12 mapped[f{dst}.mlp.down_proj.weight] raw[f{src}.mlp.w3.weight] mapped[f{dst}.input_layernorm.weight] raw[f{src}.norm1.weight] mapped[f{dst}.post_attention_layernorm.weight] raw[f{src}.norm2.weight] save_file(mapped, mini-stream-hf/model.safetensors)转完之后别急着进 vLLM先把 config.json 和 tokenizer 相关文件准备好目录结构如下mini-stream-hf/ ├── config.json ├── model.safetensors ├── tokenizer.json └── tokenizer_config.json这一步完成后最好先用 transformers 生态的参考实现把转换后的权重跑一遍生成一个 greedy 结果存下来。这叫“黄金输出”后面验证 vLLM 接入是否正确就用它对比。3.3 第二步写模型类并注册模型类的完整代码会包含三个部分基础模块注意力、MLP、层、顶层模型、以及注册入口。MiniStream 结构不特殊注意力可以直接用 vLLM 的 Attention 层。下面是完整实现思路。from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear from vllm.model_executor.layers.vocab_parallel_embeddings import VocabParallelEmbedding from vllm.attention.layer import Attention class MiniStreamAttention(nn.Module): def __init__(self, config): super().__init__() self.hidden_size config.hidden_size self.num_heads config.num_attention_heads self.head_dim config.head_dim self.q_proj ColumnParallelLinear(self.hidden_size, self.hidden_size, biasFalse) self.k_proj ColumnParallelLinear(self.hidden_size, self.hidden_size, biasFalse) self.v_proj ColumnParallelLinear(self.hidden_size, self.hidden_size, biasFalse) self.o_proj RowParallelLinear(self.hidden_size, self.hidden_size, biasFalse) self.attn Attention(self.num_heads, self.head_dim) def forward(self, hidden_states, positions, kv_cache, attn_metadata): q, _ self.q_proj(hidden_states) k, _ self.k_proj(hidden_states) v, _ self.v_proj(hidden_states) q q.view(-1, self.num_heads, self.head_dim) k k.view(-1, self.num_heads, self.head_dim) v v.view(-1, self.num_heads, self.head_dim) out self.attn(q, k, v, kv_cache, attn_metadata) out out.view(-1, self.hidden_size) out, _ self.o_proj(out) return out class MiniStreamMLP(nn.Module): def __init__(self, config): super().__init__() self.gate_up_proj ColumnParallelLinear( config.hidden_size, 2 * config.intermediate_size, biasFalse ) self.act_and_mul SiluAndMul() self.down_proj RowParallelLinear( config.intermediate_size, config.hidden_size, biasFalse ) def forward(self, x): x, _ self.gate_up_proj(x) x self.act_and_mul(x) x, _ self.down_proj(x) return x class MiniStreamDecoderLayer(nn.Module): def __init__(self, config): super().__init__() self.input_layernorm RMSNorm(config.hidden_size, epsconfig.rms_norm_eps) self.self_attn MiniStreamAttention(config) self.post_attention_layernorm RMSNorm(config.hidden_size, epsconfig.rms_norm_eps) self.mlp MiniStreamMLP(config) def forward(self, hidden_states, positions, kv_cache, attn_metadata): residual hidden_states hidden_states self.input_layernorm(hidden_states) hidden_states self.self_attn(hidden_states, positions, kv_cache, attn_metadata) hidden_states residual hidden_states residual hidden_states hidden_states self.post_attention_layernorm(hidden_states) hidden_states self.mlp(hidden_states) hidden_states residual hidden_states return hidden_states class MiniStreamModel(nn.Module): def __init__(self, config): super().__init__() self.embed_tokens VocabParallelEmbedding( config.vocab_size, config.hidden_size ) self.layers nn.ModuleList([ MiniStreamDecoderLayer(config) for _ in range(config.num_hidden_layers) ]) self.norm RMSNorm(config.hidden_size, epsconfig.rms_norm_eps) def forward(self, input_ids, positions, kv_caches, attn_metadata): hidden_states self.embed_tokens(input_ids) for i, layer in enumerate(self.layers): hidden_states layer(hidden_states, positions, kv_caches[i], attn_metadata) hidden_states self.norm(hidden_states) return hidden_states现在把最关键的 load_weights 补上。因为转换脚本已经把 fused 的 Wqkv 拆成了 q/k/v 三份所以 load_weights 里大部分情况是直接拷贝只有两个地方需要特殊处理lm_head 是 tied 的embedding 和 lm_head 用同一份权重还有 gate_up_proj 这个参数名要认出来。def load_weights(self, weights): params_dict dict(self.named_parameters()) for item in weights: name getattr(item, name, item[0]) loaded_weight getattr(item, loaded_weight, item[1]) if name model.embed_tokens.weight: params_dict[model.embed_tokens.weight].data.copy_(loaded_weight) params_dict[lm_head.weight].data.copy_(loaded_weight) continue if name in params_dict: params_dict[name].data.copy_(loaded_weight)为什么可以这么简单因为我们在转换脚本里已经把命名统一成了 vLLM 模型类里的参数名。这是最推荐的做法能通过转换脚本解决的问题就不要在 load_weights 里堆分支。注册方式有两种。一种是直接把 ministream.py 放进 vLLM 源码目录 vllm/model_executor/models/ 下并在该目录的init.py 里加一行 import。这个办法最省事但每次升级 vLLM 都要重新打补丁。另一种是写成独立的 Python 包在调用 LLM 之前的启动脚本里先 import 你的模型模块让 ModelRegistry.register_model 装饰器在引擎启动前完成注册。第二种更干净我后来一直用这种。3.4 第三步启动推理与正确性验证代码写完之后直接用 vLLM 启动from vllm import LLM, SamplingParams llm LLM( model./mini-stream-hf, dtypebfloat16, max_model_len2048, enforce_eagerTrue, # 调试阶段先关掉 CUDA Graph gpu_memory_utilization0.8, ) prompts [请用一句话介绍你自己。, 解释一下什么是 KV Cache。] outputs llm.generate(prompts, SamplingParams(max_tokens64, temperature0.2)) for o in outputs: print(o.outputs[0].text)如果这里能正常输出恭喜最难的 90% 已经过了。但“能输出”不等于“输出正确”。我强烈建议做一次黄金输出对比拿转换后的 model.safetensors 加载进参考实现用同一个 prompt、同样的采样参数最好 temperature0 的 greedy生成一份结果再和 vLLM 的输出逐 token 对比。我第一次接入模型时输出看起来通顺但首 token 命中率只有七成左右。后来发现是注意力缩放系数里用了 hidden_size 的平方根而原模型用的是 head_dim 的平方根。这类问题靠肉眼几乎不可能发现只能靠机械对比。4. 验证跑通之后性能与版本差异4.1 性能表现怎么看正确性过了之后把 enforce_eager 关掉恢复 CUDA Graph再重新跑一遍。这时候你会看到吞吐明显提升。再检查几个点第一确认你的注意力层确实走到了 vLLM 的 Attention。MiniStream 是标准 MHA所以 paged KV cache、前缀缓存都能直接用。如果你的模型用了 SDPA 兜底性能会差一大截那就得回头想想怎么把特殊结构标准化。第二看 KV cache 显存占用。vLLM 启动时会自动计算 KV cache 大小但如果你的 max_model_len 设得很大KV cache 会吃掉大量显存。可以从 gpu_memory_utilization、max_num_seqs、max_model_len 三个参数入手调整。第三想上多卡张量并行的话load_weights 的写法要升级。上面的例子在单卡下没问题多卡时 ColumnParallelLinear 的参数会被切分你需要处理每个 rank 拿到对应分片。最简单的方式是参考同版本 vLLM 官方模型文件里对 weight_loader 的处理把每个参数的加载交给 weight_loader而不是直接对 param.data 整体 copy。这里我建议直接抄同版本 LLaMA 类模型的做法别自己发明。4.2 不同 vLLM 版本的适配差异vLLM 迭代速度快得离谱不同版本的接口差异是我踩过最多坑的地方。核心思想是你的模型文件最好基于你部署版本里某个现成模型文件做“差分”而不是拿着老博客的代码直接套。版本区间主要差异与建议0.4.xSampler 和模型调度接口比较老模型文件结构与 0.6 差异较大按同版本现成模型改0.5/0.6.x本文的写法基本适用ModelRegistry.register_model 和 vLLM Attention 层是稳定接口0.7/0.8引入 SupportsV0/SupportsV1 协议模型类需要实现对应接口load_weights 里拿到的权重对象带 shard 信息1.x 及更新默认 V1 调度接口进一步收敛插件化机制更完善推荐把自定义模型模块打成独立包注册不管哪个版本load_weights 的“语义”从来没变过上游把 checkpoint 权重一个接一个传给你你负责翻译成模型参数并拷贝。只要抓住这个核心版本升级对你来说只是换接口写法不是换思路。5. 常见问题与排查实录5.1 注册了还是 not supported这是第一个容易翻车的点。排查顺序如下先确认 config.json 里 architectures 的值和注册时的字符串一字不差包括大小写。再确认你的模型模块确实被 import 了。如果你用的是自定义包的方式很常见的情况是启动脚本里忘了 import装饰器根本没执行注册表里自然没有。最后确认 vLLM 的模型注册表是否已经包含你的类可以在代码里临时打印出来看一眼。5.2 权重加载各种 mismatch权重相关报错是最常见的本质都是“vLLM 模型参数名”和“checkpoint 键名”对不上。解决手段是写一个小脚本同时打印两边的键名# 打印 checkpoint 键名 from safetensors import safe_open with safe_open(mini-stream-hf/model.safetensors, frameworkpt) as f: print(\n.join(sorted(f.keys()))) # 打印模型参数名 model MiniStreamForCausalLM(config) print(\n.join(sorted(dict(model.named_parameters()).keys())))对着两份名单做差集哪个名字少、哪个名字多一目了然。如果你看到 size mismatch 而不是 name mismatch那通常是某个维度算错了比如 QKV 拆分的 hidden_size 写错或者词表大小因 padding 和实际权重不一致。5.3 显存爆了与生成变慢如果是加载阶段就爆显存优先怀疑 dtype。默认加载 fp32 很容易爆改成 bfloat16 或 float16 会让显存直接减半。如果加载没问题一跑生成就爆重点查 max_model_len 和 gpu_memory_utilization。调试阶段还可以加上 enforce_eager把 CUDA Graph 的显存开销省掉。如果是生成速度越来越慢而不是一开始就慢大概率是 KV cache 没有被正确复用。检查你的注意力实现是不是每次 forward 都重新计算了历史 token。如果模型用了自定义 attention 兜底这个问题尤其常见。5.4 输出和原模型对不上输出不一致时按这个优先级排查dtype 不一致比如原模型是 fp32vLLM 里用了 bf16某些层会有一点点漂移但不会离谱。RMSNorm 的 eps。很多模型这个值是 1e-6 或 1e-5差一个数量级小概率出现明显偏差。RoPE 相关参数包括 rope_theta、rope_scaling甚至 max_position_embeddings。位置编码类的 bug 往往在长序列才暴露短 prompt 看不出问题。注意力缩放系数。检查注意力分数乘的是 head_dim 的倒数平方根还是 hidden_size 的倒数平方根。词表 padding。如果你的词表大小不是 32 的倍数某些实现会 padding 到对齐导致 lm_head 的权重形状和 embedding 对不上。以上都排完了还有一个大招在模型类里加一个临时 debug 钩子把每一层的 hidden states 导出成文件和参考实现逐层做余弦相似度。哪一层开始分叉问题就定位在哪一层。6. 我踩过的坑与几个保命习惯最后说几句实在的都是真金白银踩出来的经验。第一个习惯动手写模型类之前先到你的 vLLM 版本里找结构最接近的官方模型文件复制一份然后基于这份拷贝去改。版本接口再怎么变化你手里的官方文件就是最准确的“标准答案”。网上教程再详细也不如你自己差分一份真实代码来得靠谱。我第一次升级 vLLM 版本后旧模型文件直接跑不起来就是靠这个习惯救回来的。第二个习惯权重转换脚本一定要保留而且要像对待核心代码一样对待它。很多模型发布新版本后权重结构会微调你必然要重跑转换。没有脚本就等着哭吧。转完立刻做一次端到端小规模验证用参考实现加载转换后的权重跑一个 greedy 输出存成黄金文件。后面 vLLM 接入、参数调整都以这个黄金文件为基准。第三个习惯调试阶段永远先把 enforce_eager 打开。CUDA Graph 捕获失败、形状不一致这类问题开了 eager 之后错误信息会清晰得多。等一切正常了再关掉不要一开始就纠结性能。最后一个小技巧我觉得比什么都实用给模型类留一个 debug 开关把前向过程里的中间 hidden states 导出来跟参考实现逐层对比。位置编码、mask、缩放系数这些 bug靠看最终输出根本发现不了逐层相似度一对比差异出在哪一层立刻现形。这一招帮我省掉的排查时间比写模型类本身还多。给 vLLM 加一个自定义模型本质上就是把“训练时的假设”翻译成“推理引擎的假设”。这个翻译过程没有太多魔法就是一个又一个参数、一层又一层的核对。只要流程清晰、验证到位小众模型也能吃上工业级推理引擎的红利。
模型训练收工的那一刻通常是最开心的接下去要把模型送上线才是最磨人的。手里一个小众模型训练时用的自定义框架checkpoint 命名随心所欲架构和主流模型有那么点像又不完全一样。你想用 vLLM 上线启动脚本直接报错“Model architectures are not supported”换回原项目自带的推理脚本单请求低吞吐批处理调度、显存管理全得自己写。于是只好硬着头皮打开 vLLM 源码研究怎么把一个自定义模型塞进去。这篇文章就是完整记录这件事的做法。我会拿一个完全虚构、但结构上很典型的小众因果语言模型 MiniStream 作为例子从权重转换、模型类编写、注册机制到启动验证把“从零给 vLLM 添加自定义模型推理支持”的完整链路捋一遍。适合三类人被“不支持的架构”卡住的部署工程师想把自己训练的模型接入 vLLM 的算法同学以及想搞懂 vLLM 模型适配机制的框架学习者。1. 为什么“加模型支持”这件事绕不开1.1 小众模型上线的两条死路先说一个现实vLLM 的高吞吐能力本质上来自它对模型结构做了大量假设和预优化。连续批处理、paged KV cache、量化、算子融合这些都是按主流 decoder-only 架构的规律去设计的。你的模型越是“长得像主流模型”接入成本越低越是个性化越要自己动手补适配层。小众模型常见的困境是两条路都不好走。第一条路是用原项目自带推理脚本上线。论文代码也好、内部实验代码也好大部分只保证“能跑”而不是“跑得快”。单条请求推理时GPU 利用率经常只有个位数百分比同时来十个请求内存翻倍、速度骤降。想支持流式输出、多轮对话、并发排队都得自己搭。最要命的是显存里的 KV cache 每次请求都要重新算长对话根本扛不住。这条路对小模型或许能撑模型一上亿参数就非常难受。第二条路是用通用推理框架跑但挑一个“看起来支持”的。很多框架只支持固定名单里的模型要么你的模型压根不在列表里要么得走 ONNX/TensorRT 导出精度和算子覆盖又容易出问题。好不容易导出成功模型里一个自定义位置编码可能就让整条链路废掉。折腾半天最后还是回到原点。相比之下直接给 vLLM 写一个模型适配文件一次性投入之后就能吃到 vLLM 整套调度和显存管理红利这是性价比最高的路径。只是这活儿有个门槛你得先搞清楚 vLLM 到底是怎么“认识”一个模型的。1.2 vLLM 认识模型的最小链条很多人以为 vLLM 支持某个模型是靠“扫描模型文件自动识别”其实完全不是。它内部维护了一张注册表这张表把“模型架构名”映射到“模型类”。你调用 LLM(model...) 时引擎会走一条固定的识别链路读取模型目录下的 config.json拿到 architectures 字段比如 [MiniStreamForCausalLM]。拿这个架构名去模型注册表里查看有没有注册过对应的类。查到了就用配置文件实例化这个模型类然后调用它的 load_weights 方法把权重灌进去。查不到直接抛“Model architectures are not supported”。所以“给 vLLM 加一个模型”这个任务拆开其实就是四件事写一个模型类、把它注册到注册表、保证配置文件能被正确解析、实现权重加载逻辑。其中权重加载通常是最耗时的因为你的 checkpoint 命名几乎一定和 vLLM 期望的不一样。这个链路听起来简单实际做起来处处是细节。下面把四个关键点一个个说清楚。2. 动手前必须摸清的四个关键点2.1 模型类长什么样先看一个最小可用的模型类长什么样。以 vLLM 0.6.x 这个相对稳定的版本为例模型类一般长这样import torch import torch.nn as nn from transformers import PretrainedConfig from vllm.model_executor.models.registry import ModelRegistry from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.sampler import Sampler ModelRegistry.register_model(MiniStreamForCausalLM) class MiniStreamForCausalLM(nn.Module): def __init__(self, config: PretrainedConfig, cache_configNone): super().__init__() self.config config self.model MiniStreamModel(config) self.lm_head ColumnParallelLinear( config.hidden_size, config.vocab_size, biasFalse ) self.sampler Sampler() def forward(self, input_ids, positions, kv_caches, attn_metadata, **kwargs): hidden_states self.model(input_ids, positions, kv_caches, attn_metadata) return hidden_states def compute_logits(self, hidden_states, sampling_metadata): return self.lm_head(hidden_states) def sample(self, logits, sampling_metadata): return self.sampler(logits, sampling_metadata) def load_weights(self, weights): # 核心工作量都在这 pass几个容易被忽略的细节forward 不直接返回 logits而是返回最后一层 hidden states。后续的 logits 计算和采样由 compute_logits、sample 这两个方法负责这是 vLLM 为了做采样处理、logits processor 而设计的拆分。init里那个 cache_config 参数有的版本有有的版本没有。保险写法是给它默认值 None避免不同版本实例化时报参数数量不匹配。注册装饰器里的字符串必须和 config.json 里 architectures[0] 完全一致大小写、空格都得一样否则注册了也找不到。2.2 权重映射才是真正的重头戏模型类说白了是个空壳真正让一个模型跑起来的是权重映射。vLLM 加载权重时会把 checkpoint 里的每个参数按名字喂给 load_weights你要做的就是把 checkpoint 名字翻译成你的模型参数名再把数据拷进去。最朴素的写法是def load_weights(self, weights): params_dict dict(self.named_parameters()) for item in weights: name getattr(item, name, item[0]) loaded_weight getattr(item, loaded_weight, item[1]) if name in params_dict: params_dict[name].data.copy_(loaded_weight)为什么这里是重头戏因为绝大多数自定义 checkpoint 不会按照“model.layers.0.self_attn.q_proj.weight”这种规范命名。我见过太多情况QKV 是融合成一个 Wqkv.weight 矩阵需要拆成 q_proj、k_proj、v_proj 三份门控 MLP 的两个权重拼在一起需要拆成 gate_proj 和 up_projembedding 和 lm_head 是权重绑定的checkpoint 里只存一份层编号从 1 开始而不是从 0 开始还有直接把 bias、scale 之类参数塞进同一个文件里的。你每遇到一种情况就要在 load_weights 里加一个分支去处理和翻译。如果模型有 30 多层每一层都这样处理分支代码会很长但逻辑本身并不复杂。关键是先把 checkpoint 的键名和模型参数名两边都打印出来做一次差集心里有数再动手。2.3 config.json 别靠猜vLLM 读配置的方式是先尝试用 transformers 生态的标准配置读取再拿着里面的字段去实例化模型类。所以 config.json 里的字段质量直接决定后面是否顺利。一份能用的 config.json 至少得包含这些{ architectures: [MiniStreamForCausalLM], model_type: ministream, hidden_size: 768, num_hidden_layers: 8, num_attention_heads: 12, num_key_value_heads: 12, head_dim: 64, intermediate_size: 2048, vocab_size: 32000, max_position_embeddings: 2048, rms_norm_eps: 1e-6, rope_theta: 10000.0, bos_token_id: 1, eos_token_id: 2, tie_word_embeddings: true }我这里特别提醒三个坑。第一个坑是 head_dim。很多模型 config 里不写 head_dim而是靠 hidden_size 除以 num_attention_heads 去推算。如果你的模型里 head_dim 和这个除法结果不一致就得显式写上否则注意力计算会直接形状不匹配。第二个坑是 architectures 字段。有些实验代码导出的 config 里architectures 写的是训练脚本里的类名比如 “MiniStreamTrainerModel”跟推理模型类名完全对不上。vLLM 只会拿这个字段去注册表里查名字对不上就是 not supported。第三个坑是 tie_word_embeddings。如果你的模型 embedding 和 lm_head 是共享权重的而 vLLM 模型实现里两者是独立参数那么加载时要人为复制一份权重。这个处理不好会出现“生成时第一个词概率异常”这类极难排查的问题。2.4 注意力能不能复用决定你后面省不省心给 vLLM 写模型类最大的性能分水岭在于注意力层怎么实现。如果你的模型用的是标准 RoPE MHA/GQA直接复用 vLLM 的 Attention 层就好paged KV cache、前缀缓存、滑动窗口这些能力全部自动生效。这是最理想的情况。如果你的模型有特殊的位置编码、特殊 attention mask 或者需要往 attention score 里加额外的 bias就得斟酌了。vLLM 的 Attention 层封装了下层各种后端不是所有自定义行为都支持。有人选择回到 PyTorch 原生实现比如 F.scaled_dot_product_attention这样功能上没问题模型照样能跑但你会失去 paged KV cache 的优势KV cache 显存占用和推理延迟都会明显变差。我的建议是分两步走第一版先保证能跑通哪怕用原生注意力实现跑通之后再回过头来看能不能把特殊注意力行为改造到 vLLM Attention 支持的范围内。别一上来就追求完美性能先让整条链路活起来。3. 实操把一个非标准 checkpoint 接进 vLLM3.1 样例模型MiniStream 与原始权重结构现在进入实战。假设我们手里有个内部训练的小模型 MiniStream参数量约 1.3 亿结构不复杂standard decoder-only8 层 Transformerhidden size 76812 个注意力头SwiGLU 激活RMSNormRoPE 位置编码。结构上一点都不特殊但 checkpoint 格式非常“实验风”。原始 checkpoint 是单个 model.pt 文件键名长这样embedding.weight形状 [32000, 768]blocks.0.attn.Wqkv.weight形状 [2304, 768]Q、K、V 按行顺序拼在一起blocks.0.attn.out_proj.weight形状 [768, 768]blocks.0.mlp.w12.weight形状 [4096, 768]SwiGLU 的两个门控矩阵拼在一起blocks.0.mlp.w3.weight形状 [768, 2048]blocks.0.norm1.weight、blocks.0.norm2.weight形状 [768]没有独立的 lm_headembedding 与输出头共享权重也就是说训练框架导出时完全按“角色”命名不关心 transformers 生态的命名规范。我们要做的第一件事就是写个脚本把它翻译成 vLLM 能消化的 HF 风格权重。3.2 第一步写权重转换脚本转换脚本的逻辑很简单逐个键名做映射。需要特别注意的是 Wqkv 融合矩阵和 w12 融合矩阵的处理方式。import torch from safetensors.torch import save_file raw torch.load(mini-stream-132m/model.pt, map_locationcpu) raw raw[model] if model in raw else raw mapped {} mapped[model.embed_tokens.weight] raw[embedding.weight] mapped[lm_head.weight] raw[embedding.weight] # tied 权重存两份 hidden_size 768 mid hidden_size # 因为 MiniStream 是 MHAQ/K/V 各占 hidden_size intermediate 2048 for i in range(8): src fblocks.{i} dst fmodel.layers.{i} fused_qkv raw[f{src}.attn.Wqkv.weight] # [2304, 768] q, k, v torch.split(fused_qkv, [hidden_size, hidden_size, hidden_size], dim0) mapped[f{dst}.self_attn.q_proj.weight] q mapped[f{dst}.self_attn.k_proj.weight] k mapped[f{dst}.self_attn.v_proj.weight] v mapped[f{dst}.self_attn.o_proj.weight] raw[f{src}.attn.out_proj.weight] w12 raw[f{src}.mlp.w12.weight] # [4096, 768]对应 gate_up_proj mapped[f{dst}.mlp.gate_up_proj.weight] w12 mapped[f{dst}.mlp.down_proj.weight] raw[f{src}.mlp.w3.weight] mapped[f{dst}.input_layernorm.weight] raw[f{src}.norm1.weight] mapped[f{dst}.post_attention_layernorm.weight] raw[f{src}.norm2.weight] save_file(mapped, mini-stream-hf/model.safetensors)转完之后别急着进 vLLM先把 config.json 和 tokenizer 相关文件准备好目录结构如下mini-stream-hf/ ├── config.json ├── model.safetensors ├── tokenizer.json └── tokenizer_config.json这一步完成后最好先用 transformers 生态的参考实现把转换后的权重跑一遍生成一个 greedy 结果存下来。这叫“黄金输出”后面验证 vLLM 接入是否正确就用它对比。3.3 第二步写模型类并注册模型类的完整代码会包含三个部分基础模块注意力、MLP、层、顶层模型、以及注册入口。MiniStream 结构不特殊注意力可以直接用 vLLM 的 Attention 层。下面是完整实现思路。from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear from vllm.model_executor.layers.vocab_parallel_embeddings import VocabParallelEmbedding from vllm.attention.layer import Attention class MiniStreamAttention(nn.Module): def __init__(self, config): super().__init__() self.hidden_size config.hidden_size self.num_heads config.num_attention_heads self.head_dim config.head_dim self.q_proj ColumnParallelLinear(self.hidden_size, self.hidden_size, biasFalse) self.k_proj ColumnParallelLinear(self.hidden_size, self.hidden_size, biasFalse) self.v_proj ColumnParallelLinear(self.hidden_size, self.hidden_size, biasFalse) self.o_proj RowParallelLinear(self.hidden_size, self.hidden_size, biasFalse) self.attn Attention(self.num_heads, self.head_dim) def forward(self, hidden_states, positions, kv_cache, attn_metadata): q, _ self.q_proj(hidden_states) k, _ self.k_proj(hidden_states) v, _ self.v_proj(hidden_states) q q.view(-1, self.num_heads, self.head_dim) k k.view(-1, self.num_heads, self.head_dim) v v.view(-1, self.num_heads, self.head_dim) out self.attn(q, k, v, kv_cache, attn_metadata) out out.view(-1, self.hidden_size) out, _ self.o_proj(out) return out class MiniStreamMLP(nn.Module): def __init__(self, config): super().__init__() self.gate_up_proj ColumnParallelLinear( config.hidden_size, 2 * config.intermediate_size, biasFalse ) self.act_and_mul SiluAndMul() self.down_proj RowParallelLinear( config.intermediate_size, config.hidden_size, biasFalse ) def forward(self, x): x, _ self.gate_up_proj(x) x self.act_and_mul(x) x, _ self.down_proj(x) return x class MiniStreamDecoderLayer(nn.Module): def __init__(self, config): super().__init__() self.input_layernorm RMSNorm(config.hidden_size, epsconfig.rms_norm_eps) self.self_attn MiniStreamAttention(config) self.post_attention_layernorm RMSNorm(config.hidden_size, epsconfig.rms_norm_eps) self.mlp MiniStreamMLP(config) def forward(self, hidden_states, positions, kv_cache, attn_metadata): residual hidden_states hidden_states self.input_layernorm(hidden_states) hidden_states self.self_attn(hidden_states, positions, kv_cache, attn_metadata) hidden_states residual hidden_states residual hidden_states hidden_states self.post_attention_layernorm(hidden_states) hidden_states self.mlp(hidden_states) hidden_states residual hidden_states return hidden_states class MiniStreamModel(nn.Module): def __init__(self, config): super().__init__() self.embed_tokens VocabParallelEmbedding( config.vocab_size, config.hidden_size ) self.layers nn.ModuleList([ MiniStreamDecoderLayer(config) for _ in range(config.num_hidden_layers) ]) self.norm RMSNorm(config.hidden_size, epsconfig.rms_norm_eps) def forward(self, input_ids, positions, kv_caches, attn_metadata): hidden_states self.embed_tokens(input_ids) for i, layer in enumerate(self.layers): hidden_states layer(hidden_states, positions, kv_caches[i], attn_metadata) hidden_states self.norm(hidden_states) return hidden_states现在把最关键的 load_weights 补上。因为转换脚本已经把 fused 的 Wqkv 拆成了 q/k/v 三份所以 load_weights 里大部分情况是直接拷贝只有两个地方需要特殊处理lm_head 是 tied 的embedding 和 lm_head 用同一份权重还有 gate_up_proj 这个参数名要认出来。def load_weights(self, weights): params_dict dict(self.named_parameters()) for item in weights: name getattr(item, name, item[0]) loaded_weight getattr(item, loaded_weight, item[1]) if name model.embed_tokens.weight: params_dict[model.embed_tokens.weight].data.copy_(loaded_weight) params_dict[lm_head.weight].data.copy_(loaded_weight) continue if name in params_dict: params_dict[name].data.copy_(loaded_weight)为什么可以这么简单因为我们在转换脚本里已经把命名统一成了 vLLM 模型类里的参数名。这是最推荐的做法能通过转换脚本解决的问题就不要在 load_weights 里堆分支。注册方式有两种。一种是直接把 ministream.py 放进 vLLM 源码目录 vllm/model_executor/models/ 下并在该目录的init.py 里加一行 import。这个办法最省事但每次升级 vLLM 都要重新打补丁。另一种是写成独立的 Python 包在调用 LLM 之前的启动脚本里先 import 你的模型模块让 ModelRegistry.register_model 装饰器在引擎启动前完成注册。第二种更干净我后来一直用这种。3.4 第三步启动推理与正确性验证代码写完之后直接用 vLLM 启动from vllm import LLM, SamplingParams llm LLM( model./mini-stream-hf, dtypebfloat16, max_model_len2048, enforce_eagerTrue, # 调试阶段先关掉 CUDA Graph gpu_memory_utilization0.8, ) prompts [请用一句话介绍你自己。, 解释一下什么是 KV Cache。] outputs llm.generate(prompts, SamplingParams(max_tokens64, temperature0.2)) for o in outputs: print(o.outputs[0].text)如果这里能正常输出恭喜最难的 90% 已经过了。但“能输出”不等于“输出正确”。我强烈建议做一次黄金输出对比拿转换后的 model.safetensors 加载进参考实现用同一个 prompt、同样的采样参数最好 temperature0 的 greedy生成一份结果再和 vLLM 的输出逐 token 对比。我第一次接入模型时输出看起来通顺但首 token 命中率只有七成左右。后来发现是注意力缩放系数里用了 hidden_size 的平方根而原模型用的是 head_dim 的平方根。这类问题靠肉眼几乎不可能发现只能靠机械对比。4. 验证跑通之后性能与版本差异4.1 性能表现怎么看正确性过了之后把 enforce_eager 关掉恢复 CUDA Graph再重新跑一遍。这时候你会看到吞吐明显提升。再检查几个点第一确认你的注意力层确实走到了 vLLM 的 Attention。MiniStream 是标准 MHA所以 paged KV cache、前缀缓存都能直接用。如果你的模型用了 SDPA 兜底性能会差一大截那就得回头想想怎么把特殊结构标准化。第二看 KV cache 显存占用。vLLM 启动时会自动计算 KV cache 大小但如果你的 max_model_len 设得很大KV cache 会吃掉大量显存。可以从 gpu_memory_utilization、max_num_seqs、max_model_len 三个参数入手调整。第三想上多卡张量并行的话load_weights 的写法要升级。上面的例子在单卡下没问题多卡时 ColumnParallelLinear 的参数会被切分你需要处理每个 rank 拿到对应分片。最简单的方式是参考同版本 vLLM 官方模型文件里对 weight_loader 的处理把每个参数的加载交给 weight_loader而不是直接对 param.data 整体 copy。这里我建议直接抄同版本 LLaMA 类模型的做法别自己发明。4.2 不同 vLLM 版本的适配差异vLLM 迭代速度快得离谱不同版本的接口差异是我踩过最多坑的地方。核心思想是你的模型文件最好基于你部署版本里某个现成模型文件做“差分”而不是拿着老博客的代码直接套。版本区间主要差异与建议0.4.xSampler 和模型调度接口比较老模型文件结构与 0.6 差异较大按同版本现成模型改0.5/0.6.x本文的写法基本适用ModelRegistry.register_model 和 vLLM Attention 层是稳定接口0.7/0.8引入 SupportsV0/SupportsV1 协议模型类需要实现对应接口load_weights 里拿到的权重对象带 shard 信息1.x 及更新默认 V1 调度接口进一步收敛插件化机制更完善推荐把自定义模型模块打成独立包注册不管哪个版本load_weights 的“语义”从来没变过上游把 checkpoint 权重一个接一个传给你你负责翻译成模型参数并拷贝。只要抓住这个核心版本升级对你来说只是换接口写法不是换思路。5. 常见问题与排查实录5.1 注册了还是 not supported这是第一个容易翻车的点。排查顺序如下先确认 config.json 里 architectures 的值和注册时的字符串一字不差包括大小写。再确认你的模型模块确实被 import 了。如果你用的是自定义包的方式很常见的情况是启动脚本里忘了 import装饰器根本没执行注册表里自然没有。最后确认 vLLM 的模型注册表是否已经包含你的类可以在代码里临时打印出来看一眼。5.2 权重加载各种 mismatch权重相关报错是最常见的本质都是“vLLM 模型参数名”和“checkpoint 键名”对不上。解决手段是写一个小脚本同时打印两边的键名# 打印 checkpoint 键名 from safetensors import safe_open with safe_open(mini-stream-hf/model.safetensors, frameworkpt) as f: print(\n.join(sorted(f.keys()))) # 打印模型参数名 model MiniStreamForCausalLM(config) print(\n.join(sorted(dict(model.named_parameters()).keys())))对着两份名单做差集哪个名字少、哪个名字多一目了然。如果你看到 size mismatch 而不是 name mismatch那通常是某个维度算错了比如 QKV 拆分的 hidden_size 写错或者词表大小因 padding 和实际权重不一致。5.3 显存爆了与生成变慢如果是加载阶段就爆显存优先怀疑 dtype。默认加载 fp32 很容易爆改成 bfloat16 或 float16 会让显存直接减半。如果加载没问题一跑生成就爆重点查 max_model_len 和 gpu_memory_utilization。调试阶段还可以加上 enforce_eager把 CUDA Graph 的显存开销省掉。如果是生成速度越来越慢而不是一开始就慢大概率是 KV cache 没有被正确复用。检查你的注意力实现是不是每次 forward 都重新计算了历史 token。如果模型用了自定义 attention 兜底这个问题尤其常见。5.4 输出和原模型对不上输出不一致时按这个优先级排查dtype 不一致比如原模型是 fp32vLLM 里用了 bf16某些层会有一点点漂移但不会离谱。RMSNorm 的 eps。很多模型这个值是 1e-6 或 1e-5差一个数量级小概率出现明显偏差。RoPE 相关参数包括 rope_theta、rope_scaling甚至 max_position_embeddings。位置编码类的 bug 往往在长序列才暴露短 prompt 看不出问题。注意力缩放系数。检查注意力分数乘的是 head_dim 的倒数平方根还是 hidden_size 的倒数平方根。词表 padding。如果你的词表大小不是 32 的倍数某些实现会 padding 到对齐导致 lm_head 的权重形状和 embedding 对不上。以上都排完了还有一个大招在模型类里加一个临时 debug 钩子把每一层的 hidden states 导出成文件和参考实现逐层做余弦相似度。哪一层开始分叉问题就定位在哪一层。6. 我踩过的坑与几个保命习惯最后说几句实在的都是真金白银踩出来的经验。第一个习惯动手写模型类之前先到你的 vLLM 版本里找结构最接近的官方模型文件复制一份然后基于这份拷贝去改。版本接口再怎么变化你手里的官方文件就是最准确的“标准答案”。网上教程再详细也不如你自己差分一份真实代码来得靠谱。我第一次升级 vLLM 版本后旧模型文件直接跑不起来就是靠这个习惯救回来的。第二个习惯权重转换脚本一定要保留而且要像对待核心代码一样对待它。很多模型发布新版本后权重结构会微调你必然要重跑转换。没有脚本就等着哭吧。转完立刻做一次端到端小规模验证用参考实现加载转换后的权重跑一个 greedy 输出存成黄金文件。后面 vLLM 接入、参数调整都以这个黄金文件为基准。第三个习惯调试阶段永远先把 enforce_eager 打开。CUDA Graph 捕获失败、形状不一致这类问题开了 eager 之后错误信息会清晰得多。等一切正常了再关掉不要一开始就纠结性能。最后一个小技巧我觉得比什么都实用给模型类留一个 debug 开关把前向过程里的中间 hidden states 导出来跟参考实现逐层对比。位置编码、mask、缩放系数这些 bug靠看最终输出根本发现不了逐层相似度一对比差异出在哪一层立刻现形。这一招帮我省掉的排查时间比写模型类本身还多。给 vLLM 加一个自定义模型本质上就是把“训练时的假设”翻译成“推理引擎的假设”。这个翻译过程没有太多魔法就是一个又一个参数、一层又一层的核对。只要流程清晰、验证到位小众模型也能吃上工业级推理引擎的红利。
到O(n))
平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测)
