CANN Alpamayo-R1智驾优化

Alpamayo-R1 智驾VLA模型推理昇腾迁移-性能优化说明

【免费下载链接】cann-recipes-embodied-intelligence本项目针对具身智能业务中的典型模型、加速算法，提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-intelligence

本文档总结了Alpamayo-R1模型在昇腾NPU上的优化策略：

1. Flash Attention (FA) 替换

优化说明

本样例使用torch_npu内置的npu_fusion_attention融合算子替换源代码中的小算子实现，npu_fusion_attention详细可见Ascend社区文档。

实现方式

在模型配置中设置以下参数，会调用torch_npu内置的npu_fusion_attention融合算子：

config.attn_implementation = 'flash_attention_2'

优化位置

• 文件：src/alpamayo_r1/models/base_model.py
• 修改：在模型初始化时设置attention实现方式

2. Transformers库中布尔索引优化

优化说明

原始实现中使用了布尔索引操作，效率较低。通过优化索引计算逻辑，避免不必要的布尔运算和类型转换。

实现方式

原始实现：

if image_features is not None and inputs_embeds[special_image_mask].numel() != image_features.numel():

优化后实现：创建了独立的patched_get_placeholder_mask函数，优化了image和video token的mask生成逻辑：

if image_features is not None: n_masked_elements = special_image_mask.sum().item() if n_masked_elements != image_features.numel(): raise ValueError( f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {image_features.shape[0]}" )

优化位置

• 新增文件：src/alpamayo_r1/qwen_patches/qwen3vl_optimization.py
• 新增：patched_get_placeholder_mask函数
• 新增：apply_qwen3vl_patches()函数用于应用patch

3. RmsNorm算子优化

优化说明

本样例使用torch_npu内置的npu_rms_norm融合算子替换源代码中的小算子实现，npu_rms_norm详细可见Ascend社区文档。在src/alpamayo_r1/models/action_in_proj.py的RmsNorm类中使能npu_rms_norm融合算子：

实现方式

原始实现：

def _norm(self, x): """Normalize the input tensor.""" return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) def forward(self, x): """Normalize the input tensor.""" output = self._norm(x.float()).type_as(x) return output * self.weight

优化后实现：

def forward(self, x): input_dtype = x.dtype weight_dtype = self.weight.dtype if input_dtype != weight_dtype: weight_dtype = self.weight.to(input_dtype) else: weight = self.weight # npu_rms_norm returns output = torch_npu.npu_rms_norm(x, weight, self.eps)[0] return output

优化位置

• 文件：src/alpamayo_r1/models/action_in_proj.py
• 类：RmsNorm
• 修改：使用torch_npu.npu_rms_norm替换原始实现

4. ADD算子优化：使用融合算子torch.addcmul

优化说明

在扩散模型的积分过程中遇到了x = x + dt * v这样的操作。使用昇腾NPU提供的融合算子torch.addcmul，可以将加法和乘法融合为一个算子，减少计算步骤和内存访问。

实现方式

原始实现：

x = x + dt * v

优化后实现：

x = torch.addcmul(x, dt, v)

优化位置

• 文件：src/alpamayo_r1/diffusion/flow_matching.py
• 函数：flow_matching_euler积分步进
• 修改：使用torch.addcmul替换原始的加法和乘法组合

5. Concat优化：预分配+索引赋值

优化说明

在处理多特征拼接时，原始实现使用torch.cat进行多次拼接，这会导致多次内存分配和数据拷贝。优化方案是预先分配完整大小的tensor，然后通过索引直接赋值，避免拼接操作。

实现方式

原始实现：

action_feats = torch.cat([s(x[:, :, i]) for i, s in enumerate(self.sinus)], dim=-1) timestep_feats = self.timestep_fourier_encoder(timesteps[..., -1]) timestep_feats = timestep_feats.repeat(1, T, 1) x = torch.cat((action_feats, timestep_feats), dim=-1)

优化后实现：

batch_feats = torch.empty(B, T, self.num_input_feats, device=x.device, dtype=x.dtype) offset = 0 for i, s in enumerate(self.sinus): feat = s(x[:, :, i]) batch_feats[:, :, offset:offset+feat.size(-1)] = feat offset += feat.size(-1) timestep_feats = self.timestep_fourier_encoder(timesteps[..., -1]) batch_feats[:, :, offset:] = timestep_feats.repeat(1, T, 1) x = batch_feats

优化位置

• 文件：src/alpamayo_r1/models/action_in_proj.py
• 类：PerWaypointActionInProjV2
• 方法：forward
• 修改：预分配batch_feats，使用索引赋值替代torch.cat

优化总结

【免费下载链接】cann-recipes-embodied-intelligence本项目针对具身智能业务中的典型模型、加速算法，提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-intelligence