作者：昇腾实战派

摘要

本文详细介绍了在vLLM推理引擎中为Qwen3稠密模型实现双流并行特性的技术方案。针对大语言模型推理中Attention和MLP层的All_Reduce通信操作导致的进程阻塞和计算空泡问题，双流并行通过将单个batch拆分为多个micro_batch，并引入独立的通信流，实现计算与通信的并行执行，从而掩盖通信延迟，提升硬件利用率和推理吞吐量。

核心内容：

  1. 背景与意义：分析vLLM推理引擎中batch_size与序列长度合轴带来的计算-通信矛盾，阐述双流并行的优化价值。
  2. 环境要求：vLLM >= v0.9.0，vLLM-Ascend >= v0.9.0rc1。
  3. 特性适配流程：
    • 数据拆分：详细说明需要拆分的4类数据（positions、hidden_states、residual_post、attn_metadata）及其拆分方法，包括简单拆分和考虑负载均衡的进阶拆分策略。
    • 双流控制：通过4个事件（ATTN_COM_FINISH、ATTN_AR_FINISH、FFN_COM_FINISH、FFN_AR_FINISH）实现计算与通信的异步执行，提供完整的伪代码实现。
  4. 性能效果：在input_length=2600、batch_size=400、TP=2的测试配置下，双流并行使Mean TTFT降低6.30%，Media TTFT降低6.71%，P99 TTFT降低6.90%，显著优化了推理延迟。

该技术方案通过创新的数据拆分策略和异步流控制机制，有效解决了大规模语言模型推理中的通信瓶颈问题，为高性能推理提供了可靠的技术保障。

1. 背景与意义

  • **背景：**在大语言模型推理中，如Qwen3稠密模型，Attention层和MLP层中的All_Reduce操作是关键的通信步骤，会导致进程阻塞，产生计算空泡，降低整体效率。vLLM推理引擎通过将batch_size和序列长度合轴为bs轴，进一步放大了计算耗时与通信等待的矛盾。
    在这里插入图片描述

    batch_size、seq_len合轴示意图

  • **意义：**双流并行技术旨在解决这一问题，通过将单个batch拆分为多个micro_batch，并引入独立的通信流，实现计算操作与通信操作的并行执行，从而掩盖通信延迟，提升硬件利用率和推理吞吐量。
    在这里插入图片描述

    双流并行前后耗时示意图

2. 环境版本

vllm >= v0.9.0
vllm-ascend >= v0.9.0rc1

3. 特性适配流程

