Attention#

class sdm.nn.Attention(channels: int, num_query_heads: int, num_key_value_heads: int | 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

Multi-head attention layer with grouped-query attention support.

This module owns the query, key, value, and output projections. It performs self-attention when key_value is omitted and cross-attention when key_value is given.

Parameters:

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

  • num_query_heads (int) – The number of query attention heads. channels must be divisible by num_query_heads.

  • num_key_value_heads (int | None) – The number of key/value attention heads. Setting this below num_query_heads enables grouped-query attention (GQA); setting it to 1 enables multi-query attention (MQA). Must divide num_query_heads. Defaults to num_query_heads (standard multi-head attention).

  • 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) → 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]) → 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) → 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 attention output.

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]