Attention Layers

About

The d9d.module.block.attention package provides softmax and linear attention layers. Softmax attention layers compute scaled dot-product attention (SDPA) with a pluggable backend.

Softmax Attention

Grouped-Query Attention

GroupedQueryAttention implements Grouped-Query Attention. If num_key_value_heads equals num_attention_heads, it works as Multi-Head Attention. If num_key_value_heads is 1, it works as Multi-Query Attention.

Uses a pluggable SDPA backend.
Uses rotary position embedding, optionally on a part of each head (rope_dim).
Supports optional QK Normalization.
Supports optional sigmoid output gating.

Multi-Head Latent Attention

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

Uses a pluggable SDPA backend.
Uses rotary position embedding.
v_head_dim must not exceed qk_nope_head_dim + qk_rope_head_dim.

Scaled Dot-Product Attention Backends

Attention layers pass the scaled dot-product computation to a backend. The available backends are:

`kind`	Library	Sliding window	Learnable sinks	Explicit mask	Extra
`"flash_attention_4"`	FlashAttention 4	Yes	Yes	No	`d9d[backend-sdpa-flash-attention-4]`
`"flash_attention_2"`	FlashAttention 2	Yes	No	No	`d9d[backend-sdpa-flash-attention-2]`
`"torch"`	PyTorch `scaled_dot_product_attention`	No	No	Yes	-
`"eager"`	Plain PyTorch ops	Yes	Yes	Yes	-

The "torch" backend can restrict PyTorch to specific kernels (MATH, FLASH_ATTENTION, EFFICIENT_ATTENTION, CUDNN_ATTENTION). The "eager" backend has no extra dependencies. Use it as a fallback and as a correctness reference for the other backends.

Backend Selection

A layer picks its backend in this order:

The sdpa_backend configuration passed to the layer.
The D9D_BACKEND_AUTO_SDPA environment variable.
Auto-detection: the first installed backend that supports the layer, in the order of the table above.

D9D_BACKEND_AUTO_SDPA holds a backend configuration as JSON, keyed by the kind field:

# Force the eager backend.
export D9D_BACKEND_AUTO_SDPA='{"kind": "eager"}'

# Force the PyTorch backend and restrict it to the FlashAttention kernel.
export D9D_BACKEND_AUTO_SDPA='{"kind": "torch", "backends": ["FLASH_ATTENTION"]}'

Linear Attention

Gated DeltaNet

GatedDeltaNet implements Gated DeltaNet (GDN). It is a linear attention layer. It combines the delta rule with Mamba-style data-dependent gating and short causal convolutions.

It uses kernels from flash-linear-attention and requires the d9d[linear-attention] extra.

Usage

GroupedQueryAttention takes RoPE embeddings from a RotaryEmbeddingProvider. Modules use late initialization, so call reset_parameters() before use.

import torch

from d9d.module.block.attention import GroupedQueryAttention
from d9d.module.block.attention.sdpa import TorchSdpaBackendConfig
from d9d.module.block.positional import RotaryEmbeddingProvider, RotaryEmbeddingStyle

rope = RotaryEmbeddingProvider(
    rope_base=1_000_000,
    head_dim=128,
    max_position_ids=4096,
    style=RotaryEmbeddingStyle.HALF,
).cuda()
attention = GroupedQueryAttention(
    hidden_size=2048,
    num_attention_heads=32,
    num_key_value_heads=4,
    head_dim=128,
    qk_norm_eps=1e-6,
    is_causal=True,
    rope_style=RotaryEmbeddingStyle.HALF,
    sdpa_backend=TorchSdpaBackendConfig(),  # Omit to auto-detect
).cuda()
rope.reset_parameters()
attention.reset_parameters()

hidden_states = torch.randn(2, 16, 2048, device="cuda")
position_ids = torch.arange(16, device="cuda").expand(2, -1)
output = attention(hidden_states, attention_mask=None, position_embeddings=rope(position_ids))

API Reference

`d9d.module.block.attention`

Attention layers.

`GroupedQueryAttention`

Bases: Module, ModuleLateInit

Grouped Query Attention (GQA) layer with RoPE and optional QK normalization.

The layer runs these steps:

