TransformerBlock#

class sdm.nn.TransformerBlock(channels: int, num_query_heads: int, mlp: Module, 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, scale: float | None = None, 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) – 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.

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

  • 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, *, 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, *, 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, *, 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 (Tensor) – The query tensor with shape [..., Q, C]. Q is the query sequence length, C is the number of channels.

  • key_value (Tensor | KVCacheEntry | None) – 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 (Tensor | None) – Valid key/value lengths with shape [...] and torch.int32 dtype.

  • attn_mask (Tensor | None) – Boolean attention mask with shape [..., Q, KV]. Entries set to True participate in attention.

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

  • batch_size_limit (int | Literal['auto'] | None) – Maximum number of batch elements processed at once.

  • out (Tensor | None) – 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.

Return type:

Tensor | tuple[Tensor, KVCacheEntry]