Hacktoberfest 2026：维护者为十月标记出来的 issue，仍然开放、适合新手。 浏览 Hacktoberfest issue

Request for batched general_gemm() (or FP8-aware torch.bmm) for non-Linear GEMM workloads

未关闭
  • 哪里不对？
#2,846 0 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

维护者通常 2 天内回复

还没有人认领这个 Issue。

评估

难度
5/5
预计耗时
一周以上
新手友好度
35/100
Issue 类型
功能
描述清晰度
基本清楚
活跃度
冷清
技术栈
python, pytorch
领域
machine-learning, performance

调研方向

从 transformer_engine.pytorch.cpp_extensions 中的 general_gemm 开始，并将其 2D 契约与 torch.bmm 的 3D 输入进行比较。检查现有的 Float8Tensor 和 MXFP8Tensor 路径，包括 use_split_accumulator 的使用和对 backward 的支持。当无需循环遍历 batch slice，即可为所描述的训练工作负载提供带 FP32 累加的批量 FP8 GEMM 时，即视为完成。

由索引模型根据 Issue 内容生成。

描述

field-request

Is your feature request related to a problem? Please describe.

We’re accelerating triangular multiplication in a protein structure prediction model (AlphaFold-style tri-mul). The core operation is two large einsums over 4D pair representations that we’ve reshaped into batched matmuls:

# Input: (B, N, N, D) where N = 2048 (sequence length), D = 128
# After chunk, permute, reshape: (B*32, 2048, 2048)
x1 = torch.bmm(a, b.transpose(1, 2))  # B*32 independent N×N GEMMs

At N = 2048, this accounts for roughly 40% of the tri-mul compute and is heavily memory-bandwidth-bound. Currently we run in FP32 (4 bytes/element) or BF16 (2 bytes/element). MXFP8 inputs (1 byte/element) with FP32 accumulation would provide up to a 4× reduction in HBM reads, which is the dominant cost at these sizes.

However, there is currently no way to run FP8 batched matrix multiplication through TE:

  • te.autocast() only intercepts TE modules, not torch.bmm
  • Float8Tensor / MXFP8Tensor passed to torch.bmm silently dequantize to full precision
  • general_gemm() supports FP8 × FP8 with use_split_accumulator=True, but only accepts 2D inputs — looping over B*32 slices would likely negate the bandwidth savings

Related: #1910 describes the same gap for FP8 GEMM beyond te.Linear.

Describe the solution you’d like

A batched variant of general_gemm() that accepts 3D inputs and runs FP8 GEMMs across the batch dimension with FP32 accumulation:

from transformer_engine.pytorch.cpp_extensions import batched_general_gemm

# Quantize inputs to FP8
a_fp8 = mxfp8_quantizer.quantize(a_3d)  # (B*32, N, N)
b_fp8 = mxfp8_quantizer.quantize(b_3d)  # (B*32, N, N)

# Batched FP8 GEMM with FP32 accumulation
output = batched_general_gemm(
    a_fp8,
    b_fp8,
    out_dtype=torch.bfloat16,
    layout="NN",
    use_split_accumulator=True,  # FP8×FP8 multiply, FP32 accumulate
)
# output: (B*32, N, N) in BF16

Alternatively, making Float8Tensor / MXFP8Tensor dispatch torch.bmm to real FP8 tensor core GEMMs, instead of dequantizing, would also solve this.

Describe alternatives you’ve considered

  • GroupedLinear: Suggested in #1910, but it is designed for MoE-style use cases with different weights per group. Our use case is two arbitrary input tensors, not input × stored weight. It was also noted there may be significant overhead.
  • Looping general_gemm() over batch slices: Functionally possible, but Python loop overhead and the lack of kernel batching would likely wipe out the memory-bandwidth gains from FP8.
  • Skipping .float() and running torch.bmm in BF16: This is our current workaround. It gives a 2× memory reduction versus FP32, but still leaves another 2× on the table compared with FP8.

Additional context

  • Targeting Blackwell (MXFP8BlockScaling) and Hopper (DelayedScaling / CurrentScaling)
  • Training workload, so backward-pass support is needed
  • This batched FP8 GEMM pattern would also help other workloads with non-Linear matmuls, including attention (unfused path), structure prediction, graph neural networks, and any model with einsum contractions reshaped to bmm
  • TE v1.12+

Happy to provide a minimal repro or benchmark if helpful.

主要语言
Python
星标
3.6k
派生
862
平均合并
4 天 15 小时
30 天内合并 PR
58

环境准备

从这里开始

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

NVIDIA/TransformerEngine 的其他 Issue

查看 NVIDIA/TransformerEngine 的全部 Issue

相似的 Issue

更多 Python Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。