Linear projection to Q, K and V.
Optional RMSNorm on Q and K.
Rotary position embedding (RoPE).
Scaled dot-product attention with the configured SDPA backend.
Optional sigmoid output gating.
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 object.

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 the RMSNorm applied to Q and K. If `None`, normalization is disabled.	required
`is_causal`	`bool`	Whether to apply a causal mask.	required
`rope_style`	`RotaryEmbeddingStyle`	Rotary embedding layout style.	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`, uses 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 is auto-detected by `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 attention mask passed to the SDPA backend.	required
`position_embeddings`	`tuple[Tensor, Tensor]`	Tuple of `(cos, sin)` tensors for RoPE. Shape of each: `(batch, seq_len, rope_dim)` with partial RoPE, `(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

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 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. Shape: `(..., in_features)`.	required

Returns:

Type	Description
`Tensor`	Projected output tensor. Shape: `(..., out_features)`.

`reset_parameters()`

Resets module parameters.

`MultiHeadLatentAttention`

Bases: Module, ModuleLateInit

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

The layer runs these steps:

Projection to the query, either direct or through a low-rank bottleneck with RMSNorm.
Down-projection to a low-rank KV latent vector and a key RoPE sub-vector shared by all heads.
RMSNorm on the KV latent vector.
Up-projection of the KV latent vector into the key content (no-RoPE) and value sub-vectors.
Rotary position embedding (RoPE) on the query and key RoPE sub-vectors only.
Concatenation of the content and rotated sub-vectors into the final query and key heads.
Scaled dot-product attention with the configured SDPA backend.
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 object.

Parameters:

Name	Type	Description	Default
`hidden_size`	`int`	Hidden size.	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.	required
`rope_style`	`RotaryEmbeddingStyle`	Rotary embedding layout style.	required
`sdpa_backend`	`AnySdpaBackendConfig \| None`	Configuration for the scaled dot-product attention backend. If `None`, the backend is auto-detected by `build_sdpa_backend`.	`None`

Raises:

Type	Description
`ValueError`	If `v_head_dim` exceeds `qk_nope_head_dim + qk_rope_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 module parameters.

`d9d.module.block.attention.sdpa`

Scaled dot-product attention backends and their configurations.

`EagerSdpaBackendConfig`

Bases: BaseModel

Configuration for the eager 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 FlashAttention 2 backend.

Requires the d9d[backend-sdpa-flash-attention-2] extra.

Attributes:

Name	Type	Description
`kind`	`Literal['flash_attention_2']`	Discriminator field. Always `"flash_attention_2"`.

`FlashAttention4SdpaBackendConfig`

Bases: BaseModel

Configuration for the FlashAttention 4 backend.

Requires the d9d[backend-sdpa-flash-attention-4] extra.

Attributes:

Name	Type	Description
`kind`	`Literal['flash_attention_4']`	Discriminator field. Always `"flash_attention_4"`.

`SdpaBackend`

Bases: Protocol

Protocol for scaled dot-product attention (SDPA) backends.

Backends are nn.Module objects that are called directly, so the protocol declares __call__.

`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, num_heads, head_dim)`.	required
`key_states`	`Tensor`	Key tensor. Shape: `(batch, seq_len, num_kv_heads, head_dim)`.	required
`value_states`	`Tensor`	Value tensor. Shape: `(batch, seq_len, num_kv_heads, head_dim)`.	required
`attention_mask`	`Tensor \| None`	Optional attention mask, or `None`.	required
`is_causal`	`bool`	If `True`, applies a causal mask.	required
`scale`	`float`	Softmax scaling factor.	required

Returns:

Type	Description
`Tensor`	Attention output. Shape: `(batch, seq_len, num_heads, head_dim)`.

`SdpaParameters` `dataclass`

