omlx_support/qwen4_exp.py

# Adapted from oMLX PR #3161 and modified for mixed-bit SSD-backed PLE.
# Licensed under the Apache License 2.0. See LICENSE.

from __future__ import annotations

import json
import math
import mmap
import os
import struct
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional

import mlx.core as mx
import mlx.nn as nn
import numpy as np

from mlx_lm.models.activations import swiglu
from mlx_lm.models.base import BaseModelArgs, create_ssm_mask
from mlx_lm.models.cache import ArraysCache, KVCache, QuantizedKVCache
from mlx_lm.models.gated_delta import gated_delta_update
from mlx_lm.models.switch_layers import SwitchGLU


@dataclass
class ModelArgs(BaseModelArgs):
    model_type: str = "qwen4_exp"
    vocab_size: int = 248320
    hidden_size: int = 2048
    num_hidden_layers: int = 40
    num_attention_heads: int = 16
    num_key_value_heads: int = 2
    hidden_act: str = "silu"
    max_position_embeddings: int = 32768
    rms_norm_eps: float = 1e-6
    tie_word_embeddings: bool = False
    attention_bias: bool = False
    attention_dropout: float = 0.0
    head_dim: int = 256
    linear_conv_kernel_dim: int = 4
    linear_key_head_dim: int = 128
    linear_value_head_dim: int = 128
    linear_num_key_heads: int = 16
    linear_num_value_heads: int = 32
    moe_intermediate_size: int = 512
    shared_expert_intermediate_size: int = 512
    num_experts_per_tok: int = 10
    num_experts: int = 512
    layer_types: List[str] = field(default_factory=list)
    hc_count: int = 4
    hc_lowrank: int = 320
    ple_layer_ids: List[int] = field(default_factory=list)
    ple_embed_dim: Optional[int] = None
    ple_conv_kernel_size: int = 4
    ngram_size: int = 3
    heads_per_ngram: int = 8
    ngram_vocab_size_base: int = 20_000_000
    make_ngram_vocab_size_divisible_by: int = 128
    seed: int = 1234
    split_ngram_parts: int = 512
    indexer_n_heads: Optional[int] = None
    indexer_kv_heads: Optional[int] = None
    indexer_head_dim: Optional[int] = None
    indexer_budget: Optional[int] = None
    indexer_compress_ratio: Optional[int] = None
    norm_topk_prob: bool = True
    output_gate_type: Optional[str] = None
    eos_token_id: Optional[int | List[int]] = None
    rope_theta: float = 10000.0
    partial_rotary_factor: float = 1.0
    rope_parameters: Optional[Dict[str, Any]] = None

    @classmethod
    def from_dict(cls, params):
        source = dict(params.get("text_config", params))
        source["model_type"] = params.get("model_type", source.get("model_type", "qwen4_exp"))
        if params.get("eos_token_id") is not None:
            source["eos_token_id"] = params["eos_token_id"]
        rope = source.get("rope_parameters") or {}
        source.setdefault("rope_theta", rope.get("rope_theta", 10000.0))
        source.setdefault("partial_rotary_factor", rope.get("partial_rotary_factor", 1.0))
        return super().from_dict(source)

    def __post_init__(self):
        if self.ple_embed_dim is None:
            self.ple_embed_dim = self.hidden_size
        if not self.layer_types:
            self.layer_types = [
                "linear_attention" if (i + 1) % 4 else "full_attention"
                for i in range(self.num_hidden_layers)
            ]
        if len(self.layer_types) != self.num_hidden_layers:
            raise ValueError("layer_types must contain one entry per hidden layer")
        if self.hc_count <= 1:
            raise ValueError("hc_count must be greater than one")
        if self.linear_num_value_heads % self.linear_num_key_heads:
            raise ValueError("linear value heads must be divisible by key heads")
        if not 0 < self.num_experts_per_tok <= self.num_experts:
            raise ValueError("num_experts_per_tok must select existing experts")
        if any(layer < 1 or layer > self.num_hidden_layers for layer in self.ple_layer_ids):
            raise ValueError("PLE layer ids are one-indexed hidden-layer ids")
        if self.ple_layer_ids and self.eos_token_id is None:
            raise ValueError("PLE requires eos_token_id")
        qsa = (
            self.indexer_n_heads,
            self.indexer_kv_heads,
            self.indexer_head_dim,
            self.indexer_budget,
            self.indexer_compress_ratio,
        )
        if any(value is None for value in qsa):
            raise ValueError("QSA requires every indexer field")
        if self.indexer_kv_heads != 1:
            raise ValueError("QSA requires one indexer key head")
        if self.indexer_budget % self.indexer_compress_ratio:
            raise ValueError("indexer_budget must divide into complete compressed blocks")
        ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram
        if self.ple_layer_ids and self.ple_embed_dim % ngram_heads:
            raise ValueError("ple_embed_dim must be divisible by its n-gram heads")


class Qwen4RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float, group_size: Optional[int] = None):
        super().__init__()
        self.weight = mx.zeros(dim)
        self.eps = eps
        self.group_size = group_size

    def __call__(self, x):
        dtype = x.dtype
        value = x.astype(mx.float32)
        if self.group_size is not None:
            value = value.reshape(*value.shape[:-1], -1, self.group_size)
        value = value * mx.rsqrt(mx.mean(mx.square(value), axis=-1, keepdims=True) + self.eps)
        if self.group_size is not None:
            value = value.reshape(*x.shape)
        return (value * (1.0 + self.weight.astype(mx.float32))).astype(dtype)


class Qwen4RMSNormGated(nn.Module):
    def __init__(self, dim: int, eps: float, activation: str):
        super().__init__()
        self.weight = mx.ones(dim)
        self.eps = eps
        self.activation = activation

    def __call__(self, x, gate):
        dtype = x.dtype
        value = mx.fast.rms_norm(x, self.weight, self.eps)
        gate = gate.astype(mx.float32)
        gate = mx.sigmoid(gate) if self.activation == "sigmoid" else nn.silu(gate)
        return (value.astype(mx.float32) * gate).astype(dtype)


