CANN/cann-recipes-infer：DeepSeek-R1预填充优化

基于Atlas A3集群的DeepSeek-R1模型prefill阶段推理性能优化实践【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer概述本文主要介绍基于Atlas A3 系列产品的DeepSeek-R1模型prefill阶段并行部署和性能优化策略。主要采用了混合切分策略、Double-Routing、micro-batch等关键优化技术最终在Atlas A3 系列产品实现perfect eplb下6.5k tokens/卡/s的吞吐性能。MLA (Multi-Head Latent Attention)部署策略优化DP (Data Parallelism)切分策略prefill在MLA部分可以采用纯DP计算策略每卡部署完整模型权重不对模型做任何切分该策略有以下优点纯DP切分可以避免多卡冗余KV-Cache (Key-Value Cache)每卡只需要存储本卡所接收请求的KV-Cache数据量各卡之间的KV-Cache互不相同无冗余数据。MLA部分无通信算子可以避免通信开销prefill阶段更容易用满算力在请求数据量较大时更易获得高吞吐。但是该切分策略会导致以下问题计算负载不均实际场景中输入请求的序列长度往往参差不齐处理较短序列的卡会提前完成计算随后必须等待批次中最长序列的处理完成导致计算周期被浪费。TTFT (Time to First Token)被延长若实时请求数量低于卡数部分卡将处于空闲状态若为凑齐完整批次而延迟处理会延长首token生成时间若直接处理不完整批次则会导致计算资源利用不充分。针对上述问题通过对TP (Tensor Parallelism)切分策略改进本样例提出SP-TP-SP (Sequence Parallelism, SP)混合切分策略对以上问题进行优化。SP-TP-SP混合切分策略针对输入序列采用batch和seq_len合轴的方式转换为token数量对token进行SP切分均匀的分发到各个卡上计算过程如下图所示将多个请求序列拼接后并padding到卡数的最小整数倍进行SP切分此时每张卡获取的token数相同。FlashAttentionFA计算部分采用TP策略通过将Attention head均匀分布到每张计算卡上来实现确保计算负载在多卡间均衡分布。FlashAttention计算完成之后通过AlltoAll通信将TP转换为SP切分通信算子的通信量更小性能更优此时o_proj由于需要全载权重占用内存会变高。上述切分策略每张卡的MLA阶段计算量相等计算负载均匀。假设有3个输入input0input1和input2把batch和seq_len合轴后将token均匀分布到Rank0~Rank3上down_proj包含q_a_proj和kv_a_proj计算q_up_proj和kv_up_proj分别表示q_b_proj和kv_b_proj计算此时每张卡计算数据量一致计算负载均匀。具体计算流程如下图所示上述策略有以下优势TTFT更短尤其是长序列场景比如1个128K的请求此时可以将计算均匀分布到各个卡上完成计算获得更短时延。MLA部分的计算负载均匀输入序列只需要padding到卡数的倍数即可各个卡的计算量相等几乎无冗余计算。但是该策略会引入以下问题需要通过相应的优化手段解决每卡需要存储完整的KV-Cache各卡之间的KV-Cache存在冗余占用内存高可通过将KV-Cache数据分卡存储进行规避。相对于纯DP策略MLA计算引入AllGather和AlltoAll通信该通信时间可以进一步通过micro-batch进行优化。总结当前在prefill场景中当请求量较大或KV-Cache内存占用较高时可以考虑采用DP切分策略以提升吞吐量当请求数量较少、分布不均匀或单次请求序列长度较长时可以考虑采用SP-TP-SP混合切分策略以降低时延。MoE (Mixture of Experts)部署策略优化MoE计算阶段采用EP (Expert Parallelism)切分策略即将路由专家均匀分布到每张卡上同时每张卡上同步存放一个共享专家。路由专家采用Double-Routing的计算策略完成计算其核心在于每张卡需要完成2次init-routing和finalize-routing操作。该方案有以下优点当负载相对较均匀时相对于AllGather的方案峰值内存较低每张卡只需要获取对应专家计算的token即可不存在数据冗余。相比于每张卡上获取完整token数据量通信量更低通信耗时更优。该方案的主要计算流程如下图所示主要计算步骤如下调用npu_moe_init_routing_v2接口对本卡上的token进行expand并排序得到expand_x同时获取每个专家需要计算的token个数即tokens_per_expert以及用于恢复token原始排序的参数expanded_row_idx第一次routing操作。expanded_x, expanded_row_idx, tokens_per_expert, _ torch_npu.npu_moe_init_routing_v2(...)通过对tokens_per_expert进行AlltoAll操作和ReduceSum操作获得每张卡需要向其他卡发送的token数量(input_splits)以及每张卡需要从其他卡上拿到的tokens数量output_split该参数主要用于AlltoAll通信各卡分别获取对应的token。调用AlltoAll通信接口对每张卡上的token进行交互每张卡上的专家获取本专家需要计算的token。dist.all_to_all_single(gathered_tokens, expanded_x, output_splits, input_splits, groupmoe_ep_group)此时每张卡上已经获取到本卡专家需要计算的token但是处于乱序。因此本卡需要再完成一次re_routing操作得到排序后的token分布即hidden_states_ordered_by_experts以及本卡上每个专家需要计算的token数量tokens_per_local_expert同时得到用于恢复本卡token顺序的参数gathered_ids_unsort第二次routing操作。hidden_states_ordered_by_experts, _, gathered_ids_unsort, tokens_per_local_expert torch_npu.npu_moe_re_routing(...)完成路由专家和共享专家计算。hidden_states_ordered_by_experts experts(**args)对完成计算的token进行排序恢复第一次finalize-routing操作与第二次routing操作相对应此处采用torch.index_select()接口性能更优。new_x torch.index_select(hidden_states_ordered_by_experts, 0, gathered_ids_unsort.float().argsort().int())调用AlltoAll通信接口进行token多卡交互把本卡上的token获取回来。dist.all_to_all_single(gathered_tokens, new_x, input_splits, output_splits, groupmoe_ep_group)调用finalize_routing接口对本张卡上的token进行顺序恢复得到原始排序的token第二次finalize-routing操作与第一次routing操作相对应同时可直接通过该接口完成路由专家和共享专家的加权求和。hidden_states torch_npu.npu_moe_finalize_routing( gathered_tokens, skip1hidden_states_share, skip2None, biasNone, scalestopk_weight.to(gathered_tokens.dtype), expanded_src_to_dst_rowexpanded_row_idx, export_for_source_rowNone, drop_pad_mode2 )TND格式支持常规FlashAttention采用BNSD的格式[batch, num_of_head, seq_len, head_dim]计算而对于每个请求长度均不确定的场景可能需要padding操作导致计算资源的浪费。而TND格式可以对batch和seq_len进行合轴处理即[batch*seq_len, num_of_head, head_dim]转换为token数无需padding数据即可完成计算此功能依赖于FlashAttention的功能支持。针对TND[token, num_of_head, head_dim]格式需要把所有输入的请求合成一个list进行处理。假设有4个请求seq_len长度分别为128256512和1024在FlashAttention计算时需要把所有的请求合成一个list即actual_seq_lengths_kv[128, 256, 512, 1024]即可完成对应计算无需把请求padding到统一的长度。尤其对于不定长输入的场景TND格式解决了易用性以及计算资源浪费的问题。但是实际计算时TND格式会引入以下问题数据排布不亲和对于FlashAttention计算不亲和数据访问不连续FlashAttention带宽效率较低。引入额外计算开销原始BNSD的格式需要使用transpose变成TND格式引入一些非必要的计算开销。为避免上述问题实际在FlashAttention计算中采用NTD的格式提升带宽效率的同时也减少了transpose等计算开销。在prefill阶段MLA部分实际的计算流程如下图所示采用该计算方式的优点如下相对于TND格式FlashAttention的计算效率提升15%。配合matmul权重拆分以及FlashAttention支持独立的NoPE(No Positional Encoding)和RoPE(Rotary Positional Encoding)输入可以完全消除split/concat/transpose等额外计算开销提升MLA阶段的计算效率。micro-batch优化prefill阶段通常属于计算bound的场景同时又有一定比例的通信耗时开销通信耗时占比可达30%。而计算和通信由不同的硬件单元完成如果此时计算和通信能够并行运行将进一步提升prefill阶段的性能。理想情况下通信耗时可以完全被计算掩盖这样计算效率将极大提升。因此可以把输入拆分成两个micro-batch采用双流的方式构造两条流水线通过定制化的流水排布把计算和通信并行起来最大化prefill阶段的计算效率。将输入拆分成两个micro-batch数据被一分为二会改变prefill阶段每一层计算的shapeshape的变动会导致计算和通信的效率发生变化同时在eager模式下task下发的数量也会多一倍对host要求更高。因此在prefill阶段做micro-batch优化时需要确认以下问题的影响计算和通信性能线性度较好要求输入数据切分前计算和通信的性能近似等于切分后两份数据的计算和通信性能加和这样可以最小化避免因shape变化导致的性能劣化。计算未出现host bound如果模型本身就是host bound场景此时通过双流实现micro-batch会导致host性能劣化整体耗时劣化。MLA-DP和MoE-EP场景micro-batch策略MLA采用DP切分的情况下整个prefill阶段的通信算子主要来源于MoE阶段的AlltoAll通信dispatch阶段和combine阶段。因此该阶段micro-batch的核心目标就是通过计算掩盖dispatch和combine两个阶段的通信耗时最大化prefill阶段的计算效率。 该场景的流水排布如下图所示其主要步骤如下计算和通信划分将prefill的计算划分为6个阶段分别为atten、init_routing、dispatch、mlp、combine、finalize_routing 图中 atten主要包含MLA和o_proj计算dispatch主要包含将token分发到对应专家的AlltoAll通信mlp主要包含路由专家和共享专家计算combine主要包含把token恢复到各卡上的AlltoAll通信。micro-batch预处理将输入划分成2个micro-batch进行处理假设每卡上面有2个4k请求则每个micro-batch处理1个4k请求当2个micro-batch全部计算完成时则完成所有的计算处理。通信处理由于通信和计算并行此时不能使能HCCL_OP_EXPANSION_MODE的AIV模式采用默认SDMA模式。如果采用AIV模式则会导致通信和计算硬件资源抢占时延变长。创建流针对DPEP的切分策略场景将所有的计算放在默认计算流上。curr_stream torch.npu.current_stream()新建一条流用于通信stream1 torch.npu.Stream()流水排布将计算放到当前流上进行如下所示with torch.npu.stream(curr_stream): # 在curr_stream进行atten计算 atten()将通信放到stream1上进行计算with torch.npu.stream(stream1): # 在stream1进行atten计算 dispatch()针对需要同步的位置可以采用以下方式进行两条流的同步操作event0 curr_stream.record_event() ... stream1.wait_event(event0)通过使能micro-batch计算dispatch和combine阶段的通信完全被计算掩盖整个prefill阶段完全处于计算bound的状态芯片侧无空闲时间因此端到端的计算效率得到提升。基于Atlas A3 系列产品prefill采用2机共计16卡环境部署输入序列为64*4K假设MoE层专家负载均匀条件下性能测试结果如下切分策略时延(ms)总吞吐(tokens/s)平均吞吐(tokens/卡/s)MLA-DP32 MoE-EP323128838055237MLA-DP32 MoE-EP32 micro-batch25121043566522该切分策略下开启micro-batch特性可获得25%的性能提升。Matmul类算子使能weightNZ格式针对DeepSeek模型中的Matmul类计算可以将模型权重中的weight Tensor处理成昇腾AI处理器亲和的格式weightNZ这样可以进一步提升Matmul的计算效率。weightNZ格式为512Bytes对齐的格式可以有效提升硬件的带宽访问效率以int8类型的Tensorshape是[16384, 7168]为例转换为NZ格式为[224, 1024, 16, 32]附录环境部署以及样例执行【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考