跳转至

Justitia: 任务并行大模型智能体的公平高效调度与选择性饱和服务

会议: NeurIPS 2026
arXiv: 2510.17015
领域: LLM 安全
关键词: 大模型智能体调度, 任务并行, 显存瓶颈建模, 选择性饱和服务, 虚拟时间公平队列

一句话总结

针对多任务并行大模型智能体在共享 GPU 集群中的调度冲突,Justitia 提出基于虚拟时间公平完成顺序的「选择性饱和服务」(Selective Pampering)机制,结合显存时空占用的二次方成本建模与轻量级单智能体 MLP 预测,在严格保障最坏完成延迟有界的前提下,将平均作业完成时间缩短了 57.5%。

研究背景与动机

随着大语言模型在复杂长流程任务中的广泛应用,利用包含多个并发推理任务的「任务并行智能体」(Task-parallel LLM Agents)解决实际问题已成为主流范式。典型的现实场景包括大文档分块并行总结(MapReduce Summarization)、多文档并行融合打分(Document Merging)、长链条幻觉声明并行检索核验(Fact Verification)以及多分支采样多数投票(Self-Consistency)。在多租户共享的 GPU 集群中,智能体的最终结果取决于其最后一个子任务的完成,因此整体端到端作业完成时间(Job Completion Time, JCT)是用户体验的核心指标;同时,集群还需要防止恶意或长任务挤占资源,保障租户间的服务公平性。

然而,现有 LLM 服务系统在面向任务并行智能体时面临着效率与公平难以兼顾的严峻困境。主流系统如 vLLM 和 SGLang 采用先来先服务(FCFS)策略,极易引发队头阻塞(Head-of-Line Blocking);虽然引入最短作业优先(SJF)或多级反馈队列(MLFQ)能够加速短任务并提升吞吐,但会使大规模「大象智能体」遭受无限期饥饿(Starvation);而像 VTC 和 FairServe 等追求瞬时公平份额(Instantaneous Fairness)的公平调度器,则将每个智能体的瞬时并发度强行限制在平均切片内,导致原本可并发加速的子任务被强行拉长,损害了端到端 JCT。

深入分析可以发现,用户对公平性的本质诉求是「最坏完成时间有确定性保障」的长期完成时间公平(Finish-time Fairness),而非每个毫秒都机械均分算力。如果能在理想公平共享算出的预期完成顺序基础上,打破瞬时配额限制,允许排在前面的智能体集中占用全部资源快速跑完并及早释放显存,就能在不推迟任何智能体理论完工时间的前提下大幅降低平均 JCT。核心 idea:提出选择性饱和服务(Selective Pampering)机制,以广义处理器共享(GPS)下的虚拟完成时间为优先级,配合显存时空占用(KV Token-Time)的二次方成本建模与轻量 MLP 需求预测,实现兼具近最优吞吐与常数级最坏延迟保障的智能体调度。

方法详解

整体框架

Justitia 针对多租户共享 GPU 环境下任务并行智能体端到端执行的特点,设计了由「轻量预测器」、「虚拟时间调度器」与「非抢占式推理执行引擎」构成的流水线。当用户提交包含多个并行子任务的智能体时,系统首先提取输入 Prompt 并通过 TF-IDF 向量化输入至对应智能体类型的专用 4 层 MLP,预测其全生命周期总显存开销;随后,调度器根据系统全局虚拟时间 \(V(t)\) 瞬时计算该智能体的虚拟完工时间标号 \(\bar{f}_j\),并将其作为其所有子推理请求的全局队列优先级;最后,执行引擎按该优先级非抢占式批量派发推理任务至底层 vLLM 实例,在不触发高开销 KV Cache 内存换入换出的前提下实现单智能体的饱和执行。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["用户提交任务并行智能体<br/>MRS / DM / FV / SC 等"] --> B["基于输入特征的专用 MLP 需求预测<br/>TF-IDF 向量化 + 4层轻量回归网络"]
    B --> C["以显存时空占用量化真实服务代价<br/>Prefill + Decode 累积 KV Token-Time"]
    C --> D["虚拟时间公平队列排序<br/>计算虚拟完工时间标号 f_j = V(a_j) + C_j"]
    D --> E["选择性饱和服务派发<br/>单智能体子任务连续饱和执行"]
    E --> F["底层 vLLM 非抢占执行与显存状态流转<br/>Running 队列饱和运行 / Swapped 队列保底"]
    F --> G["智能体全部子任务完成并输出结果"]

关键设计

