CANN SuperKernel Attention-Only Scope 模板实战：以最小风险将 Attention 模块纳入算子二进制融合

CANN SuperKernel Attention-Only Scope 模板实战以最小风险将 Attention 模块纳入算子二进制融合【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer导读本模板是 cann-recipes-infer 项目中 SuperKernel 适配技能model-infer-superkernel提供的三种 Scope 模板之一面向首次尝试 SuperKernel性能瓶颈集中在 Attention 计算需要最小风险优化方案等场景给出仅将 Attention 模块纳入 SuperKernel 融合范围的完整实现方案。读完本文后你将掌握superkernel_scope上下文管理器的两种标记写法Attention 内部标记与 Decoder 层标记、配套的 config.yaml 配置方法以及从编译验证、功能对等到性能验证的完整落地流程并能对照仓库源码理解其底层实现原理。一、模板定位与适用场景在启用 SuperKernel 算子二进制融合技术时首要决策是确定融合范围Scope。Attention-Only 模板将融合范围严格限制在 Attention 模块Q/K/V 投影、Attention 计算、O 投影内部其余模块Add Norm、MLP/MoE 等全部排除在 Scope 之外。该模板适用场景包括首次尝试 SuperKernelScope 范围小问题易定位适合验证 SuperKernel 在当前模型上的有效性模型性能瓶颈主要在 Attention 计算针对性融合收益最直接需要最小风险的优化方案对现有推理链路影响面最小标准 Transformer 架构模型Attention 占比高且结构规整。优缺点对照维度说明优点风险最小Scope 范围小问题易定位易于调试编译快、调试简单稳定性高不影响其他模块适合首次尝试快速验证 SuperKernel 是否有效缺点性能提升有限仅优化 Attention 部分未充分利用其他计算密集型模块如 MoE/FFN 未优化从本技能整体视角看Attention-Only 是三种模板中风险最低、验证最快的切入点。完整决策树与各模板对比见 SKILL.mdScope 分析方法详见 scope-analysis-guide.md。二、前置条件SuperKernel 生效的硬性约束在使用本模板前必须确认满足以下条件来自 SKILL.md 与 docs/cann/zh/super_kernel.md约束项要求说明执行模式exe_mode: ge_graph不支持eager与aclgraph模式硬件Atlas A3 系列通过npu-smi info确认框架PyTorch基于 torchair 作用域接口标定生效阶段仅 Decode 阶段Prefill 阶段输入长度动态变化SuperKernel 自动禁用标定方式手动标记 Scope使用superkernel_scope上下文管理器融合范围连续可融合算子遇到 TBE 等不可融合算子会自动断段仓库源码层面executor/utils/common_utils.py 给出了superkernel_scope的实现启用时透传至torchair.scope.super_kernel(scope, options)禁用时返回FakeContextManager空实现——这也是为什么enable_superkernelFalse时代码可以零成本保持原逻辑def superkernel_scope(enable: bool, scope: str, options: str None): if enable: return tng.scope.super_kernel(scope, options) else: return FakeContextManager()关于 SuperKernel 的技术原理ICache Preload、Early-Start、同步优化、子 Kernel 复制、Tiling 下沉与 Weight 预取、双流并发融合等算子级/网络级优化手段以及npugraph_ex/ GE 图模式两种开启方式可参阅 docs/cann/zh/super_kernel.md。三、实现方案 A在 Attention 模块内部标记适用条件Attention 模块独立实现如Attention(nn.Module)单独成类。在models/{model_name}/models/modeling_*.py中于 Attention 的forward方法内部用superkernel_scope包裹 Q/K/V 投影、Attention 计算与 O 投影# models/{model_name}/models/modeling_*.py from executor.utils import superkernel_scope class Attention(nn.Module): def __init__(self, config): super().__init__() self.enable_superkernel config.enable_superkernel # ... 其他初始化 def forward( self, hidden_states, attention_maskNone, position_idsNone, past_key_valueNone, is_prefillFalse, **kwargs ): # SuperKernel scope 仅包含 Attention 计算 with superkernel_scope( self.enable_superkernel and not is_prefill, labelattention, optionstream-fusion1 ): # Q, K, V 投影 query_states self.q_proj(hidden_states) key_states self.k_proj(hidden_states) value_states self.v_proj(hidden_states) # Attention 计算 attn_output self._compute_attention( query_states, key_states, value_states, attention_mask, past_key_value ) # O 投影 attn_output self.o_proj(attn_output) return attn_output, past_key_value关键点说明开关与阶段双重控制self.enable_superkernel and not is_prefill保证仅在配置开启且处于 Decode 阶段时才真正进入融合 Scope。enable_superkernel来自模型配置中的custom_params由各模型 runner 读取并注入。label语义同一label表示属于同一个融合范围由用户指定不同 Attention 模块如attention、attention_layer_{idx}会形成各自的融合段。option语义SuperKernel 编译选项模板中stream-fusion1表示在融合范围内开启多流并发对应双流并发融合能力。四、实现方案 B在 Decoder 层中标记适用条件Attention 作为 Decoder 层的一部分DecoderLayer中包含 Attention 与 MLP 等子模块。将 Scope 标记放在 Decoder 层的forward中仅包裹 Self-Attention 相关计算Add Norm 与 MLP 留在 Scope 外# models/{model_name}/models/modeling_*.py from executor.utils import superkernel_scope class DecoderLayer(nn.Module): def __init__(self, config, layer_idx): super().__init__() self.enable_superkernel config.enable_superkernel self.self_attn Attention(config, layer_idx) self.mlp MLP(config) # ... 其他初始化 def forward( self, hidden_states, attention_maskNone, position_idsNone, past_key_valueNone, is_prefillFalse, **kwargs ): residual hidden_states # SuperKernel scope 仅包含 Attention with superkernel_scope( self.enable_superkernel and not is_prefill, labelfattention_layer_{self.layer_idx}, optionstream-fusion1 ): # Self-Attention hidden_states self.input_layernorm(hidden_states) attn_output, past_key_value self.self_attn( hidden_states, attention_maskattention_mask, position_idsposition_ids, past_key_valuepast_key_value, is_prefillis_prefill ) # Add Norm在 Scope 外 hidden_states residual attn_output hidden_states self.post_attention_layernorm(hidden_states) # MLP在 Scope 外 residual hidden_states hidden_states self.mlp(hidden_states) hidden_states residual hidden_states return hidden_states, past_key_value两种方案对比维度方案 A模块内部标记方案 BDecoder 层标记标记位置Attention 类forward内部DecoderLayer 类forward内部融合内容QKV 投影 Attention 计算 O 投影input_layernorm Self-Attentionlabel 建议attentionattention_layer_{layer_idx}适用结构Attention 独立成类Attention 嵌在 DecoderLayer 中两种方案的共同原则是Add Norm、MLP/MoE 等非 Attention 计算一律保持在 Scope 外以此保证融合范围最小、风险最低。仓库真实案例印证上述标记方式与仓库中 DeepSeek-R1 样例的实现思路一致。在 models/deepseek_r1/models/modeling_deepseek.py 中模型 forward 外层通过superkernel_scope包裹整个 Decode 层循环并根据是否启用多流选择不同编译选项label fdecode_layer if self.enable_multi_streams: option stream-fusion1 # if multi_streams is enabled, enable multi stream in superkernel else: option option_xxx2 with superkernel_scope(self.enable_superkernel and not forward_metadata.is_prefill, label, option):这说明stream-fusion1选项在仓库实践中正是配合多流并行enable_multi_streams使用而 Attention-Only 模板中的写法self.enable_superkernel and not is_prefill与仓库中forward_metadata.is_prefill的控制逻辑完全同构。关于 DeepSeek-R1 decode 优化整体方案可参考 docs/models/deepseek_r1/deepseek_r1_decode_optimization.md。五、配置文件在models/{model_name}/config.yaml中开启 SuperKernel 并保持其余优化项按需配置# models/{model_name}/config.yaml model_config: exe_mode: ge_graph # 必须是 ge_graph enable_cache_compile: False # 可选缓存编译 custom_params: enable_superkernel: True # 启用 SuperKernel enable_multi_streams: False # 可选多流并行参数说明参数取值说明exe_modege_graph必须为ge_grapheager/aclgraph不支持 SuperKernelenable_cache_compileFalse/True可选开启编译缓存加速二次编译enable_superkernelTrue/FalseSuperKernel 总开关对应代码中self.enable_superkernelenable_multi_streamsFalse/True可选开启多流并行与stream-fusion1选项联动源码对应关系enable_superkernel在 executor/utils/graph_utils.py 中被读取并映射为 GE 编译配置super_kernel_optimizeenable_superkernel model_config.custom_params.get(enable_superkernel, False) ... super_kernel_optimize: enable_superkernel,仓库真实配置示例models/deepseek_r1/config/decode_r1_rank_128_128ep_a8w8.yaml进一步印证了该配置结构exe_mode: ge_graph # [eager, npugraph_ex, ge_graph] enable_cache_compile: False # [False, True] custom_params: enable_multi_streams: True # [False, True] enable_superkernel: True # [False, True]配置修改属于模型侧适配。各模型的具体推理入口如bash infer.sh或python infer.py --yaml_file_path config.yaml以及完整配置字段说明可查看对应模型的 README 与 docs/common/inference_config_guide.md。六、验证步骤1. 编译验证cd models/{model_name} bash infer.sh 21 | tee compile.log # 检查编译日志 grep -i superkernel compile.log grep -i attention compile.log验证要点编译成功无报错日志中能看到 SuperKernel / attention 相关融合信息确认 Scope 范围被正确识别。2. 功能验证对比启用前后的输出确保精度一致# 对比启用前后的输出 # 1. 禁用 SuperKernel sed -i s/enable_superkernel: True/enable_superkernel: False/ config.yaml bash infer.sh baseline_output.txt # 2. 启用 SuperKernel sed -i s/enable_superkernel: False/enable_superkernel: True/ config.yaml bash infer.sh optimized_output.txt # 3. 对比输出 diff baseline_output.txt optimized_output.txt # 应该完全一致3. 性能验证# 使用性能测试脚本 python scripts/benchmark.py --config config.yaml --output performance.json # 查看性能提升 cat performance.json | jq .improvement_percent更规范的基线方法论工程上建议采用更完整的禁用 → 基线 → 启用 → 优化 → 对比流程核心指标包括 Decode 单步耗时ms、吞吐量tokens/s、首 token 延迟、内存占用、编译时间正式测试前先预热运行再运行 3~5 次取平均值并用变异系数CV 5% 为稳定确认基线可靠性。提升比例的计算方式为性能提升 (基线耗时 - 优化后耗时) / 基线耗时 × 100% 吞吐量提升 (优化后吞吐量 - 基线吞吐量) / 基线吞吐量 × 100%详细的环境准备npu-smi info、CANN / PyTorch / torch_npu 版本确认、性能数据 JSON 记录格式与稳定性判断标准见 performance-baseline-guide.md。该指南还提供了基于 NPU Profiler 与自定义计时的性能分析方法可用于确认 Attention 是否确为瓶颈。七、预期性能提升参考模型类型预期提升说明标准 Transformer10-15%Attention 占比约 30-40%MoE 模型5-10%Attention 占比较小长序列模型15-20%Attention 计算密集说明上表为模板文档给出的参考区间实际收益取决于模型结构、Attention 占比、硬件与编译选项需以本机实测为准。八、常见问题排查Q1: 编译失败提示不支持的算子可能原因Attention 模块中使用了不支持的算子。解决方法检查 Attention 实现确认使用的算子如果使用了 Tiling 下沉算子将其移出 Scope参考调试指南进行排查debugging-guide.md位于 references 目录。Q2: 性能提升不明显可能原因Attention 不是性能瓶颈模型结构特殊Attention 占比小。解决方法使用 profiler 分析性能瓶颈方法见 performance-baseline-guide.md如果 MoE 是瓶颈尝试 MoE-Only 模板moe-only.md如果整体都是瓶颈尝试 Full-Model 模板full-model.md。Q3: 精度不一致可能原因DCache 一致性问题如果使用自定义算子数值计算顺序变化。解决方法检查是否使用了自定义算子如果使用了添加 DCache 刷新调用kvcache-fa-precision-debugskill 进行排查。九、下一步演进路径Attention-Only 验证成功后可按以下顺序渐进扩大优化面扩大 Scope 范围尝试 full-model.md 模板将整个 Decoder 层纳入 Scope结合其他优化启用多流并行enable_multi_streams: Truestream-fusion1、融合算子等性能调优调整编译选项如option参数进一步优化性能。每一步扩大后都建议按编译验证 → 功能验证 → 性能对比流程回归并参考 scope-analysis-guide.md 中的自底向上 / 自顶向下 / 基于性能分析三种 Scope 分析方法确定最优融合范围。十、参考资源Scope 分析指南references/scope-analysis-guide.md性能基线指南references/performance-baseline-guide.mdMoE-Only 模板resources/scope-templates/moe-only.mdFull-Model 模板resources/scope-templates/full-model.mdSuperKernel 技能总览SKILL.mdSuperKernel 原理与约束docs/cann/zh/super_kernel.md上下文管理器实现executor/utils/common_utils.pyGE 编译配置映射executor/utils/graph_utils.py仓库真实用例DeepSeek-R1 decodemodels/deepseek_r1/models/modeling_deepseek.py 与 models/deepseek_r1/config/decode_r1_rank_128_128ep_a8w8.yaml【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考