"""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``.
"""
from __future__ import annotations
import torch
from torch import nn
from mstar.distributed.communication import CommGroup
from mstar.engine.resources.convenience import AttentionCallable
from mstar.model.components.distributed.linear import (
QKVParallelLinear,
RowParallelLinear,
)
from mstar.model.components.norm import RMSNorm
[docs]
class ParallelAttention(nn.Module):
def __init__(
self,
*,
comm_group: CommGroup | None = None,
hidden_size: int,
num_heads: int,
num_kv_heads: int,
head_dim: int,
qkv_bias: bool = False,
o_bias: bool = False,
qk_norm: bool = False,
rms_norm_eps: float = 1e-6,
rope_theta: float = 10_000.0,
rope_scale: float = 1.0,
rope_low_freq_factor: float = 1.0,
rope_high_freq_factor: float = 1.0,
rope_old_context_len: int = 8192,
input_hidden_size: int | None = None,
attn_key: str = "attn",
kv_key: str = "kv",
pos_key: str | None = "rope",
):
super().__init__()
# resource labels this layer calls; see components/attention.py
self._attn_key = attn_key
self._kv_key = kv_key
self._pos_key = pos_key
self.attn = None
self.kv = None
self.pos = None
if comm_group is None:
comm_group = CommGroup.trivial()
self.comm_group = comm_group
self.hidden_size = hidden_size
self.input_hidden_size = input_hidden_size or hidden_size
self.head_dim = head_dim
self.total_num_heads = num_heads
self.total_num_kv_heads = num_kv_heads
self.qkv_proj = QKVParallelLinear(
comm_group=comm_group,
hidden_size=self.input_hidden_size,
head_size=head_dim,
total_num_heads=num_heads,
total_num_kv_heads=num_kv_heads,
bias=qkv_bias,
)
self.num_heads = self.qkv_proj.num_heads
self.num_kv_heads = self.qkv_proj.num_kv_heads
self.o_proj = RowParallelLinear(
comm_group=comm_group,
input_size=num_heads * head_dim,
output_size=self.input_hidden_size,
bias=o_bias,
input_is_parallel=True,
reduce_results=True,
)
self.rope_theta = rope_theta
self.rope_scale = rope_scale
self.rope_low_freq_factor = rope_low_freq_factor
self.rope_high_freq_factor = rope_high_freq_factor
self.rope_old_context_len = rope_old_context_len
if qk_norm:
self.q_norm = RMSNorm(head_dim, eps=rms_norm_eps)
self.k_norm = RMSNorm(head_dim, eps=rms_norm_eps)
else:
self.q_norm = None
self.k_norm = None
def _project_qkv(
self, hidden_states: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
num_tokens = hidden_states.shape[0]
qkv = self.qkv_proj(hidden_states)
q_size = self.num_heads * self.head_dim
k_size = self.num_kv_heads * self.head_dim
v_size = self.num_kv_heads * self.head_dim
q, k, v = qkv.split([q_size, k_size, v_size], dim=-1)
q = q.view(num_tokens, self.num_heads, self.head_dim)
k = k.view(num_tokens, self.num_kv_heads, self.head_dim)
v = v.view(num_tokens, self.num_kv_heads, self.head_dim)
return q, k, v
def _apply_qk_norm(
self, q: torch.Tensor, k: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
if self.q_norm is None:
return q, k
q_shape, k_shape = q.shape, k.shape
q = self.q_norm(q.reshape(-1, self.head_dim)).view(q_shape)
k = self.k_norm(k.reshape(-1, self.head_dim)).view(k_shape)
return q, k
[docs]
def bind_resources(self, resources: dict) -> None:
"""Resolve the resources this layer calls. See
``NodeSubmodule.bind_node_resources``."""
self.attn = resources.get(self._attn_key)
self.kv = resources.get(self._kv_key)
self.pos = None if self._pos_key is None else resources.get(self._pos_key)
# see Attention.bind_resources
self.attend = AttentionCallable(kv=self.kv, attn=self.attn)
def _apply_rope(
self,
q: torch.Tensor, k: torch.Tensor, label: str,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Standard 1D RoPE through the position resource. Override for
non-standard schemes (3D MRoPE, etc.).
The llama31 kwargs go through unconditionally, as they did on the old
cache-handle path — passing them selects flashinfer's llama31 kernel."""
if self.pos is None:
return q, k
return self.pos.apply_qk(
q, k, label=label,
rope_theta=self.rope_theta,
rope_scale=self.rope_scale,
low_freq_factor=self.rope_low_freq_factor,
high_freq_factor=self.rope_high_freq_factor,
old_context_len=self.rope_old_context_len,
)
[docs]
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
"""Label and layer index come off the resources' cursors; see
``Attention.forward``."""
num_tokens = hidden_states.shape[0]
q, k, v = self._project_qkv(hidden_states)
q, k = self._apply_qk_norm(q, k)
q, k = self._apply_rope(q, k, self.attend.label)
attn_output = self.attend(q, k, v)
attn_output = attn_output.reshape(num_tokens, self.num_heads * self.head_dim)
return self.o_proj(attn_output)