3.1 数据拆分

  • 需要拆分哪些数据？
    将原本的1个batch拆分成2个micro_batch。需要拆分的数据有以下4个：
  1. positions：模型输入，表示当前token的位置信息，shape=[1, bs]；
  2. hidden_states：模型输入，表示当前token的特征向量，shape=[bs, hidden_dim]；
  3. residual_post：模型输入，用于残差计算，需要注意，在第0层Encoder的时候值为None，需要特殊处理；
  4. attn_metadata：全局信息，AscendMetadata类的对象，在Attention内部会自动通过get_forward_context().attn_metadata获取。其内部需有多个数据信息，其中在Qwen3稠密模型中被使用的有：
    • num_actual_tokens：当前步骤，文本输入的token数（即bs），int类型；
    • slot_mapping：用于将各个token位置映射到各自的slot中进行attention计算，是一个shape=[bs]的Tensor；
    • attn_state：表示当前attention的状态，Enum类型；有这几种状态：
      • AscendAttentionState.DecodeOnly：表示当前模型仅进行解码，不进行prefill
      • AscendAttentionState.ChunkedPrefill：表示当前模型正在进行prefill，并且是分块的填充
      • AscendAttentionState.ChunkedDecode：表示模型正在进行分块的解码
    • attn_mask：注意力掩码，是一个shape=[bs, seq_len]的Tensor；
    • seq_lens：表示每个样本的真实序列长度（Prefill+Decode），是一个shape=[batch_size]的Tensor
    • query_lens：表示每个样本当前的序列长度（Prefill或Decode），是一个shape=[batch_size]的Tensor
    • block_table：在模型中，token 通常被分为多个块来进行并行计算。该张量包含了每个 token 对应的块（block）的信息，是一个shape=[batch_size, num_blocks]的Tensor
  • 如何正确拆分数据？

  1. 选择"拆分点"

    • 简单方法。直接通过 batch_size // 2计算拆分点，将数据一分为二：
      由于vllm会将batch_size和seq_len合轴，所以无法直接拿到batch_size的大小。但可以通过attn_metadata中的query_lens数据来获取batch_size数和每个输入的query_len。

      在这里插入图片描述

      按 batch_size寻找拆分点，拆分数据示意图。
      batch_size_split_point = batch_size // 2
      token_split_point = sum(query_lens[:batch_size_split_point])

      • 缺点：decode+prefill混合情况下，会出现负载不均衡的情况。
        在这里插入图片描述

        在这里插入图片描述

  • 进阶方法。考虑负载均衡，寻找最佳切分点：
    在这里插入图片描述

    ① best_split_point = total_tokens // 2；
    ② 遍历query_lens，记录当前位置的累计token数cumu_tokens；
    ③ 通过 cumu_tokens 和 best_split_point 的差值，来决定最佳切分点。
    ④ 设置“平衡阈值“ balance_threshold。如果两个micro_batch的长度差异abs(seq_len1-seq_len2) / min(seq_len1, seq_len2)超过balance_threshold，认为负载不均衡，则不走双流并行。
    ⑤ 设置“长度阈值”length_threshold，如果切分后的两个micro_batch的长度有一个小于length_threshold，则不走双流并行。

  1. 拆分数据

    • hidden_states、positions、residual：这3个输入数据的shape=[bs, hidden_dim]，可以直接使用token_split_point进行拆分。

    hidden_states_pre = hidden_states[:token_split_point]
    hidden_states_post = hidden_states[token_split_point:]

    • attn_metadata：

      • num_actual_token：文本输入的token数，使用token_split_point拆分。

        num_actual_token_pre = token_split_point
        num_actual_token_post = attn_metadata.num_actual_token - token_split_point

      • slot_mapping：用于token到slot的映射，是一个shape=[bs]的Tensor，使用split_token_index拆分。

        slot_mapping_pre = attn_metadata.slot_mapping[:token_split_point]
        slot_mapping_post = attn_metadata.slot_mapping[token_split_point:]

      • seq_lens：每个样本的真实序列长度，是一个shape=[batch_size]的Tensor，使用batch_size_split_point拆分。

        seq_lens_pre = attn_metadata.seq_lens[:batch_size_split_point]
        seq_lens_post = attn_metadata.seq_lens[batch_size_split_point:]

      • query_lens：每个样本当前的序列长度，是一个shape=[batch_size]的Tensor，使用batch_size_split_point拆分。

        query_lens_pre = attn_metadata.query_lens[:batch_size_split_point]
        query_lens_post = attn_metadata.query_lens[batch_size_split_point:]

      • attn_state：attn_metadata.attn_state基本都是AscendAttentionState.ChunkedPrefill，保持不变。

      • attn_mask：是一个shape=[bs, seq_len]的Tensor，使用token_split_point进行拆分，同时更新各自的seq_len轴。

        attn_mask_pre = attn_metadata.attn_mask[:token_split_point, :max(seq_lens_pre)]
        attn_mask_post = attn_metadata.attn_mask[token_split_point:, :max(seq_lens_post)]

      • block_table：是一个shape=[batch_size, num_blocks]的Tensor，使用batch_size_split_point拆分。

        block_tables_pre = attn_metadata.block_tables[:batch_size_split_point]
        block_tables_post = attn_metadata.block_tables[batch_size_split_point:]

3.2 双流控制

在这里插入图片描述

