Skip to content

Attention Layers

About

The d9d.module.block.attention package provides optimized attention mechanism implementations.

Softmax Attention

Grouped-Query Attention

GroupedQueryAttention is a Grouped-Query Attention implementation.

Due to its abstract nature it is also can be used as a Multi-Head Attention and Multi-Query Attention module.

Multi-Head Latent Attention

MultiHeadLatentAttention is an implementation of the Multi-Head Latent Attention (MLA) mechanism introduced in DeepSeek-V2.

Scaled Dot-Product Attention Backends

The attention modules delegate the core scaled dot-product computation to a pluggable backend. Available backends are:

  • "flash_attention_4" - FlashAttention 4. Supports sliding windows and learnable sinks; does not accept explicit masks.
  • "flash_attention_2" - FlashAttention 2. Supports sliding windows but not sinks or explicit masks.
  • "torch" - PyTorch's fused scaled_dot_product_attention. Optionally pins a specific SDPBackend (MATH, FLASH_ATTENTION, EFFICIENT_ATTENTION, CUDNN_ATTENTION). Does not support sinks or sliding windows.
  • "eager" - a portable, dependency-free pure-PyTorch reference. Supports GQA/MQA, causal masking, sliding windows, explicit attention masks, and learnable attention sinks. Primarily useful as a fallback and as a correctness reference for the other backends.

The backend is chosen via build_sdpa_backend(). When no explicit configuration is given, it is auto-detected based on installed kernels and the structural requirements (sinks, windows, explicit attention masks) of the layer. The auto-detected choice can be overridden through the D9D_BACKEND_AUTO_SDPA environment variable (a JSON backend configuration).

The JSON is the serialized form of the corresponding backend config, keyed by a kind discriminator. For example:

1
2
3
4
5
# Force the pure-PyTorch eager backend.
export D9D_BACKEND_AUTO_SDPA='{"kind": "eager"}'

# Force PyTorch's fused SDPA, pinning a specific kernel.
export D9D_BACKEND_AUTO_SDPA='{"kind": "torch", "backends": ["FLASH_ATTENTION"]}'

This override only affects auto-detection; passing an explicit sdpa_backend configuration to a layer takes precedence over the environment variable.

d9d.module.block.attention

Provides attention layer implementations.

GroupedQueryAttention

Bases: Module, ModuleLateInit

Implements Grouped Query Attention (GQA) with RoPE and optional QK Normalization.

This module performs the full attention mechanism pipeline: 1. Linear projection to Q, K, V. 2. Optional RMS Normalization on Q and K. 3. Rotary Positional Embedding (RoPE) application. 4. Scaled Dot Product Attention (via FlashAttention). 5. Optional sigmoid output gating. 6. Output projection.

__init__(hidden_size, num_attention_heads, num_key_value_heads, head_dim, qk_norm_eps, is_causal, rope_style, rope_dim=None, enable_output_gate=False, qk_norm_zero_centered=False, sdpa_backend=None)

Constructs the GroupedQueryAttention layer.

Parameters:

Name Type Description Default
hidden_size int

Hidden size.

 required
num_attention_heads int

Number of Query heads.

 required
num_key_value_heads int

Number of Key/Value heads. If less than num_attention_heads, GQA/MQA is enabled.

 required
head_dim int

Dimensionality of a single attention head.

 required
qk_norm_eps float | None

Epsilon for LayerNorm/RMSNorm applied to Q and K. If None, normalization is disabled.

 required
is_causal bool

Whether to apply a causal mask (auto-regressive constraint).

 required
rope_style RotaryEmbeddingStyle

Rotary embedding layout style alignment.

 required
rope_dim int | None

Dimension of the RoPE sub-vector. If None, RoPE is applied to the full head_dim.

 None
enable_output_gate bool

If True, enables sigmoid output gating (Qwen 3.5 style).

 False
qk_norm_zero_centered bool

If True, utilizes zero-centered scaling weights for the optional Q and K RMSNorm layers (DeepSeek V3 style).

 False
sdpa_backend AnySdpaBackendConfig | None

Configuration for the Scaled Dot-Product Attention backend. If None, the backend will be auto-detected via build_sdpa_backend().

 None
forward(hidden_states, attention_mask, position_embeddings)

Computes the attention operation.

Parameters:

Name Type Description Default
hidden_states Tensor

Input tensor. Shape: (batch, seq_len, hidden_size).

 required
attention_mask Tensor | None

Optional mask associated with the inputs.

 required
position_embeddings tuple[Tensor, Tensor]

Tuple of (cos, sin) tensors for RoPE application. Each tensor should be of shape (batch, seq_len, rope_dim) when partial RoPE is used, or (batch, seq_len, head_dim) otherwise.

 required

Returns:

Type Description
Tensor

The attention output tensor. Shape: (batch, seq_len, hidden_size).
reset_parameters()

Resets module parameters.

