inkling_mlx/attention.py

"""Inkling attention: hybrid local/global, per-head q/k RMSNorm, relative-position
logits bias, optional log-scaling, and short-convolution on k/v.

Mirrors ``InklingAttention`` + ``InklingRelativeLogits`` from transformers PR #47347.
This implementation is prefill-oriented (full-sequence, no KV cache); an incremental
cache (including the 4 per-layer conv states) can be layered on top later.
"""

from __future__ import annotations

import math

import mlx.core as mx
import mlx.nn as nn

from .common import RMSNorm, ShortConvolution
from .config import TextConfig

NEG_INF = -1e30


class RelativeLogits(nn.Module):
    """Hidden-state-conditioned relative position bias.

    ``proj`` is a bank of bias-vs-distance profiles ``[d_rel, rel_extent]``. Each
    query's ``d_rel`` relative-state vector mixes them into one bias value per
    backward distance; the bias is zero outside ``0 <= distance < rel_extent``.
    """

    def __init__(self, d_rel: int, rel_extent: int):
        super().__init__()
        self.rel_extent = rel_extent
        self.proj = mx.zeros((d_rel, rel_extent))

    def __call__(self, relative_states, q_pos, kv_pos):
        # relative_states: [B, Lq, heads, d_rel]
        # rel_logits: [B, Lq, heads, rel_extent] -> [B, heads, Lq, rel_extent]
        rel_logits = mx.swapaxes(relative_states @ self.proj, 1, 2)
        B, H, Lq, _ = rel_logits.shape
        distance = q_pos[:, None] - kv_pos[None, :]          # [Lq, Lkv]
        gather = mx.clip(distance, 0, self.rel_extent - 1)   # [Lq, Lkv]
        gather = mx.broadcast_to(gather[None, None], (B, H, Lq, gather.shape[-1]))
        bias = mx.take_along_axis(rel_logits, gather, axis=-1)  # [B, H, Lq, Lkv]
        valid = (distance >= 0) & (distance < self.rel_extent)  # [Lq, Lkv]
        return mx.where(valid[None, None], bias, 0.0)


class Attention(nn.Module):
    def __init__(self, config: TextConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.is_sliding = config.layer_types[layer_idx] == "hybrid_sliding"

        self.head_dim = config.swa_head_dim if self.is_sliding else config.head_dim
        self.num_heads = config.swa_num_attention_heads if self.is_sliding else config.num_attention_heads
        self.num_kv_heads = config.swa_num_key_value_heads if self.is_sliding else config.num_key_value_heads
        self.n_rep = self.num_heads // self.num_kv_heads
        self.sliding_window = config.sliding_window_size if self.is_sliding else None
        self.rel_extent = config.sliding_window_size if self.is_sliding else config.rel_extent
        self.d_rel = config.d_rel
        # q/k are per-head RMS-normalized, hence 1/d rather than 1/sqrt(d)
        self.scaling = 1.0 / self.head_dim

        h = config.hidden_size
        self.wq_du = nn.Linear(h, self.num_heads * self.head_dim, bias=False)
        self.wk_dv = nn.Linear(h, self.num_kv_heads * self.head_dim, bias=False)
        self.wv_dv = nn.Linear(h, self.num_kv_heads * self.head_dim, bias=False)
        self.wr_du = nn.Linear(h, self.num_heads * self.d_rel, bias=False)
        self.wo_ud = nn.Linear(self.num_heads * self.head_dim, h, bias=False)

        self.k_sconv = ShortConvolution(self.num_kv_heads * self.head_dim, config.sconv_kernel_size)
        self.v_sconv = ShortConvolution(self.num_kv_heads * self.head_dim, config.sconv_kernel_size)
        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.rel_logits_proj = RelativeLogits(self.d_rel, self.rel_extent)

    def __call__(self, hidden_states, start_pos=0, kv_cache=None,
                 k_conv=None, v_conv=None, conv_mask=None):
        B, L, _ = hidden_states.shape

        q = self.wq_du(hidden_states)
        k = self.k_sconv(self.wk_dv(hidden_states), mask=conv_mask, cache=k_conv)
        v = self.v_sconv(self.wv_dv(hidden_states), mask=conv_mask, cache=v_conv)
        rel = self.wr_du(hidden_states)

        q = self.q_norm(q.reshape(B, L, self.num_heads, self.head_dim))
        k = self.k_norm(k.reshape(B, L, self.num_kv_heads, self.head_dim))
        v = v.reshape(B, L, self.num_kv_heads, self.head_dim)

        # -> [B, heads, L, head_dim]
        q = q.transpose(0, 2, 1, 3)
        k = k.transpose(0, 2, 1, 3)
        v = v.transpose(0, 2, 1, 3)

        q_pos = mx.arange(L) + start_pos
        if kv_cache is not None:
            k, v = kv_cache.update(k, v)      # full history
        kv_pos = mx.arange(k.shape[2])

        rel = rel.reshape(B, L, self.num_heads, self.d_rel)
        position_bias = self.rel_logits_proj(rel, q_pos, kv_pos)  # [B, heads, Lq, Lkv]

        # log-scaling (global layers only; no-op for context <= n_floor)
        if not self.is_sliding and self.config.log_scaling_n_floor is not None:
            n_floor = self.config.log_scaling_n_floor
            eff_n = (q_pos + 1).astype(mx.float32)
            tau = 1.0 + self.config.log_scaling_alpha * mx.log(
                mx.maximum(eff_n / n_floor, 1.0)
            )
            tau_q = tau.reshape(1, 1, -1, 1)
            q = (q.astype(mx.float32) * tau_q).astype(q.dtype)
            position_bias = (position_bias.astype(mx.float32) * tau_q).astype(position_bias.dtype)

        # Fused SDPA: O = softmax(scale·QKᵀ + mask)·V, softmax in fp32, native GQA
        # (k/v passed un-tiled). mask = relative-position bias + causal/sliding, additive.
        mask = position_bias + self._causal_mask(q_pos, kv_pos)   # [B, heads, Lq, Lkv]
        out = mx.fast.scaled_dot_product_attention(
            q, k, v, scale=self.scaling, mask=mask.astype(q.dtype)
        )                                                        # [B, heads, Lq, head_dim]

        out = out.transpose(0, 2, 1, 3).reshape(B, L, self.num_heads * self.head_dim)
        return self.wo_ud(out)

    def _causal_mask(self, q_pos, kv_pos):
        distance = q_pos[:, None] - kv_pos[None, :]  # [Lq, Lkv]
        allowed = distance >= 0
        if self.sliding_window is not None:
            allowed = allowed & (distance < self.sliding_window)
        mask = mx.where(allowed, 0.0, NEG_INF)
        return mask[None, None].astype(mx.float32)