modeling_parallax.py

"""Parallax HF modeling file, modular-style (inherits from Qwen3).

This module is the second, separate HF interface for the Parallax family.
It inherits almost everything from ``transformers.models.qwen3`` and only
overrides the attention class to add the centered amortized local linear
attention op (``attn_type="parallax"``).

Design notes
------------
* All boilerplate — RMSNorm, MLP, rotary embeddings, decoder-layer scaffolding,
  the base ``PreTrainedModel`` class, ``Qwen3Model``, ``Qwen3ForCausalLM``,
  mask preparation, cache handling, ``GenerationMixin`` integration, etc. —
  is inherited from ``transformers.models.qwen3``.
* The only new code is:
    - ``parallax_centered_attention_forward``: a pure-PyTorch attention kernel
      matching the math of ``torchtitan/models/flashlla/dev_ops/reference.py``
      (``centered_dense_reference``), but batched over heads and taking an
      additive attention mask in the HF shape ``(B, 1, q_len, kv_len)``.
    - ``ParallaxAttention``: inherits from ``Qwen3Attention``, adds an
      ``r_proj`` / ``r_norm`` pair when ``config.attn_type == "parallax"``,
      and overrides ``forward`` to dispatch to the parallax kernel in that
      case. For ``attn_type == "sdpa"`` the forward simply delegates to
      ``Qwen3Attention.forward`` unchanged, so all modern HF attention
      backends (sdpa, flash_attention_2, flex_attention, eager) remain
      available for that path.
    - ``ParallaxDecoderLayer``, ``ParallaxModel``, ``ParallaxPreTrainedModel``,
      ``ParallaxForCausalLM``: thin subclasses that rewire parent classes to
      the ``ParallaxAttention`` / ``ParallaxConfig`` below.

Why this file exists alongside ``torchtitan/hf_modeling/parallax/``
-------------------------------------------------------------------
The older ``torchtitan.hf_modeling.parallax`` package is fully self-contained
(no imports from ``transformers.models.qwen3``). That makes it robust across
transformers versions but also duplicates a lot of boilerplate that drifts
from HF conventions. This ``parallax_qwen3`` package takes the opposite
tradeoff: minimal code, surgical diffs from Qwen3, tracks HF updates for
free, at the cost of depending on a stable ``transformers.models.qwen3`` API.

Both produce numerically equivalent outputs on the same Parallax checkpoint
weights; a parity test verifies this.
"""

from __future__ import annotations

from typing import Optional

import torch
import torch.nn as nn
from transformers.cache_utils import Cache
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.models.qwen3.modeling_qwen3 import (
    Qwen3Attention,
    Qwen3DecoderLayer,
    Qwen3ForCausalLM,
    Qwen3Model,
    Qwen3PreTrainedModel,
    Qwen3RMSNorm,
    apply_rotary_pos_emb,
    repeat_kv,
    rotate_half,
)

from .configuration_parallax import ParallaxConfig


# ---------------------------------------------------------------------------
# Parallax centered attention kernel (pure PyTorch reference)
# ---------------------------------------------------------------------------


def parallax_centered_attention_forward(
    module: nn.Module,
    query: torch.Tensor,
    r: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: Optional[torch.Tensor],
    scaling: float,
    dropout: float = 0.0,
    **kwargs,
) -> tuple[torch.Tensor, None]:
    """Centered amortized local linear attention.

    This is a batched generalization of ``centered_dense_reference`` from
    ``torchtitan/models/flashlla/dev_ops/reference.py``:

        p = softmax((Q Kᵀ) * scale + mask)
        rk = R Kᵀ
        s = p V
        c = sum(p * rk, -1, keepdim=True)
        r_term = (p * rk) V
        out = s + c * s - r_term

    Shapes
    ------
    query : (B, H, q_len, D)
    r     : (B, H, q_len, D)          — Parallax-specific extra projection
    key   : (B, H_kv, kv_len, D)      — GQA-compressed, expanded internally
    value : (B, H_kv, kv_len, D)
    attention_mask : (B, 1, q_len, kv_len) additive, 0 or -inf (HF convention)

    Returns
    -------
    attn_output : (B, q_len, H, D)    — transposed to match Qwen3 eager impl
    attn_weights : None               — not exposed (centered attention has no
                                         meaningful "attention matrix" per head)
    """
    # Expand GQA-compressed k/v to match q heads
    key = repeat_kv(key, module.num_key_value_groups)
    value = repeat_kv(value, module.num_key_value_groups)

    # Upcast to float32 for numerical stability (matches the training kernel,
    # which uses bfloat16 for math but accumulates in fp32 within the softmax).
    out_dtype = query.dtype
    qf = query.float()
    rf = r.float()
    kf = key.float()
    vf = value.float()

    qk = torch.einsum("bhqd,bhkd->bhqk", qf, kf) * scaling
    q_len = qk.shape[-2]
    kv_len = qk.shape[-1]
    if attention_mask is not None:
        # HF causal masks may cover a larger key span than our current
        # kv_len during decode; slice to match.
        qk = qk + attention_mask[..., :q_len, :kv_len].to(qk.dtype)
    elif q_len > 1:
        # HF's `create_causal_mask` returns None when the model's default
        # attention backend (e.g. sdpa) handles causality implicitly. Our
        # parallax kernel is hand-written and needs the explicit additive
        # mask, so build one here. Decode (q_len == 1) is naturally
        # unmasked: the single query can attend to every cached key.
        offset = kv_len - q_len
        row = torch.arange(q_len, device=qk.device)[:, None] + offset
        col = torch.arange(kv_len, device=qk.device)[None, :]
        causal = row >= col
        bias = torch.zeros(q_len, kv_len, dtype=qk.dtype, device=qk.device)
        bias = bias.masked_fill(~causal, float("-inf"))
        qk = qk + bias

    p = torch.softmax(qk, dim=-1)
    rk = torch.einsum("bhqd,bhkd->bhqk", rf, kf)

    s = torch.einsum("bhqk,bhkd->bhqd", p, vf)
    c = (p * rk).sum(dim=-1, keepdim=True)
    r_term = torch.einsum("bhqk,bhkd->bhqd", p * rk, vf)

    attn_output = (s + c * s - r_term).to(out_dtype)
    # Match Qwen3 eager_attention_forward: return transposed (B, q_len, H, D)
    attn_output = attn_output.transpose(1, 2).contiguous()
    return attn_output, None


