Performance bottleneck in TransformerLayer and Attention when using attention mask

维护者通常 2 天内回复

评估

这个 Issue 还没有评估数据。

描述

Describe the bug

Dear Nvidia experts,

When attention mask is used, these three lines will cause very significant slow down since they are looping over all the items in the batch.

Performance bottleneck in TransformerLayer and Attention.

function get_indices

assertion the type of attention mask

similarly another assertion the type of attention mask

See below for profiling trace:

As you can see from the largest bar from the bottom rows, operations taking most of the time in each transformer is now transformer_engine/pytorch/attention/dot_product_attention/utils.py(1518): get_indices, Similarly the other two assertion is also cause significant slow down.

Steps/Code to reproduce bug

This problem should manifest in any profiling run with data that have attention mask in the input of forward()

Expected behavior

Multiple times of slow down when calling forward of TransformerLayer.

环境准备

从这里开始

1. 先读完整个 Issue，再读项目的贡献指南。
2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
3. Fork 仓库，在一个分支上完成修改。
4. 提交 Pull Request，并在描述里引用这个 Issue 编号。

NVIDIA/TransformerEngine 的其他 Issue

• [PyTorch] fp8_cs_quantize fake implementation returns a vector inverse scale instead of a scalar

  未关闭

  ⋯反馈这个 Issue 的问题

  • 哪里不对？

  • 已被占
  • 不适合新手
  • 已失效
  • 其他…
  • 发送 取消

  难度 2/5 1-3 小时 新手友好度 82/100

  NVIDIA/TransformerEngine#3636 · 2026年10月6日 创建

  维护者通常 2 天内回复

• [Bug] Backend selection picks FA3 for training with head_dim_qk=192 / v_head_dim=128, but FA3 backward cannot run it

  可能已有人在做 @yuweih205 于 30 天前认领。 未关闭

  ⋯反馈这个 Issue 的问题

  • 哪里不对？

  • 已被占
  • 不适合新手
  • 已失效
  • 其他…
  • 发送 取消
  attention

  难度 2/5 1-3 小时 新手友好度 85/100

  NVIDIA/TransformerEngine#3481 · 4 条评论 · 2026年9月4日 创建

  维护者通常 2 天内回复

• Increase MAX_TENSOR_NUM

  未关闭

  ⋯反馈这个 Issue 的问题

  • 哪里不对？

  • 已被占
  • 不适合新手
  • 已失效
  • 其他…
  • 发送 取消
  bug

  难度 2/5 1-3 小时 新手友好度 68/100

  NVIDIA/TransformerEngine#2189 · 7 条评论 · 5 个 reaction · 2025年9月19日 创建

  维护者通常 2 天内回复

• Enabling grouped MLP fusion breaks previously working 128-aligned inputs

  未关闭

  ⋯反馈这个 Issue 的问题

  • 哪里不对？

  • 已被占
  • 不适合新手
  • 已失效
  • 其他…
  • 发送 取消
  bug

  难度 4/5 3-5 天 新手友好度 54/100

  NVIDIA/TransformerEngine#3640 · 5 条评论 · 2026年10月6日 创建

  维护者通常 2 天内回复

• [BUG] Grouped MXFP8 quantization is not concurrency safe with multiple streams

  可能已有人在做 @kainzhong 于 1 天前认领。 未关闭

  ⋯反馈这个 Issue 的问题

  • 哪里不对？

  • 已被占
  • 不适合新手
  • 已失效
  • 其他…
  • 发送 取消
  bug

  难度 4/5 3-5 天 新手友好度 25/100

  NVIDIA/TransformerEngine#3630 · 2026年10月5日 创建

  维护者通常 2 天内回复

查看 NVIDIA/TransformerEngine 的全部 Issue

相似的 Issue

• Change profile's df column order to be ("Weight", "Voter Set")

  未关闭

  ⋯反馈这个 Issue 的问题

  • 哪里不对？

  • 已被占
  • 不适合新手
  • 已失效
  • 其他…
  • 发送 取消

  难度 1/5 1 小时以内 新手友好度 82/100

  mggg/VoteKit#400 · 2026年10月7日 创建

• Permanently Delete Temporary Uploads in AGOL

  未关闭

  ⋯反馈这个 Issue 的问题

  • 哪里不对？

  • 已被占
  • 不适合新手
  • 已失效
  • 其他…
  • 发送 取消

  难度 2/5 1-3 小时 新手友好度 64/100

  agrc/palletjack#208 · 2026年10月7日 创建

• [Bug]: Explicit .bed artifacts with 8–12 columns are parsed as VCF

  未关闭

  ⋯反馈这个 Issue 的问题

  • 哪里不对？

  • 已被占
  • 不适合新手
  • 已失效
  • 其他…
  • 发送 取消
  status/needs-triage type/bug

  难度 2/5 1-3 小时 新手友好度 85/100

  PKU-YuanGroup/OpenAI4S#218 · 2026年10月7日 创建

  维护者通常 1 天内回复

• Test Leakage

  可能已有人在做 @garland3 今天认领。 未关闭

  ⋯反馈这个 Issue 的问题

  • 哪里不对？

  • 已被占
  • 不适合新手
  • 已失效
  • 其他…
  • 发送 取消

  难度 1/5 1 小时以内 新手友好度 90/100

  sandialabs/atlas-ui-3#1030 · 2026年10月7日 创建

  维护者通常 1 天内回复

• Claude Code guard for destructive commands never blocks: $CLAUDE_TOOL_INPUT is not set

  未关闭

  ⋯反馈这个 Issue 的问题

  • 哪里不对？

  • 已被占
  • 不适合新手
  • 已失效
  • 其他…
  • 发送 取消

  难度 1/5 1 小时以内 新手友好度 84/100

  aws-samples/sample-ai-persona#151 · 2026年10月7日 创建

  维护者通常 1 天内回复

更多 Python Issue