def _l2_normalize(x, eps=1e-6):
    return x * mx.rsqrt(mx.sum(x * x, axis=-1, keepdims=True) + eps)


def _apply_rope(x, positions, rotary_dim: int, theta: float):
    if rotary_dim == 0:
        return x
    dtype = x.dtype
    inv_freq = theta ** (-mx.arange(0, rotary_dim, 2, dtype=mx.float32) / rotary_dim)
    angles = positions.astype(mx.float32)[..., None] * inv_freq
    cos = mx.concatenate([mx.cos(angles), mx.cos(angles)], axis=-1)
    sin = mx.concatenate([mx.sin(angles), mx.sin(angles)], axis=-1)
    while cos.ndim < x.ndim:
        cos = mx.expand_dims(cos, axis=-2)
        sin = mx.expand_dims(sin, axis=-2)
    rotated, remainder = x[..., :rotary_dim], x[..., rotary_dim:]
    first, second = mx.split(rotated, 2, axis=-1)
    rotated = rotated * cos + mx.concatenate([-second, first], axis=-1) * sin
    return mx.concatenate([rotated, remainder], axis=-1).astype(dtype)


class QSAKVCache(KVCache):
    def __init__(self):
        super().__init__()
        self.index_keys = None

    def update_indexer(self, keys):
        previous = self.offset
        length = keys.shape[1]
        end = previous + length
        if self.index_keys is None or end > self.index_keys.shape[1]:
            batch, _, head_dim = keys.shape
            steps = (self.step + length - 1) // self.step
            extension = mx.zeros(
                (batch, steps * self.step, head_dim),
                dtype=keys.dtype,
            )
            if self.index_keys is None:
                self.index_keys = extension
            else:
                if previous % self.step:
                    self.index_keys = self.index_keys[:, :previous, :]
                self.index_keys = mx.concatenate([self.index_keys, extension], axis=1)
        self.index_keys[:, previous:end, :] = keys
        return self.index_keys[:, :end, :]

    @property
    def state(self):
        index_state = (
            self.index_keys
            if self.index_keys is None
            else self.index_keys[:, : self.offset, :]
        )
        if self.keys is None:
            return self.keys, self.values, index_state
        return (
            self.keys[..., : self.offset, :],
            self.values[..., : self.offset, :],
            index_state,
        )

    @state.setter
    def state(self, value):
        self.keys, self.values, self.index_keys = value
        self.offset = 0 if self.keys is None else self.keys.shape[2]

    def trim(self, count):
        return super().trim(count)

    def to_quantized(self, group_size=64, bits=4):
        cache = QSAQuantizedKVCache(group_size=group_size, bits=bits)
        cache.offset = self.offset
        cache.index_keys = self.index_keys
        if self.keys is not None:
            cache.keys = mx.quantize(self.keys[..., : self.offset, :], group_size=group_size, bits=bits)
            cache.values = mx.quantize(self.values[..., : self.offset, :], group_size=group_size, bits=bits)
        return cache

    @property
    def nbytes(self):
        size = super().nbytes
        return size + (0 if self.index_keys is None else self.index_keys.nbytes)


class QSAQuantizedKVCache(QuantizedKVCache):
    def __init__(self, group_size=64, bits=4):
        super().__init__(group_size=group_size, bits=bits)
        self.index_keys = None

    def update_indexer(self, keys):
        previous = self.offset
        length = keys.shape[1]
        end = previous + length
        if self.index_keys is None or end > self.index_keys.shape[1]:
            batch, _, head_dim = keys.shape
            steps = (self.step + length - 1) // self.step
            extension = mx.zeros(
                (batch, steps * self.step, head_dim),
                dtype=keys.dtype,
            )
            if self.index_keys is None:
                self.index_keys = extension
            else:
                if previous % self.step:
                    self.index_keys = self.index_keys[:, :previous, :]
                self.index_keys = mx.concatenate([self.index_keys, extension], axis=1)
        self.index_keys[:, previous:end, :] = keys
        return self.index_keys[:, :end, :]

    @property
    def state(self):
        quantized_state = (self.keys, self.values) if self.keys is None else super().state
        index_state = (
            self.index_keys
            if self.index_keys is None
            else self.index_keys[:, : self.offset, :]
        )
        return quantized_state, index_state

    @state.setter
    def state(self, value):
        quantized_state, self.index_keys = value
        QuantizedKVCache.state.fset(self, quantized_state)
        self.offset = 0 if self.keys is None else self.keys[0].shape[2]

    def trim(self, count):
        return super().trim(count)

    def to_quantized(self, group_size=64, bits=4):
        if group_size == self.group_size and bits == self.bits:
            return self
        raise ValueError("QSA KV cache is already quantized")

    @property
    def nbytes(self):
        size = super().nbytes
        return size + (0 if self.index_keys is None else self.index_keys.nbytes)


