feat(v4): gated_attention block — direct re-export of mlx-lm Qwen3NextAttention

The 25% softmax slot in the Qwen3-Next / Qwen3.5 / Qwen3.6 hybrid stack
(3 GDN + 1 Gated Attention per block) is Qwen3NextAttention from upstream
mlx-lm. It has four features that distinguish it from vanilla SDPA:

  - asymmetric GQA (num_attention_heads != num_key_value_heads)
  - per-head RMSNorm on Q and K (post-projection, Qwen3-Next-specific)
  - partial RoPE (partial_rotary_factor controls fraction of head_dim)
  - sigmoid output gate split from a 2x-sized q_proj output

Under the hood it calls mx.fast.scaled_dot_product_attention (Apple MPS
fused softmax SDPA) — there is no faster Metal path on M-series, so
re-implementing it adds nothing. Per the user instruction: import
directly from mlx-lm, do not vendor or rewrite.

This commit:
  - cppmega_v4/nn/gated_attention.py: GatedAttentionBlock wrapper +
    GatedAttentionConfig. Builds the minimal Qwen3Next ModelArgs the
    upstream class reads at __init__ and instantiates Qwen3NextAttention
    directly. The wrapper exposes `self.attn` so weight loaders that walk
    sub-parameter names (q_proj / k_proj / v_proj / o_proj / q_norm /
    k_norm) find them at the expected paths.
  - cppmega_v4/models/unified_superblock_v4.py: register
    "gated_attention" -> _build_gated_attention in BLOCK_BUILDERS so
    UnifiedSuperblock specs can use the Qwen 3.6 / Qwen3-Next attention
    slot directly.
  - tests/v4/test_gated_attention.py: 8 tests — register check, direct-
    import-from-mlx-lm assertion, forward-shape, 2x q_proj gate split,
    asymmetric GQA, Q/K norm presence, unified-superblock round-trip,
    Qwen 3.6-sized config (24 Q / 4 KV / head_dim=128 ratio).

Suite: 379 passed / 0 skipped (was 371).

Author: apstenku123

cppmega_v4/models/unified_superblock_v4.py

@@ -137,6 +137,34 @@ def _build_mla(hidden_size: int, params: dict) -> nn.Module:
     return MLABlock(cfg)
 
 
