From Tokens to Layers: Redefining Stall-Free Scheduling for LLM Serving with Layered Prefill

文章标题与作者
- 标题: 从令牌到层:用分层预填充重新定义大型语言模型服务的无停顿调度 (FROM TOKENS TO LAYERS: REDEFINING STALL-FREE SCHEDULING FOR LLM SERVING WITH LAYERED PREFILL)
- 作者: Gunjun Lee, Jiwon Kim, Jaiyoung Park, Younjoo Lee, Jung Ho Ahn

A1 主要贡献

本文旨在解决在生产环境中大型语言模型(LLM)推理服务所面临的核心挑战:如何在固定的计算、内存和互连预算下,同时满足严格的服务水平目标(SLO),即首次令牌时间(TTFT)和令牌间时间(TBT),并最大化吞吐量。

核心问题: 现代LLM服务系统采用分块预填充(chunked prefill)等无停顿调度技术,通过将长提示(prompt)沿令牌维度切分,并将预填充(prefill)与进行中的解码(decode)迭代交错执行,有效稳定了TBT。然而,这种方法在专家混合(MoE)模型中引入了巨大的开销:冗余的专家权重加载导致内存流量增加高达39%,并增加了能耗。

研究目标: 提出一种新的调度范式,以消除由分块预填充引起的MoE权重重复加载问题,从而在维持无停顿解码的同时,降低内存带宽需求、TTFT、端到端延迟以及每令牌能耗。

创新点与主要贡献:
本文提出了分层预填充(layered prefill),将Transformer的层组作为主要调度单元。其核心思想是:
1. 将层作为一等调度单元:将模型垂直划分为连续的层组,并在每个迭代中,只让一个层组同时执行预填充和解码,而其他组只执行解码。这种“每迭代一组”的规则保证了无停顿解码。
2. 消除分块引起的MoE权重重复加载:通过分层预填充,每个输入提示的预填充过程对每一层只遍历一次,从根本上避免了因分块而导致的专家权重重复加载,系统性地减少了冗余的参数流量,并改善了在长上下文、小批量场景下的TTFT-TBT权衡。
3. 降低服务能耗:通过削减冗余的MoE重载,减少了片外参数移动和GPU DRAM活动,从而降低了能耗。

实验评估表明,与分块预填充相比,分层预填充在TTFT-TBT的帕累托前沿上实现了持续改进,减少了专家加载流量和能源成本,同时保持了无停顿解码。总体而言,将调度轴从令牌转向层,为在共置环境中实现高效、节能的LLM服务开辟了一个新的操作范式,可将TTFT降低高达70%,端到端延迟降低41%,每令牌能耗降低高达22%。

A3 背景知识与关键观察

2.1 LLM 架构

当代LLM的通用架构。当前如GPT【GPT-4 Technical Report, 2024, https://arxiv.org/abs/2303.08774】、Llama 【The llama 3 herd of models, 2024, https://arxiv.org/abs/2407.21783】和Qwen 【Qwen3 Technical Report, 2025, https://arxiv.org/abs/2505.09388】等大型语言模型通常采用仅解码器(decoder-only)的Transformer架构 【Attention is all you need, 2017, Advances in neural information processing systems】。模型将一系列令牌ID映射为连续的向量表示,通过一个Transformer解码器块堆栈进行处理,最后将得到的隐藏状态投影到词汇表的logits上,以预测下一个令牌。

解码器块的内部结构。每个解码器块结构相同:一个自注意力层后跟一个前馈网络(FFN),两者周围都应用了残差连接和归一化。自注意力机制使每个令牌(表示为向量)能够整合序列中较早令牌的上下文信息。它通过学习到的投影权重($W_Q$, $W_K$, $W_V$)将输入隐藏状态投影到查询、键和值矩阵,应用因果掩码以确保令牌只关注过去和当前的位置,并通过查询和键之间的softmax归一化点积来导出注意力权重。这些权重引导对值的加权求和,从而产生一个具有上下文感知能力的令牌表示。已有大量工作致力于加速注意力算子,例如FlashAttention及其后续工作【Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022, Advances in neural information processing systems; FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2024, The Twelfth International Conference on Learning Representations; Flashattention-3: Fast and accurate attention with asynchrony and low-precision, 2024, Advances in Neural Information Processing Systems】,它们减少了注意力的I/O并提高了在现代AI加速器(如GPU【NVIDIA Hopper H100 GPU: Scaling Performance, 2023, IEEE Micro; Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings, 2023, Proceedings of the 50th annual international symposium on computer architecture】和TPU【The design process for Google’s training chips: TPUv2 and TPUv3, 2021, IEEE Micro; Tensor Processing Units (TPU): A Technical Analysis and Their Impact on Artificial Intelligence, 2023, Tech4Future Information Technology Report】)上的内核效率。