1. 显存中心型服务成本建模:捕捉 KV 缓存时空膨胀的二次方代价 传统 LLM 调度系统(如 VTC)通常以计算 token 数量的加权和(例如 \(p + 2d\))作为服务开销度量,然而在大规模并发服务实践中,系统的真正性能瓶颈在于 GPU 显存中的 KV Cache 容量,而非纯计算单元。随着自回归解码的逐 token 推进,序列占用的显存空间随时间持续累积。若单个推理任务的预填充长度为 \(p\)、生成长度为 \(d\),其在整个生命周期内所霸占的显存时空累积占用(KV Token-Time)应定义为所有迭代步骤占用块数的积分: $\(c = \sum_{i=1}^{d} (p + i) = pd + \frac{d^2}{2}\)$ 其中 \(d^2/2\) 精确刻画了解码长度对显存容量与持有时间的双重膨胀效应。整个智能体 \(j\) 的总服务开销 \(C_j\) 则为其所有并行或依赖子任务的 KV Token-Time 之和 \(\sum c\)。这种建模避免了纯算力模型将显存敏感型长解码任务过低估计的缺陷。

2. 基于智能体先验的轻量级 MLP 需求预测:极低开销下的精准回归 实施调度决策必须在智能体到达时刻获知其预估总代价 \(C_j\)。现有预测方法(如 \(S^3\))微调包含 6600 万参数的 Distillbert 判别模型,不仅单次预测开销高达 55 ms,且单一模型难以拟合多样化智能体的异构生成分布。Justitia 发现,同一类业务智能体(如 Fact Verification 中的 query 生成、MapReduce 的分块总结)其提示词与输出长度分布呈现出极强的自身统计稳定性。因此,Justitia 为每类智能体独立训练一个紧凑的 4 层多层感知机(MLP)。输入端采用轻量级的 TF-IDF 算法将输入 Prompt 转化为词频统计向量,首层神经元数量与智能体平均输入规模正相关。该方案使得单次推理耗时降至 2.16 ms,且仅需每个智能体 100 条历史样本即可在 1 分钟内完成训练收敛,兼顾了调度实时性与预测泛化力。

3. 虚拟时间公平队列与常数级最坏延迟边界:消除饥饿的选择性饱和调度 为了在不产生队头阻塞的前提下实现顺序服务,Justitia 借鉴了网络报文调度中的广义处理器共享(GPS)与虚拟时间公平排队思想。系统维护虚拟时间变量 \(V(t)\),其随物理时间 \(t\) 的推进速率取决于当前活跃智能体总数 \(N_t\) 与系统总显存容量 \(M\): $\(V(0) = 0, \quad \frac{\mathrm{d}}{\mathrm{d}t} V(t) = \frac{M}{N_t}\)$ 当智能体 \(j\) 在物理时刻 \(a_j\) 到达且预估代价为 \(C_j\) 时,其在理想公平系统中的虚拟完工时间标号直接计算为: $\(\bar{f}_j = V(a_j) + C_j\)$ 该虚拟完成时间具有一次性确定、无需随后续新到达任务动态重算的关键性质。Justitia 严格按照 \(\bar{f}_j\) 递增顺序为各智能体子任务分配 GPU 资源,使得排在前面的智能体获得全部可用 KV Cache 迅速饱和执行。作者从理论上证明了(Theorem B.1),智能体 \(j\) 在 Justitia 下的实际完工时间 \(f_j\) 相比于理想公平系统 \(\bar{f}_j\) 的最大偏差严格满足: $\(f_j - \bar{f}_j \le 2 c_{\max} + \frac{C_{\max}}{M}\)$ 其中 \(c_{\max}\) 为任意单推理任务的最大开销,\(C_{\max}\) 为任意智能体的最大总开销。这一常数上界从数学上保证了小任务无法无休止延迟大任务,彻底杜绝了 SJF 类算法的饿死漏洞。

一个完整示例

假设集群总 KV 块容量 \(M = 100\)。时刻 \(t=0\),系统处于空闲状态,此时到达一个大象智能体 \(A\)(包含 10 个子推理,总代价 \(C_A = 200\),最大单子任务开销 \(c_{\max} = 20\)),此时活跃数 \(N_0 = 1\),虚拟时间速率 \(\mathrm{d}V/\mathrm{d}t = 100/1 = 100\)。根据公式计算,\(\bar{f}_A = V(0) + 200 = 200\)。系统将全部 100 个 KV 块分配给智能体 \(A\) 的前 5 个并发任务。

