QueryScaling#

class sdm.nn.QueryScaling#

Bases: Module, ABC

Base class for query scaling modules in SDPA.

Query scaling modules transform projected query heads before scaled dot-product attention. They may use the effective key length to implement length-aware temperature scaling while preserving compatibility with torch.nn.functional.scaled_dot_product_attention().

abstractmethod forward(query: Tensor, *, key_len: Tensor | int) → Tensor#

The forward pass.

Parameters:

  • query (Tensor) – The query tensor to scale, with shape [..., S, H, C]. S is the query sequence length, H is the number of attention heads, and C is the channels per head.

  • key_len (Tensor | int) – The number of valid keys used to scale each query. An int denotes one shared key length for every query. A tensor shaped [..., 1] denotes one key length per batch item, shared across all query positions in S. A tensor shaped [..., S] denotes one individual key length per query position.

Returns:

Tensor with shape [..., S, H, C].

Return type:

Tensor