vLLM-Ascend常见模型部署样例：Qwen3稠密模型双流并行-优快云博客

作者：昇腾PAE技术支持团队

昇腾案例库简介：https://agent.blog.youkuaiyun.com/article/details/155446713
昇腾案例抢鲜预览：https://gitcode.com/invite/link/8791cccc43cb4ee589e8
（如对本文有疑问，请移步案例库提交issue，专人答疑）

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个：

positions：模型输入，表示当前token的位置信息，shape=[1, bs]；
hidden_states：模型输入，表示当前token的特征向量，shape=[bs, hidden_dim]；
residual_post：模型输入，用于残差计算，需要注意，在第0层Encoder的时候值为None，需要特殊处理；
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

如何正确拆分数据？

选择"拆分点"
- 简单方法。直接通过 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，则不走双流并行。

拆分数据
- 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.88	3844.99	8084.73
双流	4046.86	3587.15	7526.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）机制，利用独立通信流与主流进行同步，安全可靠地实现了计算算子与通信算子的并行执行。