mstar.model.components.distributed.attention

mstar.model.components.distributed.attention#

TP-aware multi-head attention.

Mirrors mstar.model.components.Attention but with the QKV projection sharded across heads via QKVParallelLinear and the output projection all-reduced via RowParallelLinear. QK-norm (if enabled) and RoPE each operate on this rank’s local slice of heads — no cross-rank communication beyond the AllReduce hidden inside o_proj.

Worker integration:

  • The per-rank num_heads / num_kv_heads come from self.qkv_proj (already computed by QKVParallelLinear based on the comm group’s world size and GQA replica count).

  • The KV resource’s KVConfig must carry the per-rank head counts so paged attention reads / writes the right slice. KVConfig.shard() narrows them to one rank’s slice at build time, from the instance world size (tp * sp).

For non-standard RoPE (qwen3’s 3D MRoPE), subclass and override _apply_rope — same shape as the non-parallel Attention.

Classes

ParallelAttention(*[, comm_group, qkv_bias, ...])

class mstar.model.components.distributed.attention.ParallelAttention(*, comm_group=None, 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:

  • comm_group (CommGroup | None)

  • 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.

Parameters:

resources (dict)

Return type:

None

forward(hidden_states)[source]#

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

Parameters:

hidden_states (Tensor)

Return type:

Tensor