class Qwen4GatedDeltaNet(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.num_v_heads = args.linear_num_value_heads
        self.num_k_heads = args.linear_num_key_heads
        self.head_k_dim = args.linear_key_head_dim
        self.head_v_dim = args.linear_value_head_dim
        self.key_dim = self.num_k_heads * self.head_k_dim
        self.value_dim = self.num_v_heads * self.head_v_dim
        self.conv_kernel_size = args.linear_conv_kernel_dim
        self.conv_dim = 2 * self.key_dim + self.value_dim
        self.conv1d = nn.Conv1d(
            self.conv_dim,
            self.conv_dim,
            self.conv_kernel_size,
            groups=self.conv_dim,
            bias=False,
        )
        self.in_proj_qkv = nn.Linear(args.hidden_size, self.conv_dim, bias=False)
        self.in_proj_z = nn.Linear(args.hidden_size, self.value_dim, bias=False)
        self.in_proj_b = nn.Linear(args.hidden_size, self.num_v_heads, bias=False)
        self.in_proj_a = nn.Linear(args.hidden_size, self.num_v_heads, bias=False)
        self.dt_bias = mx.ones(self.num_v_heads)
        self.A_log = mx.log(mx.random.uniform(low=0.01, high=16.0, shape=(self.num_v_heads,)))
        self.norm = Qwen4RMSNormGated(
            self.head_v_dim,
            args.rms_norm_eps,
            args.output_gate_type or args.hidden_act,
        )
        self.out_proj = nn.Linear(self.value_dim, args.hidden_size, bias=False)

    def __call__(self, x, mask=None, cache=None):
        batch, length, _ = x.shape
        mixed = self.in_proj_qkv(x)
        z = self.in_proj_z(x).reshape(batch, length, self.num_v_heads, self.head_v_dim)
        b = self.in_proj_b(x)
        a = self.in_proj_a(x)
        if mask is not None:
            mixed = mx.where(mask[..., None], mixed, 0)
        if cache is not None and cache[0] is not None:
            conv_state = cache[0]
        else:
            conv_state = mx.zeros((batch, self.conv_kernel_size - 1, self.conv_dim), dtype=x.dtype)
        conv_input = mx.concatenate([conv_state, mixed], axis=1)
        if cache is not None:
            cache[0] = mx.contiguous(conv_input[:, -(self.conv_kernel_size - 1) :, :])
        mixed = nn.silu(self.conv1d(conv_input))
        q, k, value = mx.split(mixed, [self.key_dim, 2 * self.key_dim], axis=-1)
        q = q.reshape(batch, length, self.num_k_heads, self.head_k_dim)
        k = k.reshape(batch, length, self.num_k_heads, self.head_k_dim)
        value = value.reshape(batch, length, self.num_v_heads, self.head_v_dim)
        state = None if cache is None else cache[1]
        q = _l2_normalize(q) * (self.head_k_dim**-0.5)
        k = _l2_normalize(k)
        output, state = gated_delta_update(
            q,
            k,
            value,
            a.astype(mx.float32),
            b,
            self.A_log,
            self.dt_bias,
            state,
            mask,
            use_kernel=not self.training,
        )
        if cache is not None:
            cache[1] = state
            cache.advance(length)
        output = self.norm(output, z).reshape(batch, length, -1)
        return self.out_proj(output)


class Qwen4QSAIndexer(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.n_heads = args.indexer_n_heads
        self.head_dim = args.indexer_head_dim
        self.budget = args.indexer_budget
        self.compress_ratio = args.indexer_compress_ratio
        self.block_topk = self.budget // self.compress_ratio
        self.rotary_dim = int(args.head_dim * args.partial_rotary_factor)
        self.rope_theta = args.rope_theta
        self.index_qk_proj = nn.Linear(
            args.hidden_size,
            (args.indexer_n_heads + args.indexer_kv_heads) * args.indexer_head_dim,
            bias=False,
        )
        self.q_layernorm = Qwen4RMSNorm(self.head_dim, args.rms_norm_eps)
        self.k_layernorm = Qwen4RMSNorm(self.head_dim, args.rms_norm_eps)

    def __call__(self, hidden_states, cache, offset):
        batch, length, _ = hidden_states.shape
        qk = self.index_qk_proj(hidden_states)
        split = self.n_heads * self.head_dim
        query, raw_keys = mx.split(qk, [split], axis=-1)
        query = self.q_layernorm(query.reshape(batch, length, self.n_heads, self.head_dim))
        raw_keys = raw_keys.reshape(batch, length, self.head_dim)
        raw_keys = cache.update_indexer(raw_keys) if cache is not None else raw_keys
        positions = offset + mx.arange(length)
        query = _apply_rope(query, positions[None], self.rotary_dim, self.rope_theta)
        return query, raw_keys, positions

    def select(self, query, raw_keys, positions):
        batch, length, _, _ = query.shape
        key_length = raw_keys.shape[1]
        ratio = self.compress_ratio
        block_count = key_length // ratio
        selected_parts = []
        valid_parts = []
        if block_count:
            pooled = raw_keys[:, : block_count * ratio].reshape(
                batch, block_count, ratio, self.head_dim
            ).mean(axis=2)
            pooled = self.k_layernorm(pooled)
            block_positions = mx.arange(block_count) * ratio
            pooled = _apply_rope(
                pooled,
                block_positions[None],
                self.rotary_dim,
                self.rope_theta,
            )
            scores = mx.einsum(
                "blhd,bkd->blhk",
                query.astype(mx.float32),
                pooled.astype(mx.float32),
            )
            scores = mx.sum(mx.maximum(scores, 0), axis=2) / math.sqrt(self.head_dim)
            complete = (positions + 1) // ratio
            block_valid = mx.arange(block_count)[None, None, :] < complete[None, :, None]
            scores = mx.where(block_valid, scores, mx.finfo(scores.dtype).min)
            take = min(self.block_topk, block_count)
            if take == block_count:
                selected_blocks = mx.broadcast_to(
                    mx.arange(block_count)[None, None, :],
                    (batch, length, block_count),
                )
            else:
                selected_blocks = mx.argpartition(scores, kth=block_count - take, axis=-1)[..., -take:]
            chosen_valid = mx.take_along_axis(block_valid, selected_blocks, axis=-1)
            selected_parts.append(
                (selected_blocks[..., None] * ratio + mx.arange(ratio)).reshape(batch, length, -1)
            )
            valid_parts.append(
                mx.broadcast_to(chosen_valid[..., None], (*chosen_valid.shape, ratio)).reshape(batch, length, -1)
            )
        tail_width = max(ratio - 1, 1)
        tail_start = ((positions + 1) // ratio) * ratio
        tail = tail_start[:, None] + mx.arange(tail_width)[None]
        tail_valid = tail <= positions[:, None]
        selected_parts.append(mx.broadcast_to(tail[None], (batch, length, tail_width)))
        valid_parts.append(mx.broadcast_to(tail_valid[None], (batch, length, tail_width)))
        return mx.concatenate(selected_parts, axis=-1), mx.concatenate(valid_parts, axis=-1)


class Qwen4Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.num_heads = args.num_attention_heads
        self.num_kv_heads = args.num_key_value_heads
        self.head_dim = args.head_dim
        self.repeats = self.num_heads // self.num_kv_heads
        self.scale = self.head_dim**-0.5
        self.rotary_dim = int(self.head_dim * args.partial_rotary_factor)
        self.rope_theta = args.rope_theta
        self.q_proj = nn.Linear(args.hidden_size, self.num_heads * self.head_dim * 2, bias=args.attention_bias)
        self.k_proj = nn.Linear(args.hidden_size, self.num_kv_heads * self.head_dim, bias=args.attention_bias)
        self.v_proj = nn.Linear(args.hidden_size, self.num_kv_heads * self.head_dim, bias=args.attention_bias)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, args.hidden_size, bias=args.attention_bias)
        self.q_norm = Qwen4RMSNorm(self.head_dim, args.rms_norm_eps)
        self.k_norm = Qwen4RMSNorm(self.head_dim, args.rms_norm_eps)
        self.indexer = Qwen4QSAIndexer(args)

    def _select_cache_rows(self, values, batch_index, indices, cache):
        if not isinstance(values, (list, tuple)):
            return values[batch_index, :, indices, :].transpose(2, 0, 1, 3)
        parts = [value[batch_index, :, indices, :].transpose(2, 0, 1, 3) for value in values]
        return mx.dequantize(
            parts[0],
            parts[1],
            parts[2],
            group_size=cache.group_size,
            bits=cache.bits,
        )

    def _sparse_attention(self, query, keys, values, selected, valid, cache):
        outputs = []
        for batch_index in range(query.shape[0]):
            grouped_query = query[batch_index].reshape(
                self.num_kv_heads, self.repeats, query.shape[2], self.head_dim
            )
            chunks = []
            for start in range(0, query.shape[2], 64):
                end = min(start + 64, query.shape[2])
                key_length = (
                    keys[0].shape[2]
                    if isinstance(keys, (list, tuple))
                    else keys.shape[2]
                )
                indices = mx.clip(
                    selected[batch_index, start:end],
                    0,
                    key_length - 1,
                )
                selected_keys = self._select_cache_rows(keys, batch_index, indices, cache)
                selected_values = self._select_cache_rows(values, batch_index, indices, cache)
                local_query = grouped_query[:, :, start:end]
                scores = mx.einsum("hrld,hlmd->hrlm", local_query, selected_keys) * self.scale
                scores = mx.where(valid[batch_index, start:end][None, None], scores, mx.finfo(scores.dtype).min)
                probabilities = mx.softmax(scores, axis=-1, precise=True)
                output = mx.einsum("hrlm,hlmd->hrld", probabilities, selected_values)
                mx.eval(output)
                chunks.append(output)
            output = mx.concatenate(chunks, axis=2)
            outputs.append(output.reshape(self.num_heads, query.shape[2], self.head_dim)[None])
        return mx.concatenate(outputs, axis=0)

    def __call__(self, x, cache=None):
        batch, length, _ = x.shape
        offset = 0 if cache is None else cache.offset
        projected = self.q_proj(x).reshape(batch, length, self.num_heads, 2 * self.head_dim)
        query, gate = mx.split(projected, 2, axis=-1)
        gate = gate.reshape(batch, length, -1)
        keys = self.k_proj(x).reshape(batch, length, self.num_kv_heads, self.head_dim)
        values = self.v_proj(x).reshape(batch, length, self.num_kv_heads, self.head_dim)
        positions = offset + mx.arange(length)
        query = _apply_rope(self.q_norm(query), positions[None], self.rotary_dim, self.rope_theta)
        keys = _apply_rope(self.k_norm(keys), positions[None], self.rotary_dim, self.rope_theta)
        query = query.transpose(0, 2, 1, 3)
        keys = keys.transpose(0, 2, 1, 3)
        values = values.transpose(0, 2, 1, 3)
        index_query, raw_keys, _ = self.indexer(x, cache, offset)
        if cache is not None:
            keys, values = cache.update_and_fetch(keys, values)
        selected, valid = self.indexer.select(index_query, raw_keys, positions)
        output = self._sparse_attention(query, keys, values, selected, valid, cache)
        output = output.transpose(0, 2, 1, 3).reshape(batch, length, -1)
        return self.o_proj(output * mx.sigmoid(gate))


class Qwen4MLP(nn.Module):
    def __init__(self, dim, hidden_dim):
        super().__init__()
        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)

    def __call__(self, x):
        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))