在时刻 \(t=1\) 时,大象智能体 \(A\) 已消耗 100 单位显存时代价,此时两个短小的老鼠智能体 \(B\) 和 \(C\) 同时到达,各自预估代价为 \(C_B = 30\) 与 \(C_C = 40\)。在 \(t=1\) 时系统虚拟时间推移为 \(V(1) = 100\)。于是智能体 \(B\) 的虚拟完成标号为 \(\bar{f}_B = 100 + 30 = 130\),智能体 \(C\) 的标号为 \(\bar{f}_C = 100 + 40 = 140\)。由于 \(\bar{f}_B < \bar{f}_C < \bar{f}_A\)(130 < 140 < 200),当智能体 \(A\) 正在运行的某个单次子任务非抢占完成释放显存后,系统并未按 FCFS 规则继续把后续显存给 \(A\),而是优先饱和执行 \(B\) 和 \(C\)。智能体 \(B\) 和 \(C\) 仅耗费极短时间便彻底执行完毕,随后全部资源再次全速服务 \(A\)。最终 \(B\) 和 \(C\) 享受到了近乎独占的完成速度,而 \(A\) 的总体完工时间由于延迟受限于理论常数 \(2 c_{\max} + C_{\max}/M\),避免了无尽等待。

损失函数 / 训练策略

针对每个业务智能体类型维护的 4 层 MLP 成本预测模型,其网络参数通过带有 \(L_2\) 正则化的均方误差(Mean Squared Error, MSE)进行梯度下降优化: $\(\mathcal{L}(\theta) = \frac{1}{K} \sum_{k=1}^{K} \left( \hat{C}_k(\theta) - C_k \right)^2 + \gamma \|\theta\|_2^2\)$ 其中 \(K=100\) 为每个智能体类型采集的历史离线执行样本数量,\(\hat{C}_k\) 为模型预测值,\(C_k\) 为基于真实执行 profile 计算得出的真值(采样生成阶段固定 temperature 为 0 确保确定性),\(\gamma\) 为正则化超参数。训练在单个 NVIDIA V100 GPU 上进行,每个智能体专属模型仅需约 1 分钟即可完成训练收敛。

实验关键数据

主实验

评估基于 vLLM 框架实现,在由 300 个包含不同输入的多样化真实智能体(涵盖 MapReduce 总结、文档合并、代码检查、知识问答验证、事实核验、数学自一致性等 9 类任务,划分为小 72%、中 26%、大 2%)构成的混合负载套件上展开。硬件后端覆盖 LLaMA-7B(A100-40GB)、LLaMA-13B(4×V100-16GB 张量并行)以及 Qwen2.5-32B(H800-80GB),负载到达强度沿用阿里 Mooncake 真实生产集群轨迹(设置 1×、2×、3× 提交密度)。

下表展示了在 LLaMA-7B 及多种基线调度器下各关键维度的定量对比:

调度方案 调度粒度与策略核心 平均 JCT (s) 相对 VTC 降低 (%) 92% 任务延迟上限 恶劣饥饿风险 调度开销
vLLM (FCFS) 请求级先来先服务 388.2 - 严重延迟 (>80%) 无 极低 (<1ms)
vLLM-SJF 请求级 Distillbert 预估最短优先 215.4 39.5% 较高延迟 (>50%) 存在 高 (~55ms)
Parrot 智能体级 FCFS 412.5 增加 6.7% 严重队头阻塞 无 极低 (<1ms)
VTC 智能体级计算代币瞬时公平份额 386.4 基线 (0.0%) 1.0× (定义基准) 无 低 (~5ms)
SRJF 智能体级预估剩余代价最短优先 148.6 61.5% 较高尾部延迟 极严重 (大任务无限拖延) 低 (~3ms)
Justitia (本文) 智能体级虚拟时间选择性饱和 164.2 57.5% 仅 8% 任务延迟 (最大延迟 26.0%) 理论证明无饥饿 (有界) 极低 (<10ms)

消融实验

为了验证显存时空代价建模(KV Token-Time)与专用轻量 MLP 预测器相对于传统方法的优越性,作者设计了两组对照实验: 1. 预测器消融:对比专用 MLP 方案与 \(S^3\) 所采用的微调 Distillbert 方案(基于 V100 训练,在 2× 负载密度下测试)。 2. 代价模型消融:将 Justitia 中的显存模型替换为 VTC 的计算代币模型 \(p + 2d\)(记为 Justitia/C)。

评估维度 / 模块配置 平均相对误差 (%) 单次预测推理开销 (ms) 全局平均 JCT (s) 离线训练时长 JCT 性能衰减
Distillbert (S³ 基线) 452.0% 55.70 ms 366.7 s ~2 小时 基准对比
Justitia 专用 MLP 53.0% 2.16 ms 151.1 s ~1 分钟 大幅领先
Justitia 完整版 (显存模型) - - 164.2 s - 最优性能
Justitia/C (纯计算代币模型) - - 233.7 s - 性能暴跌 42.3%

