ComfyUI Node

Pt Bmm

The batched matrix multiply behind attention, without the Python

By HowToSD·Created about a year ago·Updated about a year ago· 7
Pt Bmm
  • tens_a
  • tens_b
  • TENSOR

If you're here because a workflow told you to use bmm and you have no idea what it does, here's the short version: it's matrix multiplication for a whole batch of matrices at once. Where a regular matmul multiplies two 2D grids, Pt Bmm multiplies two 3D tensors of shape (batch, m, n) and (batch, n, p) and returns (batch, m, p) - the same multiply-accumulate dance, run batch times in one call.

The reason this matters in this pack specifically: batched matmul is the inner workhorse of multi-head attention. The pack has a whole "build a Transformer from scratch" guide, and attention's Q @ K^T / sqrt(d) and scores @ V steps are exactly the kind of per-batch, per-head multiplies that bmm exists for. You won't wire this into a KSampler; you'll wire it into a model-building graph.

Inputs

  • tens_a - a 3D TENSOR, shape (batch, m, n).
  • tens_b - a 3D TENSOR, shape (batch, n, p), so it's compatible with tens_a.

Output: one TENSOR, shape (batch, m, p).

Where people get burned

The node is strict: both inputs must be 3D, and it says so with an explicit ValueError rather than a cryptic backend error - "torch.bmm() only supports batch matrix multiplication for 3D tensors." If your tensors are 2D, you don't want bmm at all; use the pack's PtMatMul or PtMm. And the batch dimension, the inner n dimension, and the dtypes all have to line up, or you'll get a shape mismatch deep inside torch's error messages. If a transformer graph keeps failing right around here, a PtShowSize node on both inputs is the fastest way to see what actually arrived.

Also worth knowing: there's no precision surprise here - it's standard float matmul, so use fp32 unless you've explicitly gone bf16. The pack ships PtToFloat16/PtToBfloat16 if you're chasing memory on a big model build.

Installing it

Ships in ComfyUI-Pt-Wrapper by Hide Inada (HowToSD). ComfyUI Manager → search "ComfyUI-Pt-Wrapper" → restart, or:

cd ComfyUI/custom_nodes
git clone https://github.com/HowToSD/ComfyUI-Pt-Wrapper

This pack carries a heavy requirements.txt (transformers, datasets, scikit-learn, scipy, gensim, pandas, peft, accelerate...). If Manager's install fails, pip install -r requirements.txt in the clone. No model downloads for tensor ops - only the example training workflows auto-fetch datasets like CIFAR-10.

CategoryData Analysis

Inputs (2)

NameTypeDefaultDescription
tens_aTENSOR
tens_bTENSOR

Outputs (1)

NameTypeDescription
TENSORTENSOR