TransformerBlock#

class sdm.nn.TransformerBlock(channels: int, num_query_heads: int, mlp: Module | None, num_key_value_heads: int | None = None, query_norm: Module | None = None, key_value_norm: Module | None = None, post_attn_norm: Module | None = None, query_transform: Module | None = None, key_transform: Module | None = None, query_scaling: QueryScaling | None = None, gate: Module | None = None, scale: float | None = None, dropout: float = 0.0, bias: bool = True, device: device | str | None = None, dtype: dtype | None = None)#

Bases: Module

Transformer block with normalization and feedforward residual modules.

Parameters:

  • channels (int) – The number of input and output channels.

  • num_query_heads (int) – The number of query attention heads.

  • mlp (Module | None) – Feedforward module applied after the attention residual.

  • num_key_value_heads (int | None) – The number of key/value attention heads. Defaults to num_query_heads (standard multi-head attention).

  • query_norm (Module | None) – Normalization applied to query inputs before attention.

  • key_value_norm (Module | None) – Normalization applied to key/value inputs before attention.

  • post_attn_norm (Module | None) – Normalization applied to the attention output before its residual addition.

  • query_transform (Module | None) – Transformation applied to projected query heads before scaled dot-product attention.

  • key_transform (Module | None) – Transformation applied to projected key heads before scaled dot-product attention.

  • query_scaling (QueryScaling | None) – Query scaling module to scale projected query heads before scaled dot-product attention, e.g., QASSMax.

  • gate (Module | None) – Multiplicative output gate computed from unprojected queries.

  • scale (float | None) – Scaling factor passed to torch.nn.functional.scaled_dot_product_attention(). None uses 1 / sqrt(channels_per_head).

  • dropout (float) – Attention dropout probability.

  • bias (bool) – If set to False, the module will not learn an additive bias.

  • device (device | str | None) – The device.

  • dtype (dtype | None) – The dtype.

forward(query: Tensor, key_value: Tensor | KVCacheEntry | None = None, seqused_key_value: Tensor | None = None, attn_mask: Tensor | None = None, query_pos: Tensor | None = None, key_pos: Tensor | None = None, *, is_causal: bool = False, return_key_value: Literal[False] = False, batch_size_limit: int | Literal['auto'] | None = None, out: Tensor | None = None) → Tensor#
forward(query: Tensor, key_value: Tensor | KVCacheEntry | None = None, seqused_key_value: Tensor | None = None, attn_mask: Tensor | None = None, query_pos: Tensor | None = None, key_pos: Tensor | None = None, *, is_causal: bool = False, return_key_value: Literal[True], batch_size_limit: int | Literal['auto'] | None = None, out: Tensor | None = None) → tuple[Tensor, KVCacheEntry]
forward(query: Tensor, key_value: Tensor | KVCacheEntry | None = None, seqused_key_value: Tensor | None = None, attn_mask: Tensor | None = None, query_pos: Tensor | None = None, key_pos: Tensor | None = None, *, is_causal: bool = False, return_key_value: bool, batch_size_limit: int | Literal['auto'] | None = None, out: Tensor | None = None) → Tensor | tuple[Tensor, KVCacheEntry]

The forward pass.

Parameters:

  • query – The query tensor with shape [..., Q, C]. Q is the query sequence length, C is the number of channels.

  • key_value – The key/value tensor with shape [..., KV, C] or precomputed key/value projections as a KVCacheEntry. KV is the key/value sequence length. If omitted, query is used for self-attention.

  • seqused_key_value – Valid key/value lengths with shape [...] and torch.int32 dtype.

  • attn_mask – Attention mask with shape [..., Hq, Q, KV]. Hq is the number of query heads. For boolean masks, entries set to True participate in attention. Floating-point masks act as additive bias.

  • query_pos – Query positions with shape [..., Q] passed to query_transform.

  • key_pos – Key positions with shape [..., KV] passed to key_transform.

  • is_causal – Apply an upper-left causal attention mask.

  • return_key_value – Whether to return the computed key and value projections alongside the block output.

  • batch_size_limit – Maximum number of batch elements processed at once.

  • out – The output tensor.

Returns:

Tensor with shape [..., Q, C] when return_key_value is False. Otherwise, a tuple of the output tensor and a KVCacheEntry.