SDPA#

class sdm.nn.SDPA(num_query_heads: int, num_key_value_heads: int | None = None, query_scaling: QueryScaling | None = None, scale: float | None = None)#

Bases: Module

Scaled Dot-Product Attention (SDPA).

This module wraps torch.nn.functional.scaled_dot_product_attention() and extends it by arbitrary batch dimensions, optional inference-time batch chunking, query-scaling, and padding support for key/value pairs.

Parameters:

  • num_query_heads (int) – The number of query attention 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_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 the default value of 1 / sqrt(channels).

forward(query: Tensor, key: Tensor, value: Tensor, seqused_key_value: Tensor | None = None, attn_mask: Tensor | None = None) → Tensor#

The forward pass.

Parameters:

  • query (Tensor) – The query tensor with shape [..., Q, Hq, C]. Q is the query sequence length, Hq is the number of query attention heads (num_query_heads), and C is the channels per head.

  • key (Tensor) – The key tensor with shape [..., KV, Hkv, C]. KV is the key/value sequence length and Hkv is the number of key/value heads (num_key_value_heads).

  • value (Tensor) – The value tensor with shape [..., KV, Hkv, C].

  • 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.

Returns:

Tensor with shape [..., Q, Hq, C].

Return type:

Tensor