mstar.model.components.attention

mstar.model.components.attention#

Multi-head attention with GQA, optional QK-norm, and pluggable RoPE.

Constructed with separate q_proj / k_proj / v_proj Linears matching the HF checkpoint layout. After loading, call consolidate_qkv_weight() to fuse them into a single qkv_proj_weight buffer (one fused GEMM instead of three) and null out the originals; the forward branches on whether consolidation has happened.

Variations supported:

  • GQA via num_kv_heads < num_heads.

  • Optional bias on qkv / o projections.

  • Optional per-head QK-norm (RMSNorm applied to q / k after projection, before RoPE) — used by qwen3.

  • Different input_hidden_size from the model’s nominal hidden_size (used by pi05’s action expert, which shares K/V dims with PaliGemma but has its own width).

  • Llama-style RoPE scaling parameters (rope_scale, low_freq_factor, etc.) for the cache-handle path.

For non-standard RoPE schemes (e.g. qwen3’s 3D MRoPE), subclass and override _apply_rope rather than going through the position resource.

Classes

Attention(*, hidden_size, num_heads, ...[, ...])

CrossAttention(*, hidden_size, num_heads, ...)

Multi-head cross-attention over an engine-managed encoder-context KV.
class mstar.model.components.attention.Attention(*, hidden_size, num_heads, num_kv_heads, head_dim, qkv_bias=False, o_bias=False, qk_norm=False, rms_norm_eps=1e-6, rope_theta=10_000.0, rope_scale=1.0, rope_low_freq_factor=1.0, rope_high_freq_factor=1.0, rope_old_context_len=8192, input_hidden_size=None, attn_key='attn', kv_key='kv', pos_key='rope')[source]#

Bases: Module

Parameters:

  • hidden_size (int)

  • num_heads (int)

  • num_kv_heads (int)

  • head_dim (int)

  • qkv_bias (bool)

  • o_bias (bool)

  • qk_norm (bool)

  • rms_norm_eps (float)

  • rope_theta (float)

  • rope_scale (float)

  • rope_low_freq_factor (float)

  • rope_high_freq_factor (float)

  • rope_old_context_len (int)

  • input_hidden_size (int | None)

  • attn_key (str)

  • kv_key (str)

  • pos_key (str | None)

bind_resources(resources)[source]#

Resolve the resources this layer calls. See NodeSubmodule.bind_node_resources.

.get: a layer may be bound on a node that owns only some of them.

Parameters:

resources (dict)

Return type:

None

consolidate_qkv_weight()[source]#

Fuse q_proj/k_proj/v_proj weights into a single qkv_proj_weight buffer and null out the originals. Idempotent.

Return type:

None

forward(hidden_states)[source]#

The label and layer index are cursors on the resources, set by the caller running the layer stack (attend.bind_step once, then attend.set_layer_idx per layer) rather than passed in per call.

Parameters:

hidden_states (Tensor)

Return type:

Tensor

class mstar.model.components.attention.CrossAttention(*, hidden_size, num_heads, head_dim, q_bias=True, k_bias=False, v_bias=True, o_bias=True, source='default', cross_key=None, context_kv_key=None)[source]#

Bases: Module

Multi-head cross-attention over an engine-managed encoder-context KV.

For encoder-decoder models (Whisper, etc.): the decoder attends to a fixed encoder context whose K/V are computed once at prefill and written into a KV resource of their own, under the context_label the model’s CrossAttentionConfig names (see issue #160). Q is projected per step; K/V projections are exposed via compute_kv so the submodule can write them into that cache at encode time.

Q/K/V/O are separate nn.Linear matching the HF layout. Subclasses override projection details (bias flags, a compute_kv that reshapes for a model-specific pool layout) as needed; the default matches Whisper (q/v/o biased, k unbiased).

TODO(#160): the projections are plain nn.Linear — this module is not yet TP/SP-compatible (no column/row-parallel splits over heads). A tensor-parallel cross-attention variant is needed to serve the decoder under TP alongside the self-attention path.

Parameters:

  • hidden_size (int)

  • num_heads (int)

  • head_dim (int)

  • q_bias (bool)

  • k_bias (bool)

  • v_bias (bool)

  • o_bias (bool)

  • source (str)

  • cross_key (str | None)

  • context_kv_key (str | None)

bind_resources(resources)[source]#

Resolve this source’s cross-attention resource and the cache holding its context. See NodeSubmodule.bind_node_resources.

Parameters:

resources (dict)

Return type:

None

bind_step(label)[source]#
Parameters:

label (str)

Return type:

None

compute_kv(encoder_states)[source]#

Project the encoder context to K/V for the cross-attention pool.

(enc_len, hidden) -> (k, v), each (enc_len, num_heads, head_dim). Override to reshape for a model-specific pool layout.

Parameters:

encoder_states (Tensor)

Return type:

tuple[Tensor, Tensor]

forward(hidden_states)[source]#

Label and layer index come off the resources’ cursors; see Attention.forward.

Parameters:

hidden_states (Tensor)

Return type:

Tensor

set_layer_idx(layer_idx)[source]#
Parameters:

layer_idx (int)

Return type:

None