class BeamAttention(nn.Module):
"""GQA with per-head Q/K RMSNorm, then RoPE (default) or SSMax (NoPE).
With ``attention_output_gate``, the attention output is scaled by
``gate_floor + (1 - gate_floor) * sigmoid((gate + bias) / temperature)``
before ``o_proj``.
"""
def __init__(
self,
config: BeamConfig,
cache_config: CacheConfig | None,
quant_config: QuantizationConfig | None,
prefix: str = "",
reduce_results: bool = True,
) -> None:
super().__init__()
self.layer_id = extract_layer_index(prefix)
hidden_size = config.hidden_size
tp_size = get_tensor_model_parallel_world_size()
self.total_num_heads = config.num_attention_heads
assert self.total_num_heads % tp_size == 0
self.num_heads = self.total_num_heads // tp_size
self.total_num_kv_heads = config.num_key_value_heads
if self.total_num_kv_heads >= tp_size:
assert self.total_num_kv_heads % tp_size == 0
else:
assert tp_size % self.total_num_kv_heads == 0
self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
self.head_dim = config.head_dim or (hidden_size // self.total_num_heads)
self.q_size = self.num_heads * self.head_dim
self.kv_size = self.num_kv_heads * self.head_dim
self.scaling = self.head_dim**-0.5
self.merged_qkv_gate = use_merged_qkv_gate(config, tp_size)
if self.merged_qkv_gate:
self.qkv_gate_proj = MergedColumnParallelLinear(
hidden_size,
[
self.total_num_heads * self.head_dim,
self.total_num_kv_heads * self.head_dim,
self.total_num_kv_heads * self.head_dim,
self.total_num_heads * self.head_dim,
],
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.qkv_gate_proj",
)
else:
self.qkv_proj = QKVParallelLinear(
hidden_size,
self.head_dim,
self.total_num_heads,
self.total_num_kv_heads,
bias=config.attention_bias,
quant_config=quant_config,
prefix=f"{prefix}.qkv_proj",
)
self.o_proj = RowParallelLinear(
self.total_num_heads * self.head_dim,
hidden_size,
bias=config.attention_bias,
quant_config=quant_config,
prefix=f"{prefix}.o_proj",
reduce_results=reduce_results,
)
# Affine and parameter-free per-head Q/K norms share `rms_norm_eps`.
qk_norm_has_weight = not config.use_paramfree_qk_norm
self.rms_norm_eps = config.rms_norm_eps
self.q_norm = RMSNorm(self.head_dim, has_weight=qk_norm_has_weight)
self.k_norm = RMSNorm(self.head_dim, has_weight=qk_norm_has_weight)
self.disable_rope = self.layer_id in config.nope_layer_ids
self.use_ssmax = bool(config.ssmax) and self.disable_rope
self.ssmax_mode = config.ssmax_scale_mode if self.use_ssmax else "none"
if self.ssmax_mode not in ("none", "default", "root10", "root1000"):
raise ValueError(f"Unknown ssmax_scale_mode: {self.ssmax_mode!r}")
self.ssmax_offset = config.ssmax_offset
self.ssmax_log_divisor = (
math.log(config.ssmax_offset)
if self.use_ssmax
and self.ssmax_mode == "default"
and config.ssmax_divide_by_log_offset
else 1.0
)
# The output gate has a separate bias per head.
self.attn_output_gate = bool(config.attention_output_gate)
if self.attn_output_gate:
self.gate_floor = float(config.attention_output_gate_floor)
self.gate_temperature = float(config.attention_output_gate_temperature)
if not 0 <= self.gate_floor < 1:
raise ValueError(
"attention_output_gate_floor must be in [0, 1); "
f"got {self.gate_floor}"
)
if self.gate_temperature <= 0:
raise ValueError(
"attention_output_gate_temperature must be positive; "
f"got {self.gate_temperature}"
)
if not self.merged_qkv_gate:
self.output_gate_proj = ColumnParallelLinear(
hidden_size,
self.total_num_heads * self.head_dim,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.output_gate_proj",
)
if config.attention_output_gate_bias:
self.gate_bias = nn.Parameter(torch.zeros(self.num_heads))
set_weight_attrs(
self.gate_bias, {"weight_loader": sharded_weight_loader(0)}
)
else:
self.register_parameter("gate_bias", None)
if self.disable_rope:
self.rotary_emb = None
else:
self.rotary_emb = get_rope(
self.head_dim,
max_position=config.max_position_embeddings,
is_neox_style=True,
rope_parameters=config.rope_parameters,
dtype=torch.float32,
)
# Attention masks follow `layer_types` independently of the NoPE flags.
# Backends interpret `sliding_window` as `(q - k) <= sliding_window - 1`.
is_sliding = config.layer_types[self.layer_id] == "sliding_attention"
per_layer_sw = config.sliding_window if is_sliding else None
self.attn = Attention(
self.num_heads,
self.head_dim,
self.scaling,
num_kv_heads=self.num_kv_heads,
cache_config=cache_config,
quant_config=quant_config,
per_layer_sliding_window=per_layer_sw,
attn_type=AttentionType.DECODER,
prefix=f"{prefix}.attn",
)
self.fuse_kv_insert = (
self.attn.attn_backend.get_name() in ("FLASH_ATTN", "FLASHINFER")
and self.attn.impl.dcp_world_size == 1
and self.attn.kv_sharing_target_layer_name is None
)
self.low_latency_projection = False
def warmup_projection(self) -> None:
if not self.merged_qkv_gate:
return
projection = self.qkv_gate_proj
self.low_latency_projection = isinstance(
projection.quant_method, UnquantizedLinearMethod
) and use_low_latency_gemm(projection.weight)
if self.low_latency_projection:
warmup_packed_projection(projection.weight)
def forward(
self, positions: torch.Tensor, hidden_states: torch.Tensor
) -> torch.Tensor:
gate = None
if self.merged_qkv_gate:
if self.low_latency_projection:
packed = packed_projection(
hidden_states, self.qkv_gate_proj.weight, True
)
else:
packed, _ = self.qkv_gate_proj(hidden_states)
qkv, gate = packed.split(
[self.q_size + 2 * self.kv_size, self.q_size], dim=-1
)
else:
qkv, _ = self.qkv_proj(hidden_states)
if self.attn_output_gate:
gate, _ = self.output_gate_proj(hidden_states)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
kv_cache, slot_mapping = None, None
if (
self.fuse_kv_insert
and self.attn.kv_cache.dtype == q.dtype == torch.bfloat16
):
slots = get_forward_context().slot_mapping
if isinstance(slots, dict):
slot_mapping = slots.get(self.attn.layer_name)
if slot_mapping is not None:
kv_cache = self.attn.kv_cache
rope = self.rotary_emb
q, k = beam_qk_norm_rope(
q,
k,
self.q_norm.weight,
self.k_norm.weight,
self.rms_norm_eps,
positions,
rope.cos_sin_cache if rope is not None else None,
self.head_dim,
rope.rotary_dim if rope is not None else 0,
rope.is_neox_style if rope is not None else True,
self.ssmax_mode,
self.ssmax_offset,
self.ssmax_log_divisor,
value=v,
kv_cache=kv_cache,
slot_mapping=slot_mapping,
)
# Both backends read the populated cache when K/V are omitted.
attn_output = (
self.attn(q, k, v) if kv_cache is None else self.attn(q, None, None)
)
if gate is not None:
attn_output = beam_output_gate(
attn_output,
gate,
self.gate_bias,
num_heads=self.num_heads,
head_dim=self.head_dim,
floor=self.gate_floor,
temperature=self.gate_temperature,
)
output, _ = self.o_proj(attn_output)
return output