+def _build_gated_attention(hidden_size: int, params: dict) -> nn.Module:
+    """Qwen3-Next / Qwen3.5 / Qwen3.6 Gated Attention.
+
+    Directly re-exports ``mlx_lm.models.qwen3_next.Qwen3NextAttention`` via
+    ``cppmega_v4.nn.gated_attention.GatedAttentionBlock`` — no vendoring,
+    no re-implementation. Under the hood: ``mx.fast.scaled_dot_product_attention``
+    (Apple MPS) + sigmoid output gate + partial RoPE + asymmetric GQA +
+    Q/K RMSNorm. See cppmega_v4/nn/gated_attention.py for the wrapper.
+    """
+    from cppmega_v4.nn.gated_attention import GatedAttentionBlock, GatedAttentionConfig
+    cfg = GatedAttentionConfig(
+        hidden_size=hidden_size,
+        num_attention_heads=params.get("num_attention_heads", max(1, hidden_size // 64)),
+        num_key_value_heads=params.get(
+            "num_key_value_heads",
+            max(1, params.get("num_attention_heads", max(1, hidden_size // 64)) // 8),
+        ),
+        head_dim=params.get("head_dim", 64),
+        rms_norm_eps=params.get("rms_norm_eps", 1e-6),
+        rope_theta=params.get("rope_theta", 1_000_000.0),
+        partial_rotary_factor=params.get("partial_rotary_factor", 0.25),
+        max_position_embeddings=params.get("max_position_embeddings", 262_144),
+        attention_bias=params.get("attention_bias", False),
+        rope_scaling=params.get("rope_scaling"),
+    )
+    return GatedAttentionBlock(cfg)
+
+
 def _build_attention(hidden_size: int, params: dict) -> nn.Module:
     """Standard multi-head self-attention (causal). Used for `attention`."""
     num_heads = params.get("num_heads", max(1, hidden_size // 64))
@@ -261,6 +289,11 @@ def __call__(self, x):
     "kda": _build_kda,
     "moe": _build_moe,
     "attention": _build_attention,
+    # gated_attention = Qwen3-Next / Qwen3.6 25%-softmax slot. Direct
+    # re-export of mlx_lm.models.qwen3_next.Qwen3NextAttention; output gate
+    # + partial RoPE + asymmetric GQA + Q/K RMSNorm. Underlying compute is
+    # mx.fast.scaled_dot_product_attention (Apple MPS Metal SDPA).
+    "gated_attention": _build_gated_attention,
     # mla = V3-style with LoRA Q + LoRA KV + RoPE on pe-only split.
     # mla_absorb = same block, prefers absorb fast-path at decode.
     "mla": _build_mla,

cppmega_v4/nn/gated_attention.py

@@ -0,0 +1,127 @@
+"""Gated Attention — direct re-export of mlx-lm's Qwen3NextAttention.
+
+mlx-lm already ships the production Gated Attention used by Qwen3-Next /
+Qwen3.5 / Qwen3.6 (the 25% softmax slot in their hybrid 3:1 GDN+Attn
+stack). It implements all four features that distinguish Gated Attention
+from vanilla SDPA:
+
+  - asymmetric GQA (num_attention_heads != num_key_value_heads)
+  - per-head RMSNorm on Q and K (post-projection)
+  - partial RoPE (partial_rotary_factor controls the fraction of head_dim)
+  - sigmoid output gate: ``self.o_proj(SDPA(q, k, v) * sigmoid(gate))``
+    where ``gate`` is split from a 2x-sized q_proj output
+
+Under the hood it calls ``mx.fast.scaled_dot_product_attention`` (the
+Apple MPS fused softmax SDPA kernel) — there is no faster Metal path on
+M-series, so wrapping it ourselves would add nothing. We import the
+upstream module verbatim and just expose it via our V4 block-builder
+interface.
+
+Upstream reference:
+    mlx_lm/models/qwen3_next.py::Qwen3NextAttention
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+# Direct upstream import — no vendoring, no re-implementation. The mlx-lm
+# Qwen3-Next module is part of the installed package; we use it as-is.
+from mlx_lm.models.qwen3_next import (
+    ModelArgs as _Qwen3NextModelArgs,
+    Qwen3NextAttention as _UpstreamGatedAttention,
+)
+
+
+@dataclass(frozen=True)
+class GatedAttentionConfig:
+    """Configuration for the Gated Attention block.
+
+    Defaults match the Qwen3-Next / Qwen3.6 attention slot:
+        num_attention_heads=16, num_key_value_heads=2 (8:1 GQA),
+        head_dim=128, partial_rotary_factor=0.25.
+    """
+
+    hidden_size: int
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 2
+    head_dim: int = 128
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 1_000_000.0
+    partial_rotary_factor: float = 0.25
+    max_position_embeddings: int = 262_144
+    attention_bias: bool = False
+    rope_scaling: Optional[dict] = None
+
+
+def _build_upstream_args(cfg: GatedAttentionConfig) -> _Qwen3NextModelArgs:
+    """Build the minimal mlx-lm ModelArgs that Qwen3NextAttention reads.
+
+    Qwen3NextAttention.__init__ only touches the attention-related fields
+    of ModelArgs; the linear-attn / MoE / norm-vocab fields are unused at
+    instantiation time but required by the @dataclass schema. Fill them
+    with neutral defaults so the dataclass constructor accepts.
+    """
+    return _Qwen3NextModelArgs(
+        model_type="qwen3_next",
+        hidden_size=cfg.hidden_size,
+        num_hidden_layers=1,
+        intermediate_size=cfg.hidden_size * 4,
+        num_attention_heads=cfg.num_attention_heads,
+        # Linear-attn fields (unused by Qwen3NextAttention)
+        linear_num_value_heads=cfg.num_attention_heads,
+        linear_num_key_heads=cfg.num_attention_heads,
+        linear_key_head_dim=cfg.head_dim,
+        linear_value_head_dim=cfg.head_dim,
+        linear_conv_kernel_dim=4,
+        # MoE fields (unused by Qwen3NextAttention)
+        num_experts=1,
+        num_experts_per_tok=1,
+        decoder_sparse_step=1,
+        shared_expert_intermediate_size=cfg.hidden_size,
+        mlp_only_layers=[],
+        moe_intermediate_size=cfg.hidden_size,
+        # Norm / vocab (unused for this block but required by schema)
+        rms_norm_eps=cfg.rms_norm_eps,
+        vocab_size=1,
+        # Attention-specific
+        num_key_value_heads=cfg.num_key_value_heads,
+        rope_theta=cfg.rope_theta,
+        partial_rotary_factor=cfg.partial_rotary_factor,
+        max_position_embeddings=cfg.max_position_embeddings,
+        head_dim=cfg.head_dim,
+        attention_bias=cfg.attention_bias,
+        rope_scaling=cfg.rope_scaling,
+    )
+
+
+class GatedAttentionBlock(nn.Module):
+    """Thin V4-block wrapper around mlx-lm's Qwen3NextAttention.
+
+    Exposes ``self.attn`` so weight loaders that walk ``Qwen3NextAttention``
+    sub-parameter names (q_proj/k_proj/v_proj/o_proj/q_norm/k_norm) find
+    them at the expected paths.
+    """
+
+    def __init__(self, cfg: GatedAttentionConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.attn = _UpstreamGatedAttention(_build_upstream_args(cfg))
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        return self.attn(x, mask=mask, cache=cache)
+
+
+__all__ = [
+    "GatedAttentionBlock",
+    "GatedAttentionConfig",
+]

tests/v4/test_gated_attention.py

@@ -0,0 +1,101 @@
+"""Tests for Gated Attention block (mlx-lm Qwen3NextAttention re-export)."""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import numpy as np
+import pytest
+
+from cppmega_v4.nn.gated_attention import GatedAttentionBlock, GatedAttentionConfig
+from cppmega_v4.models.unified_superblock_v4 import BLOCK_BUILDERS
+
+
+def test_block_registers_in_block_builders():
+    assert "gated_attention" in BLOCK_BUILDERS
+
+
+def test_direct_import_from_mlx_lm():
+    """The wrapper must use mlx-lm's actual Qwen3NextAttention, not a fork."""
+    from mlx_lm.models.qwen3_next import Qwen3NextAttention as _Upstream
+    cfg = GatedAttentionConfig(hidden_size=128, num_attention_heads=4,
+                                num_key_value_heads=2, head_dim=32)
+    block = GatedAttentionBlock(cfg)
+    assert isinstance(block.attn, _Upstream)
+
+
+def test_forward_shape_preserved():
+    cfg = GatedAttentionConfig(
+        hidden_size=256, num_attention_heads=8, num_key_value_heads=2, head_dim=32,
+    )
+    block = GatedAttentionBlock(cfg)
+    x = mx.random.normal((2, 16, 256))
+    y = block(x)
+    mx.eval(y)
+    assert y.shape == x.shape
+
+
+def test_qproj_is_2x_for_gate_split():
+    """q_proj outputs 2x size so half can be split off as the sigmoid gate."""
+    cfg = GatedAttentionConfig(
+        hidden_size=128, num_attention_heads=4, num_key_value_heads=1, head_dim=32,
+    )
+    block = GatedAttentionBlock(cfg)
+    # weight shape is [out, in] in mlx
+    assert block.attn.q_proj.weight.shape == (4 * 32 * 2, 128)
+
+
+def test_asymmetric_gqa():
+    """Gated Attention uses extreme GQA (e.g. 6:1 in Qwen 3.6)."""
+    cfg = GatedAttentionConfig(
+        hidden_size=512, num_attention_heads=24, num_key_value_heads=4, head_dim=128,
+    )
+    block = GatedAttentionBlock(cfg)
+    assert block.attn.q_proj.weight.shape == (24 * 128 * 2, 512)
+    assert block.attn.k_proj.weight.shape == (4 * 128, 512)
+    assert block.attn.v_proj.weight.shape == (4 * 128, 512)
+
+
+def test_q_and_k_norm_present():
+    cfg = GatedAttentionConfig(hidden_size=128, num_attention_heads=4,
+                                num_key_value_heads=1, head_dim=32)
+    block = GatedAttentionBlock(cfg)
+    assert hasattr(block.attn, "q_norm")
+    assert hasattr(block.attn, "k_norm")
+
+
+def test_unified_superblock_builder_round_trip():
+    """Building via BLOCK_BUILDERS produces a working module."""
+    builder = BLOCK_BUILDERS["gated_attention"]
+    block = builder(
+        hidden_size=128,
+        params={
+            "num_attention_heads": 4,
+            "num_key_value_heads": 2,
+            "head_dim": 32,
+        },
+    )
+    x = mx.random.normal((1, 8, 128))
+    y = block(x)
+    mx.eval(y)
+    assert y.shape == x.shape
+    assert not bool(mx.any(mx.isnan(y)).item())
+
+
+def test_qwen36_attention_slot_shape():
+    """Qwen 3.6-27B Gated Attention slot: 24 Q / 4 KV / head_dim=128 / hidden=5120.
+
+    Use shrunk dims so the test is fast but ratios match the production
+    config (8:1 GQA, 2x q_proj for the gate split).
+    """
+    cfg = GatedAttentionConfig(
+        hidden_size=512,           # /10 of 5120
+        num_attention_heads=24,
+        num_key_value_heads=4,     # 6:1 GQA ratio
+        head_dim=64,                # /2 of 128
+        partial_rotary_factor=0.25,
+    )
+    block = GatedAttentionBlock(cfg)
+    x = mx.random.normal((1, 32, 512))
+    y = block(x)
+    mx.eval(y)
+    assert y.shape == x.shape