Daft × vLLM 批量推理基准:从朴素 Batch 到前缀路由的持续批处理优化实践
Daft × vLLM 批量推理基准从朴素 Batch 到前缀路由的持续批处理优化实践【免费下载链接】DaftHigh-performance data engine for AI and multimodal workloads. Process images, audio, video, and structured data at any scale项目地址: https://gitcode.com/GitHub_Trending/da/Daft在大规模 LLM 推理场景中怎么把 prompt 喂给推理引擎直接决定了吞吐表现。Daft 仓库的 benchmarking/vllm 目录提供了一组完整的对比基准它用同一份 Parquet 数据集、同一个 Qwen3-8B 模型分别跑朴素批量推理LLM.generate、朴素批量 prompt 排序、vLLM 持续批处理关闭前缀路由、持续批处理 排序、持续批处理 前缀分桶路由以及Ray Data LLM Processor共六种方案帮助你在 Ray 集群上量化各种推理编排策略的差异。读完本文你将掌握这套基准的搭建流程、六种脚本的执行机制以及 Daftprompt()函数背后前缀路由prefix routing与持续批处理continuous batching的实现原理。一、基准环境搭建根据 README 的 Setup 说明搭建步骤如下安装依赖vllm、daft、ray三个包生成数据集运行 generate_data.ipynb 笔记本生成基准数据修改配置更新 config.py 中的数据集路径与其他参数在 Ray 集群上运行各基准脚本。1.1 基准数据是如何生成的generate_data.ipynb 使用 vLLM 自带的PrefixRepetitionRandomDataset采样器以Qwen/Qwen3-8B的 tokenizer 生成固定形状的合成 promptrandom_seed42保证可复现并通过 Daft 写入 S3 上的 Parquet 文件repartition(8)后写出。笔记本中定义了两类数据生成函数from vllm.benchmarks.datasets import PrefixRepetitionRandomDataset from vllm.transformers_utils.tokenizer import get_tokenizer import daft from daft.functions import monotonically_increasing_id import ray ray.init() daft.set_runner_ray() dataset PrefixRepetitionRandomDataset(random_seed42) tokenizer get_tokenizer(Qwen/Qwen3-8B) root_dir s3://my-bucket/vllm-prefix-caching-partitioned def generate_data(request_k: int, num_prefixes: int): sample dataset.sample( tokenizertokenizer, num_requestsrequest_k * 1000, # 请求数量 prefix_len256, # 前缀长度 suffix_len256, # 后缀长度 num_prefixesnum_prefixes, # 共享前缀的数量 output_len128, # 期望输出长度 ) sample [s.prompt for s in sample] df daft.from_pydict({prompt: sample}) df df.select(monotonically_increasing_id().alias(id), prompt) df df.repartition(8) return df.write_parquet(f{root_dir}/{request_k}k_0-5_{num_prefixes}.parquet) def generate_data_no_prefix(request_k: int): # 无前缀版本prefix_len0, suffix_len512 ...实际执行的单元格生成了 5 个数据集generate_data(200, 8)、generate_data(200, 64)、generate_data(200, 512)、generate_data(200, 4096)和generate_data_no_prefix(200)。可以看出数据的变量设计就是num_prefixes共享前缀的数量前缀越少prompt 间相似度越高前缀缓存和分桶路由的潜在收益越大generate_data_no_prefix则提供一个完全没有共享前缀的对照组。这也是为什么 config.py 中默认的INPUT_PATH指向200k_0-5_512.parquet——跑其他数据集时需要手动替换路径。1.2 全局配置 config.pyconfig.py 是所有脚本共享的配置中心内容非常精简from vllm import SamplingParams MODEL_NAME Qwen/Qwen3-8B INPUT_PATH s3://my-bucket/vllm-prefix-caching-partitioned/200k_0-5_512.parquet OUTPUT_LEN 128 SAMPLING_PARAMS SamplingParams(min_tokensOUTPUT_LEN, max_tokensOUTPUT_LEN) CONCURRENCY 128 def print_benchmark_results(script: str, start_time: float, end_time: float): ... # 打印脚本名、执行耗时及上述全部配置几个关键设计点SamplingParams(min_tokensOUTPUT_LEN, max_tokensOUTPUT_LEN)将输出长度钉死在 128 token消除采样长度差异对耗时对比的干扰保证各脚本输出量一致CONCURRENCY 128是统一的并发上限用于对齐不同方案的 actor/请求并发度print_benchmark_results在每次运行结束时打印配置摘要方便把结果与参数一一归档。二、六个基准脚本的机制对比2.1 naive-batch.py朴素批量推理基线naive-batch.py 是整套基准的对照组它不依赖 Daft 的 AI 函数而是用daft.cls直接封装 vLLM 的离线LLM引擎daft.cls(max_concurrencyCONCURRENCY, gpus1) class VLLM: def __init__(self): self.llm LLM( modelMODEL_NAME, max_model_len4096, disable_log_statsFalse, ) daft.method.batch(return_dtypestr, batch_size512) def generate(self, prompts: Series) - Series: outputs self.llm.generate(prompts.to_pylist(), SAMPLING_PARAMS) self.llm.llm_engine.do_log_stats() # 打印 vLLM 内部吞吐统计 return Series.from_pylist([o.outputs[0].text for o in outputs]) def main(): daft.set_runner_ray() df daft.read_parquet(INPUT_PATH).into_partitions(32) vllm VLLM() df df.with_column(output, vllm.generate(df[prompt])) df df.collect() ...其执行模型是Ray 集群上最多拉起 128 个 actormax_concurrency128每个占 1 张 GPU每个 actor 持有一个完整的LLM引擎实例Daft 把每 512 条 promptbatch_size512攒成一个 batch 后一次性调用llm.generate。这就是典型的静态批static batching——一个 batch 内的请求同进同出短请求要等最长请求生成完才能出队且批与批之间的空隙没有复用。脚本里刻意保留了disable_log_statsFalse并手动调用do_log_stats()就是为了在 stdout 里拿到 vLLM 官方的吞吐指标如 token/s、GPU KV cache 利用率作为第三方旁证。2.2 naive-batch-sorted.py朴素批量 prompt 排序naive-batch-sorted.py 与上文唯一的区别是在推理前加了一行df df.sort(prompt) df df.with_column(output, vllm.generate(df[prompt]))按 prompt 字典序排序后共享相同前缀的 prompt 会在物理上聚拢到一起。这个排序本身无法改善 vLLM 静态批内部的调度但它显著提升了 vLLM 引擎内置 prefix cache 的命中率——相邻请求前缀相同KV cache 更可能被复用。它与naive-batch.py的差值恰好量化了数据排布对缓存友好度的影响。2.3 continuous-batch.py持续批处理关闭前缀路由continuous-batch.py 换用了 Daft 的 AI 函数入口daft.functions.prompt和vllm-prefix-cachingproviderdf daft.read_parquet(INPUT_PATH).into_partitions(8) df df.with_column( output, prompt( df[prompt], providervllm-prefix-caching, modelMODEL_NAME, engine_args{max_model_len: 4096}, generate_args{sampling_params: SAMPLING_PARAMS}, concurrencyCONCURRENCY, do_prefix_routingFalse, # 关键关闭前缀路由 batch_size512, ), )相比 naive 版本参数从自己写一个 UDF 类收敛成了prompt()的声明式配置。do_prefix_routingFalse意味着 Daft 这一侧不做前缀分桶调度请求按到达顺序进入各 vLLM actorvLLM 引擎自身仍开启 prefix caching因此这代表只有引擎层持续批处理的基线。2.4 continuous-batch-sorted.py持续批处理 排序continuous-batch-sorted.py 与 2.3 完全相同仅多了一行df df.sort(prompt)。对比它可以回答在引擎自身已有 prefix cache 的前提下额外做全局排序还能带来多少边际收益。2.5 prefix-bucketing.py持续批处理 前缀分桶路由prefix-bucketing.py 是整套基准的主角它相对 2.3 的差异是去掉了do_prefix_routingFalse和batch_size512即启用 provider 的默认前缀路由行为prompt( df[prompt], providervllm-prefix-caching, modelMODEL_NAME, engine_args{max_model_len: 4096}, generate_args{sampling_params: SAMPLING_PARAMS}, concurrencyCONCURRENCY, )这个脚本与continuous-batch.py、continuous-batch-sorted.py一起构成一个完整的 2×2 实验矩阵{是否前缀路由} × {是否排序}从而分离出两种优化手段各自的贡献。2.6 ray-data.pyRay Data LLM Processor 对照组ray-data.py 完全不使用 Daft直接用 Ray Data 的 LLM 组件跑同样的数据作为外部框架的参照系config vLLMEngineProcessorConfig( model_sourceMODEL_NAME, engine_kwargsdict( enable_prefix_cachingTrue, max_model_len4096, ), batch_size16, concurrencyCONCURRENCY, ) processor build_llm_processor(config) ds ray.data.read_parquet(INPUT_PATH) ds ds.map(lambda row: {messages: [{role: user, content: row[prompt]}], sampling_params: SAMPLING_PARAMS}) output_ds processor(ds) output_ds.take_all()值得注意它的批大小是batch_size16Ray Data LLM Processor 的典型取值与 Daft 侧的 512 不同且它把 prompt 包成 OpenAI 风格的messages结构。这个脚本用于验证 Daft 的优化策略相对 Ray Data 官方方案的相对位置而不是做严格同参数对比。三、源码纵深vLLM provider 与前缀路由是怎么实现的3.1 prompt() 的双执行路径从源码结构看prompt()函数daft/functions/ai/init.py在解析 provider 后会判断描述符类型一旦是VLLMPrefixCachingPrompterDescriptor就不走普通的 UDF 执行路径而是直接调用底层的PyExpr.vllm()把模型名和全部路由参数透传给 Rust 侧# daft/functions/ai/__init__.py if isinstance(prompter_descriptor, VLLMPrefixCachingPrompterDescriptor): ... vllm_options prompter_descriptor.get_options() return Expression._from_pyexpr( messages._expr.vllm( prompter_descriptor.model_name, vllm_options[concurrency], vllm_options[gpus_per_actor], vllm_options[do_prefix_routing], vllm_options[max_buffer_size], vllm_options[min_bucket_size], vllm_options[prefix_match_threshold], vllm_options[load_balance_threshold], vllm_options[batch_size], vllm_options[engine_args], vllm_options[generate_args], ) )同样的分支里还明确了该路径的限制vLLM provider 不支持return_format结构化输出、不支持system_message、不支持多消息列表——因为 prompt 被直接作为纯文本送入推理引擎。3.2 VLLMPromptOptions完整的参数清单与默认值provider 的完整参数定义在 daft/ai/vllm/protocols/prompter.py 的VLLMPromptOptions中其默认值由get_options()给出。基准脚本只显式用了其中一部分完整可配项如下参数默认值含义concurrency1并行 vLLM actor 数量上限基准中设为 128gpus_per_actor1每个 actor 占用的 GPU 数do_prefix_routingTrue是否启用前缀分桶路由max_buffer_size5000前缀路由缓冲区容量min_bucket_size16前缀分桶的最小桶大小prefix_match_threshold0.33判定两条 prompt 前缀相似的阈值load_balance_threshold256负载均衡触发阈值batch_sizeNone攒批大小基准中为 512engine_args{}透传给 vLLM 引擎的参数如max_model_lengenerate_args{}透传给 generate 的参数如sampling_paramsnum_gpus1actor 的 GPU 资源声明这些参数解释了基准脚本行为的差异来源continuous-batch.py显式传do_prefix_routingFalse关闭路由而prefix-bucketing.py不传该参数、落回默认值True于是请求会被按前缀相似度分桶后再路由到对应 actor让共享前缀尽量落在同一批被处理最大化 KV cache 命中。engine_args/generate_args两个字典则实现了配置与引擎 API 的解耦——基准中的max_model_len4096和定长SamplingParams都是经此通道注入的。3.3 Provider 的实验性定位vLLM provider 的类文档 中有一句明确声明VLLMPrefixCachingProviderhighly experimental and may not work且 API 近期可能变化或被移除。使用这套基准时应当意识到脚本中的参数尤其do_prefix_routing、max_buffer_size等与该实验性 API 强绑定若 Daft 升级后行为变化需要以当前仓库版本的 provider 实现为准重新核对。四、如何运行并解读结果在 Ray 集群上确保各节点已安装vllm、daft、ray且能访问 S3 数据集INPUT_PATH与 Hugging Face 模型Qwen/Qwen3-8B若使用自建数据修改 config.py 中的INPUT_PATH建议同时调整OUTPUT_LEN与数据集生成时的output_len保持一致以维持输出长度可控依次运行六个脚本例如python benchmarking/vllm/naive-batch.py、python benchmarking/vllm/prefix-bucketing.py等每个脚本结束时都会打印BENCHMARK RESULTS块包含执行耗时和当次配置naive 系列还会在 stdout 中打印 vLLM 的BATCH INFERENCE STATS吞吐、KV cache 利用率等。解读建议naive-batch→naive-batch-sorted的差值静态批下prompt 排布对前缀缓存的影响continuous-batch→continuous-batch-sorted的差值引擎持续批处理已开启时排序的边际收益预期显著收窄;continuous-batch或 sorted 版→prefix-bucketing的差值Daft 侧前缀分桶路由带来的增益且该增益会随数据集num_prefixes增大prompt 更相似而放大——这正是generate_data生成 8/64/512/4096 多档数据集、以及无前缀对照组存在的意义任意 Daft 方案 vsray-data跨框架的相对参考注意二者batch_size不同对比时应在结果记录中注明。五、总结benchmarking/vllm 用六个结构高度对称的脚本把批量 LLM 推理性能优化拆成了可独立度量的变量静态批 vs 持续批、是否排序、是否前缀路由、Daft vs Ray Data。其底层支撑是daft.functions.prompt针对vllm-prefix-cachingprovider 的专用执行路径PyExpr.vllm以及 VLLMPromptOptions 中暴露的一套前缀路由调参接口。由于该 provider 目前被标注为高度实验性且基准脚本依赖 S3 数据集与 GPU Ray 集群复现时需要以当前仓库版本的 provider 实现和config.py参数为准。【免费下载链接】DaftHigh-performance data engine for AI and multimodal workloads. Process images, audio, video, and structured data at any scale项目地址: https://gitcode.com/GitHub_Trending/da/Daft创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
上一篇/下一篇内容由系统自动关联
返回资讯列表 →