Structural parameters of an attention layer that an SDPA backend must support.

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 passes an explicit attention mask 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`	PyTorch SDPA kernels to enable. PyTorch picks one of the enabled kernels that supports the inputs. If `None`, PyTorch uses its default kernel selection.

`TorchSdpaBackendType`

Bases: StrEnum

SDPA kernels available in PyTorch.

Each member maps to the torch.nn.attention.SDPBackend member of the same name.

Attributes:

Name	Type	Description
`MATH`		Reference implementation in plain PyTorch ops.
`FLASH_ATTENTION`		FlashAttention kernel.
`EFFICIENT_ATTENTION`		Memory-efficient attention kernel.
`CUDNN_ATTENTION`		cuDNN attention kernel.

`build_sdpa_backend(params, backend_config)`

Builds the SDPA backend module for the given configuration.

If backend_config is None, the backend config is read from the D9D_BACKEND_AUTO_SDPA environment variable (a JSON backend config). If the variable is not set, the factory picks the first installed backend that supports params, in this order: FlashAttention 4, FlashAttention 2, PyTorch SDPA, eager.

Parameters:

Name	Type	Description	Default
`params`	`SdpaParameters`	Structural layer parameters the backend must support.	required
`backend_config`	`AnySdpaBackendConfig \| None`	Explicit SDPA backend configuration, or `None` to auto-detect.	required

Returns:

Type	Description
`SdpaBackend`	The SDPA backend module.

Raises:

Type	Description
`ValueError`	If the backend configuration type is unknown.

`d9d.module.block.attention.linear`

Linear attention layers.

`GatedDeltaNet`

Bases: Module, ModuleLateInit

Gated DeltaNet (GDN) attention layer.

The layer combines linear attention based on the delta rule with Mamba-style data-dependent gating and short causal convolutions. It runs these steps:

Linear projections for Q, K, V, the output gate, the decay gate and the write strength (beta).
Causal short depthwise convolution on Q, K and V.
Data-dependent decay (Mamba-style or log-sigmoid).
GQA/MQA head expansion for Q and K.
Chunked gated delta rule, with optional L2 normalization of Q and K.
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 the GatedDeltaNet object.

Parameters:

Name	Type	Description	Default
`hidden_size`	`int`	Hidden size.	required
`num_query_key_heads`	`int`	Number of query and key heads before grouped expansion.	required
`num_value_heads`	`int`	Number of value heads.	required
`head_qk_dim`	`int`	Dimension of a single query or key head.	required
`head_v_dim`	`int`	Dimension of a single value head.	required
`norm_eps`	`float`	Epsilon for the output RMSNorm.	required
`conv_size`	`int`	Kernel size of the short causal convolution.	required
`decay_gate`	`AnyDecayGateParameters`	Decay gate configuration.	required
`use_qk_l2norm`	`bool`	Whether to L2-normalize Q and K inside the kernel.	`True`

Raises:

Type	Description
`ValueError`	If `num_value_heads` is not divisible by `num_query_key_heads`.

`forward(hidden_states, attention_mask=None)`

Computes Gated DeltaNet attention.

Parameters:

Name	Type	Description	Default
`hidden_states`	`Tensor`	Input tensor. Shape: `(batch, seq_len, hidden_size)`.	required
`attention_mask`	`Tensor \| None`	Optional padding mask. Shape: `(batch, seq_len)`.	`None`

Returns:

Type	Description
`Tensor`	Output tensor. Shape: `(batch, seq_len, hidden_size)`.

`reset_parameters()`

Resets module parameters.

`LogSigmoidDecayGateParameters`

Bases: BaseModel

Configuration for the log-sigmoid decay gate.

Attributes:

Name	Type	Description
`type`	`Literal['logsigmoid']`	Discriminator field. Always `"logsigmoid"`.
`normalizer`	`float`	Temperature that divides the log-sigmoid output.

`MambaDecayGateParameters`

Bases: BaseModel

Configuration for the Mamba-style decay gate.

Attributes:

Name	Type	Description
`type`	`Literal['mamba']`	Discriminator field. Always `"mamba"`.
`normalizer`	`float`	Upper bound of the uniform initialization of `A`.
`dt_min`	`float`	Minimum `dt` for initialization.
`dt_max`	`float`	Maximum `dt` for initialization.
`dt_init_floor`	`float`	Floor for `dt` clamping during initialization.

Attention Layers

About

Softmax Attention

Grouped-Query Attention

Multi-Head Latent Attention

Scaled Dot-Product Attention Backends

Backend Selection

Linear Attention

Gated DeltaNet

Usage

API Reference

d9d.module.block.attention

GroupedQueryAttention

__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)

forward(hidden_states, attention_mask, position_embeddings)

reset_parameters()

LowRankProjection

__init__(in_features, bottleneck, out_features, norm_eps)

forward(x)

reset_parameters()

MultiHeadLatentAttention

q_lora_rank property

__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)

forward(hidden_states, attention_mask, position_embeddings)

reset_parameters()

d9d.module.block.attention.sdpa

EagerSdpaBackendConfig

FlashAttention2SdpaBackendConfig

FlashAttention4SdpaBackendConfig

SdpaBackend

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

SdpaParameters dataclass

TorchSdpaBackendConfig

TorchSdpaBackendType

build_sdpa_backend(params, backend_config)

d9d.module.block.attention.linear

GatedDeltaNet

__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)

forward(hidden_states, attention_mask=None)

reset_parameters()

LogSigmoidDecayGateParameters

MambaDecayGateParameters

`d9d.module.block.attention`

`GroupedQueryAttention`

`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)`

`forward(hidden_states, attention_mask, position_embeddings)`

`reset_parameters()`

`LowRankProjection`

`init(in_features, bottleneck, out_features, norm_eps)`

`forward(x)`

`reset_parameters()`

`MultiHeadLatentAttention`

`q_lora_rank` `property`

`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)`

`forward(hidden_states, attention_mask, position_embeddings)`

`reset_parameters()`

`d9d.module.block.attention.sdpa`

`EagerSdpaBackendConfig`

`FlashAttention2SdpaBackendConfig`

`FlashAttention4SdpaBackendConfig`

`SdpaBackend`

`call(query_states, key_states, value_states, attention_mask, is_causal, scale)`

`SdpaParameters` `dataclass`

`TorchSdpaBackendConfig`

`TorchSdpaBackendType`

`build_sdpa_backend(params, backend_config)`

`d9d.module.block.attention.linear`

`GatedDeltaNet`

`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)`

`forward(hidden_states, attention_mask=None)`

`reset_parameters()`

`LogSigmoidDecayGateParameters`

`MambaDecayGateParameters`