加速框架学习
src/rocket/runners.py代码解读
这份文档用于解释 Rocket 框架里最核心的运行器文件：
Rocket/src/rocket/runners.py这个文件的职责不是读取数据集、计算指标或写结果文件，而是把同一批 prompt 交给不同推理后端执行，并把 HF、vLLM、SGLang 的输出统一成同一种格式，方便上层 evaluator 做公平对比。
1. 文件定位
runners.py是 Rocket 的“推理后端适配层”。它向上对接evaluator.py，向下调用具体推理引擎：
evaluator.py -> build_backend(config) -> HfBackend -> VllmTextBackend -> SglangTextBackend -> backend.load() -> backend.generate(...) -> backend.runtime_metadata() -> backend.shutdown()因此，runners.py解决的是一个工程问题：不同推理框架的参数、加载方式、输出格式都不一样，但实验表格需要同一套指标。这个文件把这些差异包起来，让上层只关心“输入 prompt、得到文本和 token 数”。
2. 核心数据结构
2.1Generation
@dataclass(slots=True)classGeneration:text:strgenerated_tokens:int这是所有后端统一返回的单条生成结果。
text：模型生成的文本。generated_tokens：本次生成的 token 数，用于计算吞吐。
HF、vLLM、SGLang 原生返回结构都不同：
- HF 返回
output_idstensor。 - vLLM 返回 request output 对象。
- SGLang 返回 dict/list/string 等不同版本格式。
Rocket 最终都转换成Generation，这样 evaluator 不需要理解每个框架的内部结构。
2.2BackendConfig
@dataclass(slots=True)classBackendConfig:model_path:strmethod:MethodSpec batch_size:intmax_model_len:intmax_new_tokens:inttemperature:floattop_p:floatseed:intdtype:strmax_num_batched_tokens:intgpu_memory_utilization:floatsglang_mem_fraction_static:floatinput_mode:str="auto"torch_compile_cache_size_limit:int=16torch_compile_mode:str="reduce-overhead"这是所有后端共享的运行配置。
这里的设计重点是“变量统一”：同一个实验方法用同一批基础参数，比如模型路径、最大生成长度、batch size、采样温度、top-p 等。后端内部再把这些参数翻译成自己框架能理解的形式。
几个关键字段：
method：来自methods.py的MethodSpec，决定使用 HF/vLLM/SGLang、是否开 Rocket bucket、是否开 n-gram。max_model_len：模型总上下文长度上限。max_new_tokens：最大生成 token 数。max_num_batched_tokens：vLLM/SGLang 的 chunked prefill 相关参数。gpu_memory_utilization：vLLM 预留显存比例。sglang_mem_fraction_static：SGLang 静态显存比例。input_mode：auto/text/vision，用于 HF 路径区分纯文本模型和视觉语言模型。torch_compile_cache_size_limit：HF + Rocket 使用torch.compile时的图缓存上限。
3. 后端工厂函数
defbuild_backend(config:BackendConfig):ifconfig.method.engine=="hf":returnHfBackend(config)ifconfig.method.engine=="vllm":returnVllmTextBackend(config)ifconfig.method.engine=="sglang":returnSglangTextBackend(config)这个函数根据method.engine创建具体后端。
例如：
hf_sdpa -> HfBackend hf_rocket -> HfBackend vllm_fa -> VllmTextBackend vllm_rocket_flash -> VllmTextBackend sglang_flashinfer -> SglangTextBackend注意：Rocket 不做“自动选择最快后端”的路由。这里的每个 method 都是固定后端、固定配置，用于公平实验对比。
4. HF 运行时兼容补丁
函数：
_patch_hf_runtime_for_installed_packages(attn_implementation)这个函数只服务 HF 后端，目的是让当前环境里的 Transformers 能正常加载模型。
它不安装依赖，不卸载依赖，也不修改外部包文件，只在当前 Python 进程里做运行时兼容。
4.1 为什么需要 patchtorchao
当前环境里transformers 5.8.1可能会尝试导入 torchao quantizer。问题是安装的 torchao 版本依赖更高版本 PyTorch 里的：
torch.utils._pytree.register_consta