class Qwen4SparseMoeBlock(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.top_k = args.num_experts_per_tok
        self.norm_topk_prob = args.norm_topk_prob
        self.gate = nn.Linear(args.hidden_size, args.num_experts, bias=False)
        self.switch_mlp = SwitchGLU(args.hidden_size, args.moe_intermediate_size, args.num_experts)
        self.shared_expert = Qwen4MLP(args.hidden_size, args.shared_expert_intermediate_size)
        self.shared_expert_gate = nn.Linear(args.hidden_size, 1, bias=False)

    def _routing_weights(self, router_logits):
        gates = mx.softmax(router_logits.astype(mx.float32), axis=-1, precise=True)
        indices = mx.argpartition(gates, kth=-self.top_k, axis=-1)[..., -self.top_k :]
        scores = mx.take_along_axis(gates, indices, axis=-1)
        if self.norm_topk_prob:
            scores = scores / mx.sum(scores, axis=-1, keepdims=True)
        return scores.astype(router_logits.dtype), indices

    def __call__(self, x):
        scores, indices = self._routing_weights(self.gate(x))
        routed = mx.sum(self.switch_mlp(x, indices) * scores[..., None], axis=-2)
        shared = mx.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x)
        return routed + shared


class Qwen4GatedResidual(nn.Module):
    def __init__(self, args: ModelArgs, combine=True):
        super().__init__()
        self.hc_count = args.hc_count
        self.hidden_size = args.hidden_size
        total = self.hc_count * self.hidden_size
        self.hc_norm = Qwen4RMSNorm(total, args.rms_norm_eps, group_size=self.hidden_size)
        self.input_mix_weight_down = nn.Linear(total, args.hc_lowrank, bias=False)
        self.input_mix_weight_up = nn.Linear(args.hc_lowrank, total, bias=False)
        self.block_inject_weight = nn.Linear(total, self.hc_count, bias=False) if combine else None

    def __call__(self, x):
        normalized = self.hc_norm(x)
        weights = nn.silu(self.input_mix_weight_down(normalized) / self.hc_count)
        weights = mx.sigmoid(self.input_mix_weight_up(weights)).reshape(
            *x.shape[:-1], self.hc_count, self.hidden_size
        )
        mixed = mx.mean(
            weights * normalized.reshape(*x.shape[:-1], self.hc_count, self.hidden_size),
            axis=-2,
        )
        if self.block_inject_weight is None:
            return mixed
        injection = 2 * mx.sigmoid(self.block_inject_weight(normalized) / self.hc_count)
        return mixed, x, injection


