Ptn Multihead Attention

This is the heart of the pack's build-a-transformer-from-scratch workflow. PtnMultiheadAttention wraps PyTorch's nn.MultiheadAttention as a node, so you get a real, optimized attention block without hand-rolling the math. You reach for it when you're assembling a transformer encoder piece by piece - attention in, contextualized sequence out - which is exactly the mha.json example workflow the README points to (85% IMDB validation accuracy from a single encoder block).

How it works

Multi-head attention projects the input into queries, keys, and values, splits them across num_heads heads, computes scaled dot-product attention per head, concatenates, and projects back out. The node is a thin wrapper over nn.MultiheadAttention - same math, same optimizer - with one convenience that matters: it expects an attention mask in the pack's convention (1 = valid token, 0 = padding) and internally inverts it to PyTorch's key_padding_mask format. So the mask you feed from the dataset nodes works directly, with no flipping.

The inputs that matter

• embed_dim (default 512) - the token feature dimension (d_model). Must be divisible by num_heads.
• num_heads (default 8) - how many attention heads to split into. 8 heads over 512 dims = 64 dims per head, the classic ratio.
• dropout (default 0) - dropout on attention weights during training.
• bias (default True) - bias on the projection layers.
• kdim / vdim (default 512 each) - dimensions for key/value projections; leave equal to embed_dim for standard self-attention, which is what this pack uses (it feeds the same tensor as query, key, and value).
• batch_first (default True) - (batch, seq, features) layout.
• add_bias_kv / add_zero_attn (default False each) - niche options; the transformer workflows leave them off.

One PTMODEL out, ready to chain into a residual block and LayerNorm.

How you'd use it

The mha.json workflow shows the full pattern: PtnEmbedding (vocab 32000, dim 256) → attention block → feedforward block → masked mean pooling → linear head, with the mask threaded through via Ptn Chained Model With Attention Mask. This node goes where the "attention block" sits, usually wrapped in Ptn Residual Connection Model With Attention Mask so the encoder block gets its residual add.

Installing

Same as every node in the pack. ComfyUI Manager → search "Pt-Wrapper", or:

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

Restart ComfyUI; the pack's requirements install on first launch.

Where people get burned

embed_dim % num_heads != 0 is the classic instant error - fix the numbers, not the graph. The subtler trap is the mask convention: if you feed a mask where 1 means padding, the node's inversion makes everything exactly backwards and attention silently attends to padding. And if you're porting this from raw PyTorch, remember batch_first=True is this pack's default, unlike nn.MultiheadAttention's.