关键发现

  • 显存建模不可或缺:若放弃空间与时间积分的 KV Token-Time 建模并退回至算力代币,由于无法反映解码序列对显存的累进占有惩罚,平均 JCT 出现高达 42.3% 的剧烈恶化。
  • 饱和调度打破平衡困境:Justitia 在 3× 密度下有 92% 的智能体完工时间早于或等于 VTC,仅 8% 的任务受到轻微延迟,且平均延迟幅度低于 10%,最坏延迟仅为 26.0%,显著优于基线方案。
  • 抗饥饿微基准验证:在持续以 1 秒/个的频率涌入短小老鼠任务(KBQAV、CC 等)的极端对抗测试中,SRJF 下的大象任务(MapReduce 总结)完工时间呈发散式上升乃至永久饿死;而 Justitia 下大象任务的完工延迟始终严格收敛在常数理论边界以内。
  • 误差容忍鲁棒性极高:当人为施加高达 \(\lambda=3\)(即在 \([1/3, 3]\) 区间内随机缩放真实开销)的剧烈预测噪声时,Justitia 的平均 JCT 相比于理想预测仅微幅上浮 9.5%,表现出极高的调度鲁棒性。

亮点与洞察

  • 从机械瞬时公平到长效完工公平的观念跃升:系统设计敏锐地捕捉到了分布式系统理论中 GPS 与 LLM 并发调度的本质契合点,揭示了「短期均分显存反倒两败俱伤,顺序饱和执行才能全员加速」的调度哲理。
  • 二次方 KV Token-Time 精准锚定显存瓶颈:公式 \(pd + d^2/2\) 直观且扎实地刻画了自回归解码随步数展开对显存持有时间与空间占用的连带效应,为 LLM 系统资源量化提供了标准参照。
  • 任务级稳定先验打消重模型迷思:避免了在调度关键路径上部署重型 Transformer 预测模型的恶性循环,利用业务智能体的结构化模板稳定性,以微秒级 TF-IDF + 极浅 MLP 达成了高精度需求预估。

局限与展望

  • 假设单节点显存充裕且忽略极端碎片化:理论分析建立在单个推理 Prompt 占用远小于总显存 \(M\) 的假设上;在超长上下文(如 128k+ tokens)单任务并发时,显存碎片或无法被完全忽略。
  • 依赖智能体类型的离线先验:若用户提交完全未见过、无对应预训练 MLP 的全新结构自由智能体,系统的预测误差可能放大,需回退至启发式冷启动估算。
  • 未结合跨任务 Prefill 缓存复用:在多任务智能体中,同轮次子任务常常共享相同的系统 Prompt 或上下文,未来可探索将 Prompt Cache(前缀缓存感知)融入虚拟时间与饱和权重的联合度量。

相关工作与启发

  • vs. VTC (OSDI 2024):VTC 引入虚拟 token 计数器追求瞬时公平均分,但在多任务并发智能体场景下会导致所有任务被卡在半饱和状态并显著拖长完成时间;Justitia 引入虚拟时间 GPS 完成标号,改用选择性饱和策略,平均 JCT 提速 57.5%。
  • vs. Parrot (OSDI 2024) / Ayo (ASPLOS 2025):Parrot 和 Ayo 关注语义变量与智能体级 DAG 编排,但在底层仍沿用 FCFS 队列,面对异构大任务容易产生严重队头阻塞;Justitia 实现了兼顾短任务优先与抗饥饿保障的动态公平排队。
  • vs. FastServe / S³:FastServe 依赖高开销的抢占与多级反馈跳轮,S³ 采用重型语言模型预测解码长度造成巨大推理开销;Justitia 采用基于智能体先验的轻量 MLP(2.16 ms)并坚持非抢占式执行,开销显著更低。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ [首次将选择性饱和执行与网络 GPS 虚拟时间公平队列迁移至大模型智能体并发调度,理论与构思十分巧妙]
  • 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 300 种多类别智能体、多模型后端及 Mooncake 真实生产集群轨迹,消融、抗饥饿微基准与鲁棒性实验完备]
  • 写作质量: ⭐⭐⭐⭐⭐ [从动机阐释、显存公式推导到延迟边界的数学证明逻辑严密,图文对照极清晰]
  • 价值: ⭐⭐⭐⭐⭐ [对共享 GPU 集群的大模型多智能体服务框架(如 vLLM/SGLang 扩展)具备重大的工程落地指导意义]