HF Blog · PyTorch 注意力机制性能分析（Profiling 实战 Part 3）

Jay 草稿 | 2026-07-21

元数据

来源: Hugging Face × PyTorch 官方 Blog
URL: https://huggingface.co/blog/torch-attention-profile
发布日期: 2026-07-10
可信度: ⭐⭐⭐⭐⭐（HF + PyTorch 官方）
分类: Performance Engineering / PyTorch / CUDA / Profiling

核心内容

PyTorch Profiling 系列 Part 3，聚焦 Transformer attention 的性能剖析方法论。

剖析路径（Naive → Inplace → SDPA → Custom Kernels）

1. Naive 实现（逐算子）

class NaiveCausalAttention(nn.Module):
    def forward(self, q, k, v, mask):
        scores = torch.matmul(q, k.transpose(-2, -1))  # matmul
        scores = scores * self.scale                      # mul（scaling）
        scores = scores.masked_fill(mask, float("-inf"))  # masked_fill
        attn = torch.softmax(scores, dim=-1)              # softmax
        return torch.matmul(attn, v)                      # matmul

Profiler trace 揭示：matmul → mul → masked_fill → softmax → matmul 五步分离执行

2. SDPA（Flash Attention 等后端）

调用 torch.nn.functional.scaled_dot_product_attention，底层 kernel 融合

3. Custom CUDA Kernels

手动编写 fused kernel，极致优化

实战工具链

uv run 04_a_naive_attention.py
uvx trace-util -f traces/ -b <hf_uname>/traces

脚本列表（全开源）

04_a_naive_attention.py — 朴素实现
04_b_inplace_ops_attention.py — Inplace 算子优化
04_c_sdpa_attention.py — SDPA 后端
04_d_kernels_attention.py — Custom CUDA Kernel

脚本地址: https://huggingface.co/datasets/ariG23498/profiling-pytorch

评价

工程价值高。Attention 是 LLM 推理和训练的第一性瓶颈，这篇文章提供了系统化的 profiler 使用方法论，可直接用于团队性能调优 SOP。适合作为 PyTorch profiling 系列收录入知识库。

建议行动

纳入性能调优 SOP
对比系列 Part 1（基础算子）和 Part 2（MLP fusion）补全知识体系