`MultiheadAttention`

Scaled dot-product attention over several heads.

class MultiheadAttention : Module— python/vkml/nn.py:850

Parameter layout follows torch.nn.MultiheadAttention exactly — a packed in_proj_weight of shape (3E, E) plus an out_proj submodule — so a torch state_dict loads without rearrangement.

That is not cosmetic. It is what lets the validation suite compare against torch's own implementation rather than against a reference written here, which would only prove the two agree with each other.

ⓘ Note

embed_dim must divide by num_heads; the constructor raises otherwise. Causal masking is built from triu plus a comparison, not a dedicated kernel.

⚠ Warning

Two deliberate divergences from torch, both pinned by tests:

batch_first defaults to True. torch defaults to False, meaning (S, B, E), a legacy layout almost every caller overrides.
It returns the output tensor alone, not (output, weights). The averaged per-head weights torch returns second are a debugging aid, and a tuple that is nearly always destructured-and-discarded is worse to use.

Construction¶

`init`¶

def __init__self, embed_dim: int, num_heads: int, bias: bool=True, batch_first: bool=True— python/vkml/nn.py:869

Forward¶

`forward`¶

def forwardself, query: V.Tensor, key: V.Tensor=None, value: V.Tensor=None, attn_mask: V.Tensor=None, is_causal: bool=False -> V.Tensor— python/vkml/nn.py:925

Internals¶

`repr`¶

def __repr__self -> str— python/vkml/nn.py:970

MultiheadAttention