LowRankProjection

Bases: Module

Implements a low-rank linear projection with an intermediate normalization layer.

__init__(in_features, bottleneck, out_features, norm_eps)

Constructs the LowRankProjection object.

Parameters:

Name Type Description Default
in_features int

Input dimensionality.

 required
bottleneck int

Intermediate low-rank dimensionality.

 required
out_features int

Output dimensionality.

 required
norm_eps float

Epsilon value for the intermediate RMSNorm layer.

 required
forward(x)

Applies the low-rank projection to the inputs.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

 required

Returns:

Type Description
Tensor

Projected output tensor.
reset_parameters()

Resets module parameters.

MultiHeadLatentAttention

Bases: Module, ModuleLateInit

Implements Multi-Head Latent Attention (MLA) from DeepSeek-V2.

This module performs the full attention mechanism pipeline:

  1. Linear projection to Query (either direct or via a low-rank bottleneck with RMSNorm).
  2. Down-projection to a low-rank KV latent vector and a shared Key RoPE sub-vector.
  3. RMSNorm application on the KV latent vector.
  4. Up-projection of the KV latent vector into Key content (NOPE) and Value sub-vectors.
  5. Rotary Positional Embedding (RoPE) application strictly to the decoupled Query and Key RoPE sub-vectors.
  6. Concatenation of the content (NOPE) and rotated (RoPE) sub-vectors to form the final Query and Key heads.
  7. Scaled Dot Product Attention (via FlashAttention).
  8. Output projection.
q_lora_rank property

Rank of the Q low-rank path, or None if Q is projected directly.

__init__(hidden_size, num_attention_heads, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, kv_lora_rank, q_lora_rank, qk_down_norm_eps, is_causal, rope_style, sdpa_backend=None)

Constructs the MultiHeadLatentAttention layer.

Parameters:

Name Type Description Default
hidden_size int

Model hidden dimension.

 required
num_attention_heads int

Number of attention heads.

 required
qk_nope_head_dim int

Per-head dimension for the content (no-RoPE) part of Q and K.

 required
qk_rope_head_dim int

Per-head dimension for the RoPE-rotated part of Q and K.

 required
v_head_dim int

Per-head dimension for values.

 required
kv_lora_rank int

Rank of the KV latent compression.

 required
q_lora_rank int | None

Rank of the Q low-rank path. If None, Q is projected directly.

 required
qk_down_norm_eps float

Epsilon for the RMSNorm applied to the KV and Q latent representations.

 required
is_causal bool

Whether to apply a causal mask (auto-regressive).

 required
rope_style RotaryEmbeddingStyle

Rotary embedding layout style alignment.

 required
sdpa_backend AnySdpaBackendConfig | None

Configuration for the Scaled Dot-Product Attention backend. If None, the backend will be auto-detected via build_sdpa_backend().

 None

Raises:

Type Description
ValueError

If v_head_dim exceeds qk_head_dim.
forward(hidden_states, attention_mask, position_embeddings)

Computes Multi-Head Latent Attention.

Parameters:

Name Type Description Default
hidden_states Tensor

Input tensor. Shape: (batch, seq_len, hidden_size).

 required
attention_mask Tensor | None

Optional attention mask.

 required
position_embeddings tuple[Tensor, Tensor]

Tuple (cos, sin) for the RoPE sub-vectors. Each tensor shape: (batch, seq_len, qk_rope_head_dim).

 required

Returns:

Type Description
Tensor

Output tensor. Shape: (batch, seq_len, hidden_size).
reset_parameters()

Resets all learnable parameters.

d9d.module.block.attention.sdpa

EagerSdpaBackendConfig

Bases: BaseModel

Configuration for the eager (pure-PyTorch) backend.

The eager backend implements attention with explicit PyTorch ops.

Attributes:

Name Type Description
kind Literal['eager']

Discriminator field. Always "eager".

FlashAttention2SdpaBackendConfig

Bases: BaseModel

Configuration for the Flash Attention 2 backend.

Attributes:

Name Type Description
kind Literal['flash_attention_2']

Discriminator field. Always "flash_attention_2".

FlashAttention4SdpaBackendConfig

Bases: BaseModel

Configuration for the Flash Attention 4 backend.

Attributes:

Name Type Description
kind Literal['flash_attention_4']

Discriminator field. Always "flash_attention_4".

SdpaBackend

Bases: Protocol

Protocol for Scaled Dot-Product Attention backends.

This acts as a structural trait for attention backend modules. Any backend must provide a __call__ method with this signature since PyTorch modules are invoked directly.

__call__(query_states, key_states, value_states, attention_mask, is_causal, scale)

Computes Scaled Dot-Product Attention.

Parameters:

Name Type Description Default
query_states Tensor

Query tensor. Shape: (batch, seq_len, n_q_heads, head_dim).

 required
key_states Tensor

Key tensor. Shape: (batch, seq_len, n_kv_heads, head_dim).

 required
