For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

AttentionWithRope

AttentionWithRope

class max.experimental.nn.common_layers.attention.AttentionWithRope(*, rope, num_attention_heads, num_key_value_heads, hidden_size, kv_params, layer_idx, scale=None, has_bias=False, stacked_qkv=False, clip_qkv=None, use_qk_norm=False, rms_norm_eps=1e-06)

source

Bases: Module

Implementation of attention that uses Rotary Position Embedding (RoPE).

This is a ModuleV3 port of the legacy AttentionWithRope class. The Q/K/V projection is a fused QKVLinear qkv_proj submodule that handles both checkpoint layouts: separate q/k/v weights (the default, name-transparent) and a pre-fused qkv_proj weight (stacked_qkv=True). It also supports optional clip_qkv clamping and optional QK normalization via RMSNorm.

Initializes the attention layer.

Parameters:

• rope (RotaryEmbedding) – The rope layer to borrow the freqs_cis value from.
• num_attention_heads (int) – The number of attention heads.
• num_key_value_heads (int) – Number of key/value heads.
• hidden_size (int) – The dimension of the hidden states.
• kv_params (KVCacheParams) – KV Cache params, including number of kv heads, head dim, and dtype.
• layer_idx (int) – The layer number associated with this Attention block.
• scale (float | None) – Optional attention scale; defaults to sqrt(1/head_dim).
• has_bias (bool) – Whether Q/K/V have bias (stacked_qkv forbids bias).
• stacked_qkv (bool) – Whether the checkpoint stores Q/K/V as a single fused qkv_proj weight (vs separate q/k/v weights).
• clip_qkv (float | None) – If provided, clamp Q/K/V weights to [-clip_qkv, clip_qkv].
• use_qk_norm (bool) – Whether to use RMSNorm on Q/K.
• rms_norm_eps (float) – Value to use for numerical stability in RMSNorm.

forward()

forward(x, kv_collection, **kwargs)

source

Computes attention over the input tensor using the KV cache.

Parameters:

• x (Tensor)
• kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue])

Return type:

Tensor

wqkv

property wqkv: Tensor

source

The fused q||k||v projection weight (clamped if clip_qkv).

wqkv_bias

property wqkv_bias: Tensor | None

source

The fused q||k||v projection bias, or None when unused.