Ptn Multihead Attention Custom

PtnMultiheadAttentionCustom is the same idea as PtnMultiheadAttention - multi-head self-attention as a node - except the attention math is written out by hand, step by step, in the pack's source. You reach for this one when you want to learn what's happening inside attention, or when you want to tweak the internals yourself. The pack's mha_custom.json workflow uses it and lands at 85–86% validation accuracy on IMDB, essentially tied with the built-in version - which is the whole point: the hand-rolled one works just as well, it's just more legible.

How it works

Instead of calling nn.MultiheadAttention, the node defines its own module. In forward it runs the input through separate q_proj, k_proj, v_proj linear layers, reshapes into (batch, heads, seq, head_dim) with einops, computes q @ kᵀ / sqrt(head_dim), applies the mask (1 = valid, 0 = padding, replacing padded positions with -inf before softmax), softmaxes, multiplies by the values, and projects back through out_proj. There's an assert that embed_dim is divisible by num_heads, and the code even documents the mask handling in comments. It's the pack's equivalent of the PyTorch tutorial's attention, minus the obscuring library call.

The inputs that matter

• embed_dim (default 512) - token feature dimension, must divide evenly by num_heads.
• num_heads (default 8) - number of heads.
• dropout (default 0) - dropout applied after the output projection.
• bias (default True) - bias on all four projection layers.
• kdim / vdim (default 512) - key/value projection input dims; equal to embed_dim for self-attention.
• batch_first (default True) - (batch, seq, features) layout.

Note this version omits add_bias_kv and add_zero_attn that the built-in wrapper has - it's the leaner, more readable subset.

How you'd use it

Same slot in the encoder as the built-in version: chained via Ptn Chained Model With Attention Mask, wrapped in a residual block, then LayerNorm and feedforward. If you're following the transformer-from-scratch guide and want to see the attention, the guide explicitly points at this node's source file as the reference implementation.

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

Same shape rules as the built-in: embed_dim not divisible by num_heads fails an assert, and the mask must be 1=valid/0=padding or attention attends to the wrong positions. Because it's hand-rolled, it's also a touch less numerically optimized than PyTorch's fused implementation - fine for a single encoder block, something to keep in mind if you scale to 12 blocks and wonder why training is slower than you expected.