value_states Tensor

Value tensor. Shape: (batch, seq_len, n_kv_heads, head_dim).

 required
attention_mask Tensor | None

Mask tensor or None.

 required
is_causal bool

If True, applies an auto-regressive causal mask.

 required
scale float

Softmax scaling factor.

 required

Returns:

Type Description
Tensor

Attention output. Shape: (batch, seq_len, n_q_heads, head_dim).

SdpaParameters dataclass

Internal structural parameters passed to SDPA backend factories.

Attributes:

Name Type Description
num_sinks int | None

Number of learnable sink scalars (one per query head). None disables sinks and gives plain attention.
window_size tuple[int | None, int | None]

Sliding-window size for local attention as a tuple (left, right). (None, None) disables the window and uses full attention.
needs_attention_mask bool

Whether the layer will pass an explicit attention mask tensor to the backend at runtime. When True, auto-detection excludes backends that cannot accept explicit masks.

TorchSdpaBackendConfig

Bases: BaseModel

Configuration for the PyTorch SDPA backend.

Attributes:

Name Type Description
kind Literal['torch']

Discriminator field. Always "torch".
backends list[TorchSdpaBackendType] | None

A list of backends to enable during SDPA. If multiple are provided, PyTorch will try them in order or select the best one based on heuristics. If None, relies on PyTorch's default behavior.

TorchSdpaBackendType

Bases: StrEnum

Available SDPA backend implementations in PyTorch.

build_sdpa_backend(params, backend_config)

Builds the selected SDPA backend module based on the provided configuration.

If no explicit configuration is provided, it falls back to auto-detection (either from the D9D_BACKEND_AUTO_SDPA environment variable or programmatic defaults).

The factory resolves the appropriate module implementation, passing along the backend configuration and structural layer parameters.

Parameters:

Name Type Description Default
params SdpaParameters

Structural layer requirements needed by the backend.

 required
backend_config AnySdpaBackendConfig | None

Explicit SDPA backend configuration, or None to auto-detect.

 required

Returns:

Type Description
SdpaBackend

An instantiated SDPA module implementing the SdpaBackend protocol.

Raises:

Type Description
ValueError

If an unknown backend configuration type is encountered.

Linear Attention

Gated DeltaNet

GatedDeltaNet is an implementation of the Gated DeltaNet (GDN) attention mechanism.

It acts as a linear attention alternative that combines the Delta Rule with Mamba-style data-dependent gating and short causal convolutions.

d9d.module.block.attention.linear

GatedDeltaNet

Bases: Module, ModuleLateInit

Implements Gated DeltaNet (GDN) attention mechanism.

This module combines linear attention based on the Delta Rule with Mamba-style data-dependent gating and short causal convolutions.

Pipeline
  1. Linear projections for Q, K, V, output gate (G), decay gate (GK), and write strength (Beta).
  2. Causal short depthwise convolution applied to Q, K, V.
  3. Data-dependent decay computation (Mamba-style or log-sigmoid).
  4. GQA/MQA head expansion for Q and K.
  5. Chunked Gated Delta Rule (with optional internal L2 norm on Q/K).
  6. Per-head RMSNorm and SiLU-gated output projection.
__init__(hidden_size, num_query_key_heads, num_value_heads, head_qk_dim, head_v_dim, norm_eps, conv_size, decay_gate, use_qk_l2norm=True)

Constructs a GatedDeltaNet object.

Parameters:

Name Type Description Default
hidden_size int

Hidden size.

 required
num_query_key_heads int

Number of query and key attention heads before grouped expansion.

 required
num_value_heads int

Number of value attention heads.

 required
head_qk_dim int

Dimension allocated for a single query or key per head.

 required
head_v_dim int

Dimension allocated for a single value per head.

 required
norm_eps float

Small constant added for numerical stability to the normalization layer.

 required
conv_size int

Size of the causal convolution kernel context.

 required
decay_gate AnyDecayGateParameters

Structured parameters to initialize the selected decay gate mechanism.

 required
use_qk_l2norm bool

Whether to enable L2 normalization applied to Q/K internally.

 True

Raises:

Type Description
ValueError

When num_value_heads is not uniformly divisible by num_query_key_heads.
forward(hidden_states, attention_mask=None)

Runs forward pass.

Parameters:

Name Type Description Default
hidden_states Tensor

Input tensor sequence of shape (batch, seq_len, hidden_size).

 required
attention_mask Tensor | None

Optional padding mask tensor of shape (batch, seq_len).

 None

Returns:

Type Description
Tensor

Processed tensor possessing the identical shape as the input.
reset_parameters()

Resets learnable parameters of this module.

LogSigmoidDecayGateParameters

Bases: BaseModel

Configuration parameters for the LogSigmoid decay gate.

MambaDecayGateParameters

Bases: BaseModel

Configuration parameters for the Mamba-style decay gate.