前馈网络(FFN)的功能。FFN独立地提炼每个令牌的隐藏状态。它由两个或多个带有非线性激活函数(如ReLU或SwiGLU【Glu variants improve transformer, 2020, arXiv preprint arXiv:2002.05202】)的全连接(FC)层组成。通常,第一个FC层将隐藏维度扩大2-4倍以增加表示能力,而最后一个层则将其投影回原始大小,以保持层间的维度一致性。

2.2 LLM 推理结构:预填充和解码

自回归推理过程。普通的LLM推理(除非使用扩散或类似方法)是自回归的;它逐个生成输出令牌,同时依赖于先前生成的令牌。该过程包括多次通过完整的解码器块堆栈,这些过程被称为迭代。这些迭代分为两个不同阶段:预填充(prefill)和解码(decode)。

预填充阶段的特点。预填充阶段对应于第一次迭代。它沿序列维度并行处理整个提示令牌序列,将所有令牌传递给每个解码器块。由于其并行性,该阶段能够实现AI加速器的高计算利用率【PrefillOnly: An Inference Engine for Prefill-only Workloads in Large Language Model Applications, 2025, arXiv preprint arXiv:2505.07203】;然而,由于必须在生成第一个输出令牌(TTFT)之前处理所有提示令牌,因此其延迟也很高。

解码阶段的特点。解码阶段包括所有后续迭代,每个迭代为每个请求生成一个新令牌。由于每次迭代的序列长度只有一个令牌,计算单元利用率不足【AttAcc! Unleashing the Power of PIM for Batched Transformer-based Generative Model Inference, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2】,效率在很大程度上依赖于内存访问模式。为了提高利用率,系统通常将来自多个请求的解码迭代进行批处理。

2.3 先前的LLM服务方法与挑战

LLM服务系统的演进。LLM服务系统已迅速发展,以在严格的延迟服务水平目标(SLO)下支持不断增长的请求量。一系列近期工作【Orca: A distributed serving system for {TransformerBased} generative models, 2022, 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22); Efficient Memory Management for Large Language Model Serving with PagedAttention, 2023, Proceedings of the 29th Symposium on Operating Systems Principles; Taming Throughput-Latency tradeoff in LLM inference with Sarathi-Serve, 2024, 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24); {DistServe}: Disaggregating prefill and decoding for goodput-optimized large language model serving, 2024, 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24); Splitwise: Efficient generative llm inference using phase splitting, 2024, 2024 ACM/IEEE 51st Annual International Symposium on Computer Architecture (ISCA); Beyond the Buzz: A Pragmatic Take on Inference Disaggregation, 2025, arXiv preprint arXiv:2506.05508; Pod-attention: Unlocking full prefilldecode overlap for faster llm inference, 2025, Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2; Prefill-Decode Aggregation or Disaggregation? Unifying Both for Goodput-Optimized LLM Serving, 2025, arXiv preprint arXiv:2508.01989】提出了日益复杂的调度策略,旨在提高利用率同时保持响应性。

早期系统的局限性。像FasterTransformer【FasterTransformer, 2019, https://github.com/NVIDIA/FasterTransformer】这样的早期系统依赖于静态或按请求批处理:每个批次是固定的并独立处理,不会动态合并到达的请求。虽然实现简单,但这种设计在解码阶段导致硬件利用率低下,因为同时处理的向量数量很少,并且解码批次可能因等待长预填充而停顿,导致并发请求的令牌间时间(TBT)延迟很高 。

连续批处理的引入。为了缓解这些低效问题,像Orca【Orca: A distributed serving system for {TransformerBased} generative models, 2022, 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)】这样的系统引入了连续批处理,它合并预填充和解码请求,并在迭代级别上一起调度它们,从而大幅提高了推理吞吐量。然而,当长预填充延迟了后续解码步骤时,这种方法仍然会遭受停顿。

分块预填充的提出。Sarathi-Serve【Taming Throughput-Latency tradeoff in LLM inference with Sarathi-Serve, 2024, 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24)】通过分块预填充解决了这个问题。它不是一次性处理整个长输入,而是将序列分割成更小的块,并与解码请求交错形成混合批次。这种重构实现了无停顿执行,因为解码请求不再被冗长的预填充阻塞。此外,通过设置块大小(通常为256或512个令牌),调度器既能满足TBT约束,又能更均匀地在迭代间分配预填充工作,从而提高了硬件利用率。

分块预填充的工作流程。如图1(左)所示,分块预填充在每次迭代中处理一个块,处理完当前块的所有解码器层后才移至下一个块。每个块重用前一个块的KV缓存,直到整个提示被处理完毕。为了提高效率,短请求可以合并成一个单独的块。块大小受到限制以满足尾延迟SLO,从而限制了每次迭代的批次大小。因此,分块预填充显著改善了服务水平目标(例如TTFT和TBT),我们在评估中将其作为基线。

2.4 LLM架构和使用方式的最新趋势

专家混合(MoE)架构。MoE架构【Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017, International Conference on Learning Representations; Gshard: Scaling giant models with conditional computation and automatic sharding, 2021, International Conference on Learning Representations; Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022, Journal of Machine Learning Research】已成为一种有前景的方法,可以在不按比例增加计算复杂性的情况下扩展模型大小。与稠密的FFN(即使在小解码批次下也激活所有参数)不同,MoE通过路由器将令牌路由到一小部分独立的FFN(称为专家)。这种稀疏性通过每个令牌只激活模型参数的一小部分来有效地扩展参数数量;然而,它也给调度和内存访问带来了挑战【St-moe: Designing stable and transferable sparse expert models, 2022, arXiv preprint arXiv:2202.08906; Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale, 2022, International conference on machine learning】。

长上下文处理。最近的LLM工作负载涉及处理比以往长得多的输入提示。以前的系统通常处理数百个令牌的提示【Language Models are Few-shot Learners, 2020, Proceedings of the 34th International Conference on Neural Information Processing Systems】,而由多轮对话、检索增强生成和思维链式推理驱动的现代应用需要数万个令牌的上下文。随着像Claude 3【The Claude 3 Model Family: Opus, Sonnet, Haiku, 2024, https://api.semanticscholar.org/CorpusID:268232499】 、Gemini 2.5【Gemini 2.5: Pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capabilities, 2025, arXiv preprint arXiv:2507.06261】、GPT-OSS【gpt-oss-120b & gpt-oss-20b Model Card, 2025, arXiv preprint arXiv:2508.10925】和DeepSeek-R1【Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning, 2025, arXiv preprint arXiv:2501.12948】等长上下文LLM的部署,这一趋势正在加速,其上下文窗口已超过10万个令牌。因此,预填充阶段在计算上变得更加昂贵和内存密集,加剧了当代共置服务系统中的瓶颈,因为这些系统中预填充和解码都由相同的AI加速器处理。

新趋势带来的挑战。LLM架构和工作负载的这些进步增加了LLM服务的计算和内存需求,凸显了为保持SLO合规性而需要新的调度和并行技术的必要性。

2.5 LLM服务中的能耗

能耗成为首要问题。功率和能源效率是大规模LLM服务的首要考虑因素。除了满足面向延迟的SLO外,服务提供商还必须最小化服务请求所需的能源,以节省运营成本和减少碳足迹【Dynamollm: Designing llm inference clusters for performance and energy efficiency, 2025, 2025 IEEE International Symposium on High Performance Computer Architecture (HPCA)】。