4个事件：

  • ATTN_COM_FINISH：用来记录Attention的计算操作是否完成；
  • ATTN_AR_FINISH：用来记录Attention的ALL_Reduce操作是否完成；
  • FFN_COM_FINISH：用来记录MLP的计算操作是否完成；
  • FFN_AR_FINISH：用来记录MLP的ALL_Reduce操作是否完成。

伪代码：

"""DecoderLayer"""
def _forward_ms_layer(layer_index):
	"""ATTN"""
	for i in range(num_micro_batches):
		FFN_AR_FINISH[layer_index-1][i].wait()
		attention_comp_func(...)	# 执行 attention的计算
		ATTN_COM_FINISH[layer_index][i].record()
		with torch.npu.stream(current_ms_metadata.comm_stream):	# 切到通信流
			ATTN_COM_FINISH[layer_index][i].wait()
			attention_all_reduce_func(...)	# 执行attention的all_reduce通信
			ATTN_AR_FINISH[layer_index][i].record()

	"""MLP"""
	for i in range(num_micro_batches):
		ATTN_AR_FINISH[layer_index][i].wait()
		mlp_comp_func(...)	# 执行 mlp 的计算
		FFN_COM_FINISH[layer_index][i].record()
		with torch.npu.stream(current_ms_metadata.comm_stream):	# 切到通信流
			FFN_COM_FINISH[layer_index][i].wait()
			attention_all_reduce_func(...)	# 执行attention的all_reduce通信
			FFN_AR_FINISH[layer_index][i].record()

"""Qwen3Model"""
for layer_index in range(start_layer, end_layer):
	self.layers[layer_index]._forward_ms_layer(layer_index)

4. 特性适配结果

测试配置： input_length=2600，batch_size=400，TP=2。

Mean TTFT(ms)Media TTFT(ms)P99 TTFT(ms)
单流4318.883844.998084.73
双流4046.863587.157526.99
耗时下降6.30%6.71%6.90%

结果表明，双流并行特性有效优化了推理流程，显著降低了响应延迟，提升了模型的响应速度。

5. 总结

本文档系统地阐述了在vLLM推理引擎中，为Qwen3稠密模型实现双流并行 特性的完整技术方案与优化效果。其核心价值在于通过创新的计算-通信掩盖技术，显著提升了长序列文本下的推理特性。
1. 核心目标与价值：
针对Qwen3模型因Attention和MLP层中的All_Reduce集体通信操作而产生的“计算空泡”问题，双流并行技术旨在将通信耗时与计算耗时进行最大程度的并行化处理，从而减少进程空闲等待，充分挖掘硬件算力潜力。
2. 关键技术方案：

  • ​数据拆分​：创新性地提出了将单个Batch拆分为多个micro_batch的方案，并详细定义了需拆分的所有关键数据（如positions, hidden_states, attn_metadata等）。
  • ​负载均衡拆分点选择​：不仅提供了简单的对半拆分方法，更提出了基于总Token数的“进阶方法”，通过设置平衡阈值 与长度阈值 来智能判断并确保两个微批次的负载均衡，这是保证性能收益的关键。
  • ​异步流控制​：设计了精巧的多事件（如ATTN_COM_FINISH, ATTN_AR_FINISH）机制，利用独立通信流与主流进行同步，安全可靠地实现了计算算子与通信算子的并行执行。

3. 已验证的性能收益：
离线测试与Benchmarks结果均表明，该特性取得了实质性优化效果。在并发请求场景下，双流并行在TTFT（首Token延时）、等关键指标上表现稳定，证明了其在真实负载下的有效性与可靠性。

# python # 人工智能
Logo
AtomGit开源社区

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念，把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起，为开发者提供从开发、训练到部署的一站式体验。

更多推荐

cover

报名开启｜CANN 开发者 Meetup PyPTO 算子编程专场，解锁复杂算子开发新范式

avatar AtomGit开源社区
cover

共筑昇腾算子开源生态｜CANNBench Season1 沙龙收官，Season2 全新评测集发布

avatar AtomGit开源社区
cover

一念上手，“码道”成功｜OPC 进化论北京站 Meetup 报名开启

avatar AtomGit开源社区