# ---------------------------------------------------------------------------
# Attention / DecoderLayer / Model / ForCausalLM
# ---------------------------------------------------------------------------


class ParallaxAttention(Qwen3Attention):
    """Qwen3 attention plus an optional Parallax (centered amortized LLA) path."""

    def __init__(self, config: ParallaxConfig, layer_idx: int) -> None:
        super().__init__(config, layer_idx)
        self.attn_type = config.attn_type
        self.amortized_rho_scale = config.amortized_rho_scale
        self.learn_amortized_rho_scale = config.learn_amortized_rho_scale
        self.rope_r = getattr(config, "rope_r", "correct")
        self._buggy_rope_cos: torch.Tensor | None = None
        self._buggy_rope_sin: torch.Tensor | None = None
        if self.attn_type == "parallax":
            # Extra projection + RMSNorm for the amortized local linear attn
            self.r_proj = nn.Linear(
                config.hidden_size,
                config.num_attention_heads * self.head_dim,
                bias=config.attention_bias,
            )
            self.r_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
            # Learned per-head gate on r (sigmoid), mirroring torchtitan `wa`.
            if self.learn_amortized_rho_scale:
                self.a_proj = nn.Linear(
                    config.hidden_size,
                    config.num_attention_heads,
                    bias=False,
                )
            else:
                self.a_proj = None

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor],
        past_key_values: Optional[Cache] = None,
        **kwargs,
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        if self.attn_type != "parallax":
            # Plain Qwen3 attention — delegate to the parent, which supports
            # sdpa / flash_attention_2 / flex_attention / eager backends.
            return super().forward(
                hidden_states=hidden_states,
                position_embeddings=position_embeddings,
                attention_mask=attention_mask,
                past_key_values=past_key_values,
                **kwargs,
            )

        # ---- Parallax path ----
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        query_states = self.q_norm(
            self.q_proj(hidden_states).view(hidden_shape)
        ).transpose(1, 2)
        key_states = self.k_norm(
            self.k_proj(hidden_states).view(hidden_shape)
        ).transpose(1, 2)
        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

        # r: same num_heads as q (not num_key_value_heads), rope applied.
        # Note hidden_shape uses -1 for the heads dim, so r_proj (which produces
        # num_attention_heads*head_dim) will correctly view as (..., n_heads, D).
        r_states = self.r_norm(
            self.r_proj(hidden_states).view(hidden_shape)
        ).transpose(1, 2)

        cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb(
            query_states, key_states, cos, sin
        )
        # Apply rope to r according to rope_r mode.
        if self.rope_r == "correct":
            r_states, _ = apply_rotary_pos_emb(r_states, r_states, cos, sin)
        elif self.rope_r == "buggy":
            # Reproduce the old torchtitan bug: reshape_for_broadcast read
            # (B, H, S, D) as (bz, seqlen=H, _, head_dim=D), so it applied
            # positions 0..H-1 along the heads axis.  Every token in the
            # sequence got the SAME rotation (the one for position h, where
            # h is the head index).  This is a fixed per-head rotation that
            # does not vary with sequence position.
            #
            # We lazily compute cos/sin for positions 0..H-1 and cache them.
            # Shape: (1, H, 1, D) so head h gets rotation for position h,
            # broadcast across all sequence positions.
            if self._buggy_rope_cos is None or self._buggy_rope_cos.device != r_states.device:
                H = r_states.shape[1]
                # Reuse the same RoPE frequencies: extract from cos/sin for
                # positions 0..H-1.  cos/sin are (B,1,T,D) during prefill
                # (T >= H always since T >= 1 and H = num_heads).
                # For robustness, recompute from config's rope_theta.
                head_dim = r_states.shape[-1]
                rope_theta = self.config.rope_parameters["rope_theta"]
                inv_freq = 1.0 / (rope_theta ** (
                    torch.arange(0, head_dim, 2, device=r_states.device, dtype=torch.float32) / head_dim
                ))
                pos = torch.arange(H, device=r_states.device, dtype=torch.float32)
                freqs = torch.outer(pos, inv_freq)  # (H, D/2)
                emb = torch.cat([freqs, freqs], dim=-1)  # (H, D)
                self._buggy_rope_cos = emb.cos().view(1, H, 1, head_dim).to(r_states.dtype)
                self._buggy_rope_sin = emb.sin().view(1, H, 1, head_dim).to(r_states.dtype)
            r_states = (r_states * self._buggy_rope_cos) + (rotate_half(r_states) * self._buggy_rope_sin)
        # else: rope_r == "none", skip RoPE on R entirely
        if self.a_proj is not None:
            # a: (B, n_heads, q_len), sigmoid per-head gate
            a = torch.sigmoid(self.a_proj(hidden_states)).transpose(1, 2).to(dtype=r_states.dtype)
            r_states = r_states * (self.amortized_rho_scale * a.unsqueeze(-1))
        elif self.amortized_rho_scale != 1.0:
            r_states = r_states * self.amortized_rho_scale

        if past_key_values is not None:
            key_states, value_states = past_key_values.update(
                key_states, value_states, self.layer_idx
            )

        attn_output, attn_weights = parallax_centered_attention_forward(
            self,
            query_states,
            r_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
        )

        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights


class ParallaxDecoderLayer(Qwen3DecoderLayer):
    def __init__(self, config: ParallaxConfig, layer_idx: int) -> None:
        # Call the grandparent (nn.Module) init to avoid Qwen3DecoderLayer
        # building a Qwen3Attention we'd immediately replace.
        nn.Module.__init__(self)
        self.hidden_size = config.hidden_size
        self.attention_type = config.layer_types[layer_idx]
        self.self_attn = ParallaxAttention(config=config, layer_idx=layer_idx)
        # Reuse the MLP + norms from Qwen3 via its own constructor logic.
        # Easiest is to rebuild them directly here to match Qwen3DecoderLayer.
        from transformers.models.qwen3.modeling_qwen3 import Qwen3MLP

        self.mlp = Qwen3MLP(config)
        self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = Qwen3RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )


class ParallaxPreTrainedModel(Qwen3PreTrainedModel):
    config: ParallaxConfig  # type: ignore[assignment]
    config_class = ParallaxConfig
    base_model_prefix = "model"
    _no_split_modules = ["ParallaxDecoderLayer"]
    # Parallax path is hand-written and does not plug into ALL_ATTENTION_FUNCTIONS.
    # The sdpa path (delegated to super) supports the standard backends.
    _supports_flash_attn = True
    _supports_sdpa = True
    _supports_flex_attn = True


class ParallaxModel(ParallaxPreTrainedModel, Qwen3Model):
    def __init__(self, config: ParallaxConfig) -> None:
        # Bypass Qwen3Model.__init__ (which builds Qwen3DecoderLayers) and
        # rebuild the module tree with ParallaxDecoderLayers.
        Qwen3PreTrainedModel.__init__(self, config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size
        self.embed_tokens = nn.Embedding(
            config.vocab_size, config.hidden_size, self.padding_idx
        )
        self.layers = nn.ModuleList(
            [
                ParallaxDecoderLayer(config, layer_idx)
                for layer_idx in range(config.num_hidden_layers)
            ]
        )
        self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        # Qwen3Model constructs its rotary embedding from the config; reuse.
        from transformers.models.qwen3.modeling_qwen3 import Qwen3RotaryEmbedding

        self.rotary_emb = Qwen3RotaryEmbedding(config=config)
        self.gradient_checkpointing = False
        self.has_sliding_layers = "sliding_attention" in (config.layer_types or [])
        self.post_init()


class ParallaxForCausalLM(ParallaxPreTrainedModel, Qwen3ForCausalLM):
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}

    def __init__(self, config: ParallaxConfig) -> None:
        # Bypass Qwen3ForCausalLM.__init__ (which builds a Qwen3Model) and
        # rebuild with ParallaxModel.
        Qwen3PreTrainedModel.__init__(self, config)
        self.model = ParallaxModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()


__all__ = [
    "ParallaxAttention",
    "ParallaxDecoderLayer",
    "ParallaxForCausalLM",
    "ParallaxModel",
    "ParallaxPreTrainedModel",
    "parallax_centered_attention_forward",
]