能耗的构成与核算。在AI加速器上,LLM服务的能源成本主要可视为四个部分的总和:(i)即使在空闲时也保持设备活动的静态基线功耗;(ii)由矩阵乘法/累加引擎主导的计算能耗(例如GPU中的Tensor Cores【GPU tensor cores for fast arithmetic reductions, 2020, IEEE Transactions on Parallel and Distributed Systems; Accelerating ml workloads using gpu tensor cores: The good, the bad, and the ugly, 2024, Proceedings of the 15th ACM/SPEC International Conference on Performance Engineering】和TPU中的MXU【Tensor Processing Units (TPU): A Technical Analysis and Their Impact on Artificial Intelligence, 2023, Tech4Future Information Technology Report; Benchmarking the performance and energy efficiency of AI accelerators for AI training, 2020, 2020 20th IEEE/ACM International Symposium on Cluster, Cloud and Internet Computing (CCGRID)】);(iii)通过内存层次结构(片上SRAM和片外DRAM,如HBM或GDDR)移动参数和KV缓存的内存能耗;以及(iv)连接设备或主机的互连(如NVLink, PCIe, InfiniBand, Ethernet)的通信能耗。实际上,数据移动占了能源的很大一部分,因为每个生成的令牌都会重复重新读取模型参数的大部分。通过比较4位量化模型(W4A16)与更高精度模型(Shi & Ding, 2025【Systematic Characterization of LLM Quantization: A Performance, Energy, and Quality Perspective, 2025, https://arxiv.org/abs/2508.16712】)的功耗可以验证这一观察。因此,一种实用的核算方法是跟踪在内存层次的每一级移动的总字节数,并乘以该级别的经验测量的每字节能耗,其中片外DRAM项通常决定了总体规模 。

MoE对能耗的影响。通用矩阵乘法(GEMM)是计算密集型还是内存密集型,可以通过比较批次隐含的算术强度与加速器的“脊点”(ridge point)来判断。脊点定义为峰值算术吞吐量除以峰值内存带宽,它标志着执行从内存密集型转为计算密集型的每字节操作数(Op/B)。在脊点附近操作意味着系统能高效利用内存和计算。现代AI加速器的脊点通常在100到300 Op/B的量级【The New LLM Bottleneck: A Systems Perspective on Latent Attention and Mixture-of-Experts, 2025, https://arxiv.org/abs/2507.15465】。这意味着对于2字节的数据类型,批次大小大约需要200到600。考虑到这一点,虽然MoE每个令牌只激活一部分专家,但这种稀疏性是否真的能在服务系统中减少延迟和能耗,需要仔细分析。以一个拥有128个专家和top-k为8的MoE模型(如Qwen3-30B-A3B)为例,对于一个2048个令牌的输入提示,平均每个专家处理约128个令牌。这低于当代加速器的脊点,因此每个专家的FC计算仍然是内存密集型的,端到端的延迟和能耗主要由从HBM加载MoE权重决定。要使MoE计算变为计算密集型,需要超过8192个令牌,但这会导致延迟(TBT)急剧增加 。

3. 分块预填充在含MoE的LLM上的局限性

3.1 MoE权重加载模型

分块预填充与MoE的冲突。分块预填充在MoE层上引入了一个冲突。为了满足尾延迟SLO,分块预填充调度器限制了每次迭代的批次大小(§5.2),而高效的MoE权重加载则需要更大的批次来最大化复用已加载的专家。这种不匹配导致专家权重被广泛激活但复用率低,使得执行瓶颈在于内存流量而非计算。

稀疏性侵蚀问题。假设批次大小为128,top-k为8,专家数量为128;平均每个迭代只有8个样本被路由到每个专家。8这个数字远低于现代AI加速器的脊点。因此,我们将此问题表述为“稀疏性侵蚀”问题;分块操作增加了专家覆盖率,同时抑制了复用,侵蚀了MoE之所以高效的稀疏性优势。接下来,我们将在现实的服务条件下量化专家激活情况,并分析分块如何破坏MoE的稀疏性优势。

专家覆盖率的实证分析。在Qwen3-30B-A3B(简称Qwen)【Qwen3 Technical Report, 2025, https://arxiv.org/abs/2505.09388】上,使用ShareGPT数据集 【ShareGPT, 2023, https://sharegpt.com/】的测量结果展示了这个问题(见表1)。模型和数据集的详细信息在表3和表4中提供。在小批次(小于16)下,加载的专家不到一半,即使批次大小为64,平均专家覆盖率仍低于70%。如图3所示,在稳定的服务状态(例如SLO达成率超过90%)下,arXiv数据集 【A discourse-aware attention model for abstractive summarization of long documents, 2018, arXiv preprint arXiv:1804.05685】上的平均批次大小低于32,ShareGPT数据集上低于128,这意味着每个解码批次只激活了有限比例的专家。在纯解码中,这种有限激活是有益的:它相对于密集模型减少了内存流量。

分块预填充对稀疏性的负面影响。相比之下,分块预填充创建的混合批次通常要大得多(批次大小常超过256),这会激活大多数专家,从而侵蚀了稀疏性带来的好处。同时,TBT SLO的约束限制了块大小;因此,路由到每个专家的令牌数量仍不足以饱和计算资源,使得MoE层仍然是内存密集型的。因此,内存流量主导了系统行为,放大了每次迭代的延迟和能源成本,其程度与提示长度和块数量成正比。

3.2 微基准测试:运行时与块大小的关系

实验设置。为了理解块大小如何影响MoE效率,我们隔离了它对内核性能的影响。图2报告了在Qwen上,MoE权重加载和内核运行时随预填充块大小变化的函数关系(评估设置见§5.1),输入长度固定为8192个令牌,以近似arXiv数据集中的平均提示长度。

小块大小的性能瓶颈。当块大小为512时,MoE运行时占预填充运行时的50%以上,且预填充运行时超过500毫秒。随着块大小增加,权重加载开销急剧下降,大致与块大小成反比,相应的每操作运行时也减少。较大的块将更多令牌路由到每个专家,提高了复用并减少了冗余传输。

大块大小的性能改善。当块大小达到4096-8192个令牌时,MoE加载量降至100GB以下,预填充运行时稳定在200毫秒附近,MoE不再主导执行。其他内核对块大小的依赖性较小,主要原因是小问题规模下的算术强度较低,但其影响远不如MoE内核显著。这些结果表明,块大小强烈影响MoE效率:小块引发稀疏性侵蚀和内存密集型减速,而大块通过实现高效的专家复用缓解了这个问题。


图2. (左) MoE权重加载量 vs. 块大小。阴影区域表示单个块加载的MoE权重。(右) 各内核运行时 vs. 块大小。我们将输入长度固定为8,192个令牌。

3.3 LLM服务中块大小的权衡

块大小影响的三个方面。在服务场景中,块大小塑造了系统行为的三个关键方面:专家复用、服务容量和迭代延迟。较大的块通过增加复用和减少冗余权重加载来提高MoE效率,并且它们也能支持更高的请求率。然而,与此同时,较大的块会延长迭代持续时间,从而增加尾延迟并有违反SLO的风险。

权衡的具体表现。表2总结了在arXiv工作负载上为Qwen服务时的这种权衡。每个块大小的请求率都经过调整,以使TTFT大致保持在2.5秒。随着块大小增加,预填充运行时减少。因此,在不同块大小间保持TTFT恒定意味着较大的块可以支持更高的请求率,从而使排队延迟和预填充运行时的总和保持在2.5秒。从效率角度看,较大的块减少了每令牌的能耗——定义为总能耗除以提示和生成令牌的总和【Beyond Test-Time Compute Strategies: Advocating Energy-per-Token in LLM Inference, 2025, Proceedings of the 5th Workshop on Machine Learning and Systems; StoreLLM: Energy Efficient Large Language Model Inference with Permanently Pre-stored Attention Matrices, 2025, Proceedings of the 16th ACM International Conference on Future and Sustainable Energy Systems】——从512块大小的60 mJ/tok降至2048块大小的32 mJ/tok(-46%)。这种改进是因为每个请求的块数减少,降低了专家权重的重载并增加了复用。吞吐量也随块大小增加而提高;系统在512块大小时支持1.3 req/s,但在2048时增加到2.6 req/s,这与图2中的运行时趋势一致。

与TBT约束的冲突。然而,这与Sarathi-Serve的原始配置相冲突,因为大块大小会导致违反TBT约束。随着块大小增加,p99 TBT从512令牌时的48毫秒急剧上升到2048令牌时的129毫秒,超出了SLO。这是迭代时间变长的直接后果,它延迟了并发解码请求的调度,并增加了分布的尾部。因此,从效率和吞吐量角度看最优的块大小,在严格的延迟要求下是不可行的。


表2. 在arXiv工作负载上,Qwen的块大小权衡。较大的块能改善运行时间和能源效率,但会增加尾部TBT。

A2 方法细节

4.1 概述

提出分层预填充。为了解决上述限制,我们提出了分层预填充(layered prefill)。其核心思想是将调度轴从令牌(tokens)转向层(layers)。我们不是沿着令牌维度将输入序列划分为块,而是将模型垂直划分为连续的层组(见图1)。这种基于组的层划分与仅解码器(decoder-only)的Transformer模型天然契合,因为在这些模型中,层是同质的且按顺序执行。


图1. (右上) 在每次迭代中,分块预填充将输入提示分割为多个块,并在每次迭代中按顺序处理一个块及解码任务。(右下) 对于分层预填充,只有一个层组同时执行预填充和解码,而其他组仅执行解码。预填充任务在每次迭代中前进一个组,从而保持解码的无停顿状态。

分层预填充的调度规则。在执行过程中,分层预填充遵循“每迭代一组”(one-group-per-iteration)的调度规则:在每次迭代中,只有一个指定的层组同时为新接收的请求执行解码和预填充,而所有其他组仅执行解码。因此,预填充工作被分散到多个迭代中进行,而解码则连续进行,不会出现停顿。

分层预填充的优势。这种设计保证了每个输入提示的预填充路径对每一层只遍历一次,这与分块预填充不同,后者迫使每一层为每个块重新处理一次提示。因此,跨块的冗余专家权重重载被消除,从而在保持严格TBT SLO下的无停顿解码的同时,减少了内存带宽消耗和能耗。

4.2 调度机制

逐组处理预填充。在每次迭代中,最多只有一个指定的层组会共同处理进行中的解码批次和新进入请求的预填充;其余的组只执行解码计算。因此,预填充任务逐个组地在迭代中推进,而解码过程保持不间断。

调度过程图解。图1用四个层组的例子说明了这个过程。在第一次迭代中,新请求的预填充在第1组内运行,而第2-4组只执行解码。在下一次迭代中,预填充推进到第2组,与解码共同调度,依此类推。这个过程重复进行,直到请求遍历了所有组。经过G(组的数量)次迭代后,该请求的预填充完成,而在此期间解码始终保持无停顿状态。

与分块预填充的对比。与分块预填充(它会在不同块之间引发冗余的MoE权重重载)不同,分层预填充调度确保每一层只处理一次输入提示。因此,解码在保持无停顿的同时,避免了令牌级划分所带来的内存开销放大问题。

保证无停顿解码。解码永远不会被阻塞,因为每次迭代都包含了所有层组的解码工作,所以TBT保持足够低,以满足SLO约束。简而言之,从分块预填充转向分层预填充,在保持无停顿解码的同时,显著减少了冗余的MoE权重移动和能耗。

4.3 与分块预填充的泛化结合

两种方法的正交性。分层预填充将模型垂直划分为几个连续的层组,而分块预填充则沿着令牌轴将输入提示划分为多个块。因此,这两种方法是正交的,意味着分块预填充和分层预填充可以同时应用。

结合应用的优势一:处理超长序列。将两种策略结合使用增加了调度的灵活性,并带来两个主要好处。首先,它能够高效处理非常长的输入序列。Mitra等人提出的分块流水线并行(Chunked Pipeline Parallelism)【Beyond the Buzz: A Pragmatic Take on Inference Disaggregation, 2025, arXiv preprint arXiv:2506.05508】通过将输入划分为相对较小的块并对每个块的预填充进行流水线处理,解决了处理极长输入长度的挑战。这种方法已被证明是在不需要复杂并行策略的情况下满足长输入TTFT约束的最优方法。由于分层预填充也与分块兼容,它继承了分块流水线并行在处理长输入方面的优势。

结合应用的优势二:增大块大小。其次,它允许块大小被大幅增加。增加块大小直接减少了给定输入提示的块数。因为MoE权重读取开销随块数增加而增长,所以更少的块能显著降低内存带宽消耗和能源成本。此外,当块大小变得足够大时,MoE层中每个专家的批次大小也变得足够大以饱和计算资源,使操作从内存密集型转向计算密集型。例如,在一个有128个专家和top-k=8的模型中,8192的块大小已经可以为每个专家产生512个令牌的有效批次大小,这足以使MoE层变为计算密集型。一旦块大小大到足以将MoE层转变为计算密集型状态,一个重要的优势就出现了:将输入提示划分为多个块进行MoE计算不再会导致延迟激增。

泛化结合的总结。通过这种泛化,系统可以选择性地结合分块预填充和分层预填充,利用两种方法的优点,同时减轻它们各自的缺点。

4.4 实现细节

动态调整层组数量。我们根据输入长度调整层组的数量。具体来说,我们选择组数G,使得每次迭代的预填充工作量大约与基线分块预填充中一个512令牌的块相匹配:

$$G(L) = \max(1, \lceil L/512 \rceil).$$

因此,对于一个长度为8192的长提示,我们设置G = 16,而对于一个长度为512的短提示,我们设置G = 1。这种选择将处理粒度与块大小为512的分块预填充对齐,使得不同调度器间的迭代时间和准入节奏具有可比性,从而隔离了切换调度轴(层与令牌)所带来的影响。这里的512是一个任意值,也可以是256或1024。我们在实验中使用512,以确保与配置为512块大小的分块预填充进行公平比较。当多个小输入同时到达时,我们将它们合并成一个批次以提高效率。

技术栈。我们在vLLM【Efficient Memory Management for Large Language Model Serving with PagedAttention, 2023, Proceedings of the 29th Symposium on Operating Systems Principles】的基础上实现了分层预填充。注意力机制使用FlashAttention3【Flashattention-3: Fast and accurate attention with asynchrony and low-precision, 2024, Advances in Neural Information Processing Systems】实现;大多数计算使用自定义CUDA内核,其余层使用torch.compile。我们采用CUDA Graphs来加速预填充和解码阶段。

A4 实验环境

硬件配置:
- GPU: 2块NVIDIA H100 GPU(每块80GB),通过NVLink连接。
- 并行策略: 采用张量并行(tensor parallelism)实现模型并行。

模型:
- 评估了两种具有相似架构但路由配置不同的代表性MoE模型。
- Qwen: Qwen3-30B-A3B的简称。包含128个专家,每个token激活top-8个专家。
- GPT: GPT-OSS-20B的简称。包含32个专家,每个token激活top-4个专家。


表3. 所评估MoE模型的规格,按指标重新组织。GQA代表分组查询注意力。

数据集:
- 使用了两个具有不同特性的数据集进行评估,其统计数据见下表4。
- ShareGPT: 包含用户与ChatGPT之间的多轮对话。输入长度分布广泛,标准差大,平均输入长度约是输出长度的六倍。
- arXiv Summarization (简称arXiv): 一个摘要基准数据集,输入远长于输出,平均输入长度约是输出长度的四十倍。


表4. 评估数据集(ShareGPT和arXiv Summarization)的输入和输出长度统计。

软件与工作负载:


表5. ShareGPT和arXiv数据集的SLO约束。

A4 实验结果

5.2 端到端SLO达成率

实验内容: 比较分层预填充和分块预填充在不同请求速率下,针对两种模型(Qwen和GPT)和两种工作负载(arXiv和ShareGPT)的SLO达成率。

实验结果 (图3): 分层预填充在所有模型和工作负载上均能提供更高的SLO达成率。

分析与结论:


图3. 不同请求速率下的SLO达成率。红色水平线标记了有效的SLO达成率阈值(90%)。


表6. Qwen在arXiv上的表现(请求速率为1.3 req/s)

5.3 SLO达成率分解

实验内容: 将端到端的SLO达成率分解为其两个组成部分——TTFT和TBT,并比较分层预填充与分块预填充在请求速率增加时的表现。

实验结果 (图4):
- TTFT: 在所有模型和数据集上,分层预填充在比分块预填充更宽的操作区域内保持了近乎完美的TTFT达成率。当系统接近饱和时,分层预填充的TTFT达成率下降更为平缓,意味着其排队延迟和预填充运行时相对更短。
- TBT: 分层预填充和分块预填充的TBT曲线在各种请求速率下都保持在近100%,表明两种调度器都有效地提供了无停顿调度。即使请求速率增加到饱和状态,TBT达成率也几乎没有下降。唯一的例外是在ShareGPT上评估GPT时,高请求率增大了解码批次大小,导致两种调度器都出现了TBT违规,但分层预填充的TTFT在更高请求率下(直到6.6 req/s)仍然保持稳健。

分析与结论: 分层预填充能维持更高的整体SLO达成率,主要得益于其在负载下保持TTFT的能力,而两种调度器都能很好地维持TBT。


图4. 按组件分解的SLO达成率随请求速率的变化。

5.4 MoE专家加载流量

实验内容: 通过计算处理100个请求的固定轨迹期间,专家权重加载的总字节数来量化专家激活引起的内存压力。

实验结果 (表7):
- ShareGPT: 分层预填充将专家加载量减少了12%。
- arXiv: 在长提示的arXiv工作负载上,效果更为显著,专家加载量减少了39%。

分析与结论: 分层预填充通过避免分块引起的重复加载,有效降低了专家权重加载流量,尤其是在长提示工作负载上。这些流量的减少与图3和图4中的SLO结果一致:更低的专家加载流量与更高请求率下的更高SLO达成率相关,特别是在预填充与解码比率高的工作负载上。


表7. 在Qwen上处理100个请求的总专家权重加载量。

5.5 随时间变化的令牌生成

实验内容: 在请求速率为1.3 req/s的arXiv工作负载下,使用Qwen模型比较单个请求在分层预填充和分块预填充下的累计令牌输出。

实验结果 (图5):
- 分层预填充比分块预填充更早地开始输出令牌,并保持了更高的生成速率。
- 图中陡峭上升的中间部分反映了分层预填充快速完成其他请求的预填充后,进入仅解码模式,此时令牌生成加速。
- 最终,平均端到端延迟从9.4秒减少到5.5秒(降低了41%)。

分析与结论: 分层预填充通过降低有效的TTFT和提高吞吐量,从而在保持无停顿解码的同时,降低了端到端延迟。


图5. 在arXiv上使用Qwen的令牌随时间生成情况。

5.6 每输出令牌的能耗

实验内容: 在arXiv工作负载的稳态操作点上,报告每输出令牌的能耗(mJ/tok)以及延迟指标(平均/p99 TTFT和TBT)。为每个调度器选择能够满足延迟SLO目标的最高请求速率,因此请求速率的差异反映了在相同QoS范围内的可用容量。

实验结果 (表8):
- Qwen: 相较于分块预填充(1.3 req/s),分层预填充能支持更高的请求速率(1.6 req/s,+23%),同时将每令牌能耗从56.6 mJ/tok降低到44.2 mJ/tok(-22%)。在相同请求速率下,能耗降低9%,平均TTFT下降超过50%。
- GPT: 相较于分块预填充(2.1 req/s),分层预填充能支持更高的请求速率(2.7 req/s,+29%),同时将每令牌能耗从37.4 mJ/tok降低到29.8 mJ/tok(-20%)。在相同请求速率下,能耗降低8%,平均TTFT下降超过70%。

分析与结论: 两种模型上,分层预填充都能在更高的可持续请求速率下,实现20-22%的每令牌能耗降低,同时提供相等或更好的TTFT和可比的TBT。这与微基准测试的结果一致:避免重复的专家权重重载减少了MoE层中的HBM/NVLink流量,从而在不牺牲延迟的情况下转化为更低的每令牌能耗。


表8. arXiv上的归一化能耗。我们报告了每输出令牌的能耗(J/tok;越低越好)以及在稳态、符合SLO的操作点下的平均/p99 TTFT和TBT。对于每个调度器,所列的请求速率是满足延迟SLO的最高速率,因此req/s反映了可用容量。

A7 补充细节

6.1 LLM服务系统

相关工作概述。大型语言模型的迅速普及推动了对高效服务基础设施的广泛研究。早期的系统如FasterTransformer【FasterTransformer, 2019, https://github.com/NVIDIA/FasterTransformer】和DeepSpeed-Inference 【Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale, 2022, International conference on machine learning】引入了内核级优化(如融合注意力和张量并行)以加速单模型推理。近期的工作则专注于在资源竞争下提高多用户服务的效率。例如,vLLM【Efficient Memory Management for Large Language Model Serving with PagedAttention, 2023, Proceedings of the 29th Symposium on Operating Systems Principles】提出了PagedAttention,通过细粒度的内存虚拟化来管理KV缓存,从而实现大批量解码的高吞吐量。SarathiServe【Taming Throughput-Latency tradeoff in LLM inference with Sarathi-Serve, 2024, 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24)】引入了自适应批处理和请求调度来平衡延迟和吞吐量。本文的工作建立在这一系列工作之上,通过重新审视预填充工作负载的调度,特别是解决在MoE模型中使用分块预填充时出现的低效问题。

6.2 LLM的能源效率

能源效率成为关键问题。除了原始性能,能源效率已成为大规模部署中的一个关键问题。先前的研究【Memory is all you need: An overview of computein-memory architectures for accelerating large language model inference, 2024, arXiv preprint arXiv:2406.08413; FlexNPU: A dataflow-aware flexible deep learning accelerator for energy-efficient edge devices, 2025, Frontiers in High Performance Computing】报告称,内存流量主导了现代加速器的能源足迹,这推动了减少参数移动或KV缓存访问的技术。像MLA【Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model, 2024, arXiv preprint arXiv:2405.04434】这样的架构提议通过压缩KV缓存来降低内存容量和功耗需求。服务层面的方法,包括早退策略和推测解码,也因其削减每令牌能耗的潜力而受到分析。然而,大多数先前的工作评估的是密集Transformer模型;相对较少关注MoE模型独特的能源行为,其中专家权重的加载会显著放大片外内存的能耗。我们的工作通过量化不同预填充调度策略下MoE的能源开销,并证明分层预填充可以在长上下文工作负载中同时降低延迟和能源成本,从而补充了这一研究领域。

A5 结论

传统的分块预填充调度对于MoE模型是低效的,因为它放大了冗余的专家权重加载,从而增加了延迟和能耗。本文引入了分层预填充,一种新的调度策略,它沿着层维度而不是令牌维度对模型进行划分。这种方法从根本上消除了由分块引起的冗余内存流量。

我们在代表性的MoE模型上进行的评估表明,分层预填充持续改善了SLO达成率,在长上下文工作负载上将专家加载流量减少了高达39%TTFT降低了高达70%端到端延迟降低了41%,并且每令牌能耗降低了高达22%。这些结果强调了将调度策略与新兴模型架构协同设计的重要性。

未来工作可以扩展此方法到复杂的多GPU环境,并探索自适应的层分组策略,特别是对于层数不能被组数整除的模型。进一步的研究还可以联合优化数据中心规模的吞吐量、延迟和能源效率。