_MASK64 = (1 << 64) - 1
_SPLITMIX_GAMMA = 0x9E3779B97F4A7C15
_SPLITMIX_M1 = 0xBF58476D1CE4E5B9
_SPLITMIX_M2 = 0x94D049BB133111EB


def _splitmix64(value):
    value = (value + _SPLITMIX_GAMMA) & _MASK64
    value = ((value ^ (value >> 30)) * _SPLITMIX_M1) & _MASK64
    value = ((value ^ (value >> 27)) * _SPLITMIX_M2) & _MASK64
    return (value ^ (value >> 31)) & _MASK64


def _multipliers(vocab_size, ngram_size, layer_index, seed):
    maximum = ((1 << 63) - 1) // max(vocab_size, 1)
    bound = max(1, maximum // 2)
    base = seed + 10007 * layer_index
    return [
        2 * (_splitmix64((base + _SPLITMIX_GAMMA * (index + 1)) & _MASK64) % bound) + 1
        for index in range(ngram_size)
    ]


def _is_prime(value):
    if value < 2:
        return False
    if value % 2 == 0:
        return value == 2
    return all(value % divisor for divisor in range(3, math.isqrt(value) + 1, 2))


def _nth_prime_after(start, count):
    value = start
    for _ in range(count):
        value += 1
        while not _is_prime(value):
            value += 1
    return value


class Qwen4NGramEmbedding(nn.Module):
    def __init__(self, args: ModelArgs, layer_index: int, model_layer_index=None):
        super().__init__()
        self.ngram_size = args.ngram_size
        self.context_len = self.ngram_size - 1
        self.heads_per_ngram = args.heads_per_ngram
        self.ngram_heads = self.context_len * self.heads_per_ngram
        self.eos_token_id = args.eos_token_id[0] if isinstance(args.eos_token_id, list) else args.eos_token_id
        sizes = [
            _nth_prime_after(args.ngram_vocab_size_base - 1, head + 1)
            for head in range(self.ngram_heads)
        ]
        offsets = []
        total = 0
        for size in sizes:
            offsets.append(total)
            total += size
        padded = math.ceil(total / args.make_ngram_vocab_size_divisible_by) * args.make_ngram_vocab_size_divisible_by
        self.layer_multipliers = mx.array(
            _multipliers(args.vocab_size, self.ngram_size, layer_index, args.seed),
            dtype=mx.int64,
        )
        self.ngram_heads_vocab_sizes = mx.array(sizes, dtype=mx.int64)
        self.ngram_heads_offsets = mx.array(offsets, dtype=mx.int64)
        model_path = os.environ.get("OMLX_QWEN4_PLE_MODEL_PATH")
        mode = os.environ.get("OMLX_QWEN4_PLE_MODE", "resident")
        embedding_args = (
            padded,
            args.ple_embed_dim // self.ngram_heads,
            args.split_ngram_parts,
        )
        if mode == "mmap":
            if model_path is None or model_layer_index is None:
                raise RuntimeError("SSD-backed PLE requires its model path and decoder layer index")
            prefix = f"model.layers.{model_layer_index}.ple.ple_embedding.ngram_embedding"
            self.ngram_embedding = DiskBackedShardedEmbedding(
                model_path,
                prefix,
                *embedding_args,
            )
        elif mode == "resident":
            self.ngram_embedding = ShardedEmbedding(*embedding_args)
        else:
            raise ValueError("OMLX_QWEN4_PLE_MODE must be resident or mmap")

    def _shift(self, tokens, shift):
        if shift == 0:
            return tokens
        batch, length = tokens.shape
        positions = mx.arange(length)
        eos_positions = mx.where(tokens == self.eos_token_id, positions[None], -1)
        inclusive = mx.cummax(eos_positions, axis=1)
        previous = mx.concatenate([mx.full((batch, 1), -1, dtype=mx.int64), inclusive[:, :-1]], axis=1)
        source = positions - shift
        gathered = tokens[:, mx.maximum(source, 0)]
        valid = (positions[None] - previous - 1 >= shift) & (source[None] >= 0)
        return mx.where(valid, gathered, self.eos_token_id)

    def __call__(self, input_ids, cache=None):
        input_ids = input_ids.astype(mx.int64)
        if cache is not None and cache[3] is not None:
            previous = cache[3]
        else:
            previous = mx.full((input_ids.shape[0], self.context_len), self.eos_token_id, dtype=mx.int64)
        history = mx.concatenate([previous, input_ids], axis=1)
        if cache is not None:
            cache[3] = mx.contiguous(history[:, -self.context_len :])
        shifted = [self._shift(history, shift) for shift in range(self.ngram_size)]
        blocks = []
        for ngram in range(2, self.ngram_size + 1):
            start = (ngram - 2) * self.heads_per_ngram
            end = start + self.heads_per_ngram
            mixed = shifted[0] * self.layer_multipliers[0]
            for position in range(1, ngram):
                mixed = mx.bitwise_xor(mixed, shifted[position] * self.layer_multipliers[position])
            ids = mixed[..., None] % self.ngram_heads_vocab_sizes[start:end]
            blocks.append(ids + self.ngram_heads_offsets[start:end])
        ids = mx.concatenate(blocks, axis=-1)[:, -input_ids.shape[1] :]
        return self.ngram_embedding(ids).reshape(input_ids.shape[0], input_ids.shape[1], -1)


class ShardedEmbedding(nn.Module):
    def __init__(self, num_embeddings, dims, shard_count):
        super().__init__()
        if num_embeddings % shard_count:
            raise ValueError("the padded n-gram vocabulary must divide evenly into shards")
        self.shard_size = num_embeddings // shard_count
        self.dims = dims
        self.shards = [nn.Embedding(self.shard_size, dims) for _ in range(shard_count)]

    def __call__(self, ids):
        shape = ids.shape
        flat_ids = ids.reshape(-1)
        if flat_ids.size == 0:
            return self.shards[0](flat_ids).reshape(*shape, self.dims)
        shard_ids = flat_ids // self.shard_size
        local_ids = flat_ids % self.shard_size
        mx.eval(shard_ids)
        host_shards = shard_ids.tolist()
        output = None
        for shard_index in sorted(set(host_shards)):
            positions = mx.array(
                [index for index, value in enumerate(host_shards) if value == shard_index],
                dtype=mx.int32,
            )
            values = self.shards[shard_index](local_ids[positions])
            if output is None:
                output = mx.zeros((flat_ids.size, self.dims), dtype=values.dtype)
            output[positions] = values
        mx.eval(output)
        return output.reshape(*shape, self.dims)


_SAFETENSORS_NUMPY_DTYPES = {
    "U32": np.dtype("<u4"),
    "I32": np.dtype("<i4"),
    "I64": np.dtype("<i8"),
    "F16": np.dtype("<f2"),
    "F32": np.dtype("<f4"),
    "BF16": np.dtype("<u2"),
}


class _SafeTensorMMap:
    def __init__(self, path):
        self.path = Path(path)
        self._file = self.path.open("rb")
        raw_header_length = self._file.read(8)
        if len(raw_header_length) != 8:
            self.close()
            raise ValueError(f"Invalid safetensors header in {self.path}")
        header_length = struct.unpack("<Q", raw_header_length)[0]
        raw_header = self._file.read(header_length)
        if len(raw_header) != header_length:
            self.close()
            raise ValueError(f"Truncated safetensors header in {self.path}")
        self._header = json.loads(raw_header)
        self._data_start = 8 + header_length
        self._mapping = mmap.mmap(self._file.fileno(), length=0, access=mmap.ACCESS_READ)
        try:
            self._mapping.madvise(mmap.MADV_RANDOM)
        except (AttributeError, OSError):
            pass

    def tensor_shape(self, key):
        try:
            return tuple(self._header[key]["shape"])
        except KeyError as exc:
            raise KeyError(f"Tensor {key!r} is missing from {self.path}") from exc

    def rows(self, key, row_indices):
        try:
            entry = self._header[key]
        except KeyError as exc:
            raise KeyError(f"Tensor {key!r} is missing from {self.path}") from exc
        dtype_name = entry["dtype"]
        try:
            dtype = _SAFETENSORS_NUMPY_DTYPES[dtype_name]
        except KeyError as exc:
            raise TypeError(f"Unsupported safetensors dtype {dtype_name!r} for {key}") from exc
        shape = tuple(entry["shape"])
        if len(shape) != 2:
            raise ValueError(f"Sparse PLE tensor {key!r} must be two-dimensional")
        start, end = entry["data_offsets"]
        if end - start != math.prod(shape) * dtype.itemsize:
            raise ValueError(f"Invalid byte range for safetensors tensor {key!r}")
        view = np.ndarray(
            shape,
            dtype=dtype,
            buffer=self._mapping,
            offset=self._data_start + start,
        )
        copied = np.array(view[np.asarray(row_indices, dtype=np.intp)], copy=True)
        if dtype_name == "BF16":
            copied = (copied.astype(np.uint32) << np.uint32(16)).view(np.float32)
        return copied, dtype_name

    def close(self):
        mapping = getattr(self, "_mapping", None)
        if mapping is not None:
            mapping.close()
            self._mapping = None
        file_object = getattr(self, "_file", None)
        if file_object is not None:
            file_object.close()
            self._file = None

    def __del__(self):
        try:
            self.close()
        except Exception:
            pass


class DiskBackedShardedEmbedding(nn.Module):
    def __init__(self, model_path, prefix, num_embeddings, dims, shard_count):
        super().__init__()
        if num_embeddings % shard_count:
            raise ValueError("the padded n-gram vocabulary must divide evenly into shards")
        self.shard_size = num_embeddings // shard_count
        self.dims = dims
        self.shard_count = shard_count
        self.last_touched_shards = ()
        self.rows_read = 0
        self._prefix = prefix
        self._readers = {}
        self._tensor_readers = {}
        self._shard_quantization = []

        model_path = Path(model_path)
        index_path = model_path / "model.safetensors.index.json"
        config_path = model_path / "config.json"
        if not index_path.exists() or not config_path.exists():
            raise FileNotFoundError("SSD-backed PLE requires config.json and a safetensors index")
        weight_map = json.loads(index_path.read_text()).get("weight_map", {})
        quantization = json.loads(config_path.read_text()).get("quantization") or {}
        default = {
            "bits": int(quantization.get("bits", 4)),
            "group_size": int(quantization.get("group_size", 64)),
            "mode": quantization.get("mode", "affine"),
        }
        for shard_index in range(shard_count):
            module_key = f"{prefix}.shards.{shard_index}"
            settings = dict(default)
            settings.update(quantization.get(module_key) or {})
            bits = int(settings["bits"])
            group_size = int(settings["group_size"])
            mode = settings["mode"]
            if mode != "affine":
                raise ValueError("SSD-backed PLE supports affine quantization")
            if dims % group_size or (dims * bits) % 32:
                raise ValueError(f"Invalid PLE quantization for {module_key}")
            self._shard_quantization.append((group_size, bits, mode))
            expected_shapes = {
                "weight": (self.shard_size, dims * bits // 32),
                "scales": (self.shard_size, dims // group_size),
                "biases": (self.shard_size, dims // group_size),
            }
            for suffix, expected_shape in expected_shapes.items():
                key = f"{module_key}.{suffix}"
                try:
                    filename = weight_map[key]
                except KeyError as exc:
                    raise KeyError(f"SSD-backed PLE tensor {key!r} is absent from the index") from exc
                reader = self._readers.get(filename)
                if reader is None:
                    reader = _SafeTensorMMap(model_path / filename)
                    self._readers[filename] = reader
                if reader.tensor_shape(key) != expected_shape:
                    raise ValueError(
                        f"Unexpected shape for {key}: {reader.tensor_shape(key)} != {expected_shape}"
                    )
                self._tensor_readers[key] = reader

    def _read_rows(self, key, row_indices):
        array, dtype_name = self._tensor_readers[key].rows(key, row_indices)
        self.rows_read += len(row_indices)
        result = mx.array(array)
        return result.astype(mx.bfloat16) if dtype_name == "BF16" else result

    def __call__(self, ids):
        shape = ids.shape
        flat_ids = ids.reshape(-1).astype(mx.int64)
        mx.eval(flat_ids)
        host_ids = [int(value) for value in flat_ids.tolist()]
        if any(value < 0 or value >= self.shard_size * self.shard_count for value in host_ids):
            raise IndexError("n-gram embedding id is outside the padded vocabulary")
        touched = tuple(sorted({value // self.shard_size for value in host_ids}))
        self.last_touched_shards = touched
        self.rows_read = 0
        output = None
        for shard_index in touched:
            positions_list = [
                index for index, value in enumerate(host_ids) if value // self.shard_size == shard_index
            ]
            local_ids = [host_ids[index] % self.shard_size for index in positions_list]
            base = f"{self._prefix}.shards.{shard_index}"
            weight = self._read_rows(f"{base}.weight", local_ids)
            scales = self._read_rows(f"{base}.scales", local_ids)
            biases = self._read_rows(f"{base}.biases", local_ids)
            group_size, bits, mode = self._shard_quantization[shard_index]
            values = mx.dequantize(
                weight,
                scales=scales,
                biases=biases,
                group_size=group_size,
                bits=bits,
                mode=mode,
            )
            if output is None:
                output = mx.zeros((len(host_ids), self.dims), dtype=values.dtype)
            output[mx.array(positions_list, dtype=mx.int32)] = values
        if output is None:
            output = mx.zeros((0, self.dims), dtype=mx.bfloat16)
        mx.eval(output)
        return output.reshape(*shape, self.dims)


class Qwen4PLELayer(nn.Module):
    def __init__(self, args: ModelArgs, ple_index: int, model_layer_index: int):
        super().__init__()
        total = args.hc_count * args.hidden_size
        self.hc_count = args.hc_count
        self.hidden_size = args.hidden_size
        self.ple_embedding = Qwen4NGramEmbedding(args, ple_index, model_layer_index)
        self.key_proj = nn.Linear(args.ple_embed_dim, total, bias=False)
        self.value_proj = nn.Linear(args.ple_embed_dim, args.hidden_size, bias=False)
        self.norm_key = Qwen4RMSNorm(total, args.rms_norm_eps, args.hidden_size)
        self.norm_query = Qwen4RMSNorm(total, args.rms_norm_eps, args.hidden_size)
        self.norm_conv = Qwen4RMSNorm(total, args.rms_norm_eps, args.hidden_size)
        self.state_len = (args.ple_conv_kernel_size - 1) * args.ngram_size
        self.conv1d = nn.Conv1d(
            total,
            total,
            args.ple_conv_kernel_size,
            dilation=args.ngram_size,
            groups=total,
            bias=False,
        )

    def __call__(self, x, input_ids, cache=None, mask=None):
        embeddings = self.ple_embedding(input_ids, cache)
        key = self.norm_key(self.key_proj(embeddings)).reshape(*x.shape[:-1], self.hc_count, self.hidden_size)
        value = self.value_proj(embeddings)
        query = self.norm_query(x).reshape(*x.shape[:-1], self.hc_count, self.hidden_size)
        gate = mx.sum(key * query, axis=-1, keepdims=True) / math.sqrt(self.hidden_size)
        gate = mx.sign(gate) * mx.sqrt(mx.maximum(mx.abs(gate), 1e-6))
        gated = (mx.sigmoid(gate) * value[..., None, :]).reshape(*x.shape)
        normalized = self.norm_conv(gated)
        if mask is not None:
            gated = mx.where(mask[..., None], gated, 0)
            normalized = mx.where(mask[..., None], normalized, 0)
        if cache is not None and cache[2] is not None:
            state = cache[2]
        else:
            state = mx.zeros((x.shape[0], self.state_len, x.shape[-1]), dtype=x.dtype)
        conv_input = mx.concatenate([state, normalized], axis=1)
        if cache is not None:
            cache[2] = mx.contiguous(conv_input[:, -self.state_len :, :])
        return gated + nn.silu(self.conv1d(conv_input))


class Qwen4DecoderLayer(nn.Module):
    def __init__(self, args: ModelArgs, index: int):
        super().__init__()
        self.is_linear = args.layer_types[index] == "linear_attention"
        if self.is_linear:
            self.linear_attn = Qwen4GatedDeltaNet(args)
        else:
            self.self_attn = Qwen4Attention(args)
        self.mlp = Qwen4SparseMoeBlock(args)
        one_indexed = index + 1
        self.ple = Qwen4PLELayer(args, args.ple_layer_ids.index(one_indexed), index) if one_indexed in args.ple_layer_ids else None
        self.attn_hyper_connection = Qwen4GatedResidual(args)
        self.mlp_hyper_connection = Qwen4GatedResidual(args)

    def __call__(self, x, input_ids, mask=None, cache=None):
        if self.ple is not None:
            x = x + self.ple(x, input_ids, cache, mask)
        mixed, residual, injection = self.attn_hyper_connection(x)
        output = self.linear_attn(mixed, mask, cache) if self.is_linear else self.self_attn(mixed, cache)
        x = residual + (output[..., None, :] * injection[..., None]).reshape(*residual.shape)
        mixed, residual, injection = self.mlp_hyper_connection(x)
        output = self.mlp(mixed)
        return residual + (output[..., None, :] * injection[..., None]).reshape(*residual.shape)


class Qwen4TextModel(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.args = args
        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
        self.layers = [Qwen4DecoderLayer(args, index) for index in range(args.num_hidden_layers)]
        self.hyper_connection_mixer = Qwen4GatedResidual(args, combine=False)

    def __call__(self, inputs, cache=None):
        hidden = self.embed_tokens(inputs)
        hidden = mx.tile(hidden, (1, 1, self.args.hc_count))
        if cache is None:
            cache = [None] * len(self.layers)
        linear_index = next((i for i, layer in enumerate(self.layers) if layer.is_linear), None)
        mask = create_ssm_mask(hidden, cache[linear_index]) if linear_index is not None else None
        for layer, layer_cache in zip(self.layers, cache):
            hidden = layer(hidden, inputs, mask, layer_cache)
        return self.hyper_connection_mixer(hidden)


def _decode_block_fp8(weight, scale, block=128):
    rows, columns = weight.shape
    grid_rows, grid_columns = scale.shape
    padded_rows, padded_columns = grid_rows * block, grid_columns * block
    decoded = mx.from_fp8(weight, dtype=mx.float32)
    decoded = mx.pad(decoded, ((0, padded_rows - rows), (0, padded_columns - columns)))
    decoded = decoded.reshape(grid_rows, block, grid_columns, block)
    decoded = decoded * scale.astype(mx.float32)[:, None, :, None]
    return decoded.reshape(padded_rows, padded_columns)[:rows, :columns].astype(mx.bfloat16)


def register_oq_virtual_tensors(index, config):
    if config.get("model_type") != "qwen4_exp":
        return 0
    text_config = config.get("text_config", config)
    registrations = 0
    shard_count = int(text_config.get("split_ngram_parts", 0))
    for one_indexed_layer in text_config.get("ple_layer_ids", []):
        prefix = (
            f"model.language_model.layers.{one_indexed_layer - 1}.ple."
            "ple_embedding.ngram_embedding"
        )
        scale_key = prefix + ".weight_scale"
        if index.source_shape(scale_key) is None:
            continue
        for shard_index in range(shard_count):
            weight_key = f"{prefix}.shard_{shard_index}.weight"
            shape = index.source_shape(weight_key)
            if shape is None:
                raise ValueError(f"Missing Qwen4 PLE shard: {weight_key}")

            def materialize(weight_key=weight_key, scale_key=scale_key):
                weight = index.load_source(weight_key)
                scale = index.load_source(scale_key)
                value = mx.from_fp8(weight, dtype=mx.bfloat16) * scale
                mx.eval(value)
                return value

            index.register_virtual(
                weight_key,
                shape,
                "BF16",
                materialize,
                hides=(weight_key, scale_key),
            )
            registrations += 1
    return registrations


class Model(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.args = args
        self.model_type = args.model_type
        self.model = Qwen4TextModel(args)
        if not args.tie_word_embeddings:
            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)

    def __call__(self, inputs, cache=None):
        output = self.model(inputs, cache)
        if self.args.tie_word_embeddings:
            return self.model.embed_tokens.as_linear(output)
        return self.lm_head(output)

    @property
    def layers(self):
        return self.model.layers

    def make_cache(self):
        return [ArraysCache(size=4) if layer.is_linear else QSAKVCache() for layer in self.layers]

    def sanitize(self, weights):
        cleaned = {}
        for key, value in weights.items():
            if key.startswith("mtp.") or key.startswith("model.visual."):
                continue
            if key.startswith("model.language_model."):
                key = "model." + key[len("model.language_model.") :]
            cleaned[key] = value
        weights = cleaned
        if self.args.tie_word_embeddings:
            weights.pop("lm_head.weight", None)
        for layer_index in range(self.args.num_hidden_layers):
            prefix = f"model.layers.{layer_index}.mlp"
            if f"{prefix}.experts.0.up_proj.weight" not in weights:
                continue
            for projection in ("up_proj", "gate_proj", "down_proj"):
                values = []
                for expert in range(self.args.num_experts):
                    key = f"{prefix}.experts.{expert}.{projection}.weight"
                    scale_key = key + "_scale_inv"
                    value = weights.pop(key)
                    if scale_key in weights:
                        value = _decode_block_fp8(value, weights.pop(scale_key))
                    values.append(value)
                weights[f"{prefix}.switch_mlp.{projection}.weight"] = mx.stack(values)
        for layer_index in range(self.args.num_hidden_layers):
            base = f"model.layers.{layer_index}.ple.ple_embedding.ngram_embedding"
            layer = self.model.layers[layer_index]
            disk_backed = (
                layer.ple is not None
                and isinstance(
                    layer.ple.ple_embedding.ngram_embedding,
                    DiskBackedShardedEmbedding,
                )
            )
            if disk_backed:
                for key in list(weights):
                    if key.startswith(base + ".shards."):
                        weights.pop(key)
            shard_keys = [key for key in weights if key.startswith(base + ".shard_") and key.endswith(".weight")]
            if not shard_keys:
                continue
            shard_keys.sort(key=lambda key: int(key.split(".shard_")[1].split(".")[0]))
            scale = weights.pop(base + ".weight_scale", None)
            for shard_index, key in enumerate(shard_keys):
                value = weights.pop(key)
                if scale is not None:
                    value = mx.from_fp8(value, dtype=mx.bfloat16) * scale
                weights[f"{base}.shards.{shard_index}.weight"] = value
        for key in list(weights):
            value = weights[key]
            if key.endswith("conv1d.weight") and value.ndim == 3 and value.shape[-1] != 1:
                weights[key] = value.moveaxis(2, 1)
        return weights

    @property
    def quant_predicate(self):
        def predicate(path, module):
            if "ple_embedding.ngram_embedding.shards." in path:
                return {"group_size": 32, "bits": 4}
            if path.endswith("mlp.gate") or path.endswith("shared_expert_gate"):
                return {"group_size": 64, "bits": 8}
            return hasattr(module, "to_quantized")

        return predicate