Blame for mps_ops.py

"""MPS-specific kernels that preserve the model's mathematical operations."""

from __future__ import annotations

import math

import torch
from torch import nn
from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS, _grouped_linear
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS


def attach_expert_lora(
    experts: nn.Module,
    *,
    rank: int,
    alpha: int,
    dropout: float,
) -> None:
    """Attach packed LoRA parameters to one expert module instance."""
    if rank <= 0:
        raise ValueError("Expert LoRA rank must be positive.")
    if alpha <= 0:
        raise ValueError("Expert LoRA alpha must be positive.")
    if not 0.0 <= dropout < 1.0:
        raise ValueError("Expert LoRA dropout must be in [0, 1).")
    parameter_names = (
        "lora_gate_up_a",
        "lora_gate_up_b",
        "lora_down_a",
        "lora_down_b",
    )
    existing = tuple(hasattr(experts, name) for name in parameter_names)
    if any(existing):
        if not all(existing):
            raise RuntimeError("Expert LoRA parameters are only partially initialized.")
        return

    experts.lora_scaling = float(alpha) / rank
    experts.lora_dropout = nn.Dropout(dropout) if dropout else nn.Identity()
    experts.lora_gate_up_a = nn.Parameter(
        torch.empty(
            experts.num_experts,
            rank,
            experts.hidden_dim,
            device=experts.gate_up_proj.device,
            dtype=experts.gate_up_proj.dtype,
        )
    )
    experts.lora_gate_up_b = nn.Parameter(
        torch.zeros(
            experts.num_experts,
            2 * experts.intermediate_dim,
            rank,
            device=experts.gate_up_proj.device,
            dtype=experts.gate_up_proj.dtype,
        )
    )
    experts.lora_down_a = nn.Parameter(
        torch.empty(
            experts.num_experts,
            rank,
            experts.intermediate_dim,
            device=experts.down_proj.device,
            dtype=experts.down_proj.dtype,
        )
    )
    experts.lora_down_b = nn.Parameter(
        torch.zeros(
            experts.num_experts,
            experts.hidden_dim,
            rank,
            device=experts.down_proj.device,
            dtype=experts.down_proj.dtype,
        )
    )
    nn.init.kaiming_uniform_(experts.lora_gate_up_a, a=math.sqrt(5))
    nn.init.kaiming_uniform_(experts.lora_down_a, a=math.sqrt(5))
    experts.gate_up_proj.requires_grad_(False)
    experts.down_proj.requires_grad_(False)


def lora_aware_eager_experts_forward(
    self: nn.Module,
    hidden_states: torch.Tensor,
    top_k_index: torch.Tensor,
    top_k_weights: torch.Tensor,
) -> torch.Tensor:
    final_hidden_states = torch.zeros_like(hidden_states)
    with torch.no_grad():
        expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
        expert_mask = expert_mask.permute(2, 1, 0)
        expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()

    lora_dropout = getattr(self, "lora_dropout", nn.Identity())
    lora_scaling = getattr(self, "lora_scaling", 1.0)

    for expert_idx in expert_hit:
        expert_idx = expert_idx[0]
        if expert_idx == self.num_experts:
            continue
        top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
        current_state = hidden_states[token_idx]
        gate_up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx])
        if hasattr(self, "lora_gate_up_a"):
            update = nn.functional.linear(
                nn.functional.linear(lora_dropout(current_state), self.lora_gate_up_a[expert_idx]),
                self.lora_gate_up_b[expert_idx],
            )
            gate_up = gate_up + lora_scaling * update
        gate, up = gate_up.chunk(2, dim=-1)
        current_hidden_states = self.act_fn(gate) * up
        expert_output = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
        if hasattr(self, "lora_down_a"):
            update = nn.functional.linear(
                nn.functional.linear(lora_dropout(current_hidden_states), self.lora_down_a[expert_idx]),
                self.lora_down_b[expert_idx],
            )
            expert_output = expert_output + lora_scaling * update
        current_hidden_states = expert_output * top_k_weights[token_idx, top_k_pos, None]
        final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))

    return final_hidden_states


def grouped_expert_offsets(
    sorted_expert_ids: torch.Tensor,
    num_experts: int,
) -> torch.Tensor:
    """Build grouped-mm offsets without MPS ``histc`` host synchronization."""
    counts = torch.bincount(sorted_expert_ids, minlength=num_experts)[:num_experts]
    return torch.cumsum(counts, dim=0, dtype=torch.int32)


def mps_grouped_mm_experts_forward(
    self: torch.nn.Module,
    hidden_states: torch.Tensor,
    top_k_index: torch.Tensor,
    top_k_weights: torch.Tensor,
) -> torch.Tensor:
    """Transformers grouped-mm MoE with histogram-free MPS routing offsets."""
    num_top_k = top_k_index.size(-1)
    num_tokens = hidden_states.size(0)
    hidden_dim = hidden_states.size(-1)
    sample_weights = top_k_weights.reshape(-1)
    expert_ids = top_k_index.reshape(-1)

    expert_ids_g, perm = torch.sort(expert_ids)
    selected_hidden_states_g = hidden_states[perm // num_top_k]
    sample_weights_g = sample_weights[perm]
    offsets = grouped_expert_offsets(expert_ids_g, self.num_experts)

    sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1)
    expert_ids_g.clamp_(max=self.num_experts - 1)
    selected_hidden_states_g.masked_fill_(sentinel_mask, 0.0)

    gate_up_weights = self.gate_up_proj if self.has_gate else self.up_proj
    gate_up_biases = (
        self.gate_up_proj_bias[expert_ids_g]
        if self.has_gate and self.has_bias
        else self.up_proj_bias[expert_ids_g]
        if self.has_bias
        else None
    )
    projected = _grouped_linear(
        selected_hidden_states_g,
        gate_up_weights,
        offsets,
        bias=gate_up_biases,
        is_transposed=self.is_transposed,
    )
    projected = self._apply_gate(projected) if self.has_gate else self.act_fn(projected)

    down_biases = self.down_proj_bias[expert_ids_g] if self.has_bias else None
    projected = _grouped_linear(
        projected,
        self.down_proj,
        offsets,
        bias=down_biases,
        is_transposed=self.is_transposed,
    )
    weighted = projected * sample_weights_g.unsqueeze(-1)
    weighted.masked_fill_(sentinel_mask, 0.0)

    # Scatter directly back to the original top-k order.  Constructing the
    # inverse permutation and then gathering performs the same permutation
    # with an extra index tensor and an extra read pass over `weighted`.
    reordered = torch.empty_like(weighted)
    reordered.index_copy_(0, perm, weighted)
    weighted = reordered
    return weighted.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype)


class _IndependentExpertLoRA(torch.autograd.Function):
    """Capacity-padded reference path for independent expert LoRA."""

    @staticmethod
    def forward(
        ctx,
        inputs: torch.Tensor,
        weight_a: torch.Tensor,
        weight_b: torch.Tensor,
        expert_ids: torch.LongTensor,
        counts: torch.LongTensor,
        capacity: int,
    ) -> torch.Tensor:
        offsets = counts.cumsum(0)
        starts = offsets - counts
        slots = torch.arange(inputs.shape[0], device=inputs.device) - starts[expert_ids]
        padded = inputs.new_zeros((weight_a.shape[0], capacity, inputs.shape[-1]))
        padded[expert_ids, slots] = inputs
        low_rank = torch.bmm(padded, weight_a.transpose(1, 2))
        updates = torch.bmm(low_rank, weight_b.transpose(1, 2))
        ctx.save_for_backward(
            padded,
            low_rank,
            weight_a,
            weight_b,
            expert_ids,
            slots,
        )
        return updates[expert_ids, slots]

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        padded, low_rank, weight_a, weight_b, expert_ids, slots = ctx.saved_tensors
        grad_padded = grad_output.new_zeros(
            (weight_a.shape[0], padded.shape[1], grad_output.shape[-1])
        )
        grad_padded[expert_ids, slots] = grad_output
        grad_weight_b = torch.bmm(grad_padded.transpose(1, 2), low_rank)
        grad_low_rank = torch.bmm(grad_padded, weight_b)
        grad_weight_a = torch.bmm(grad_low_rank.transpose(1, 2), padded)
        grad_inputs_padded = torch.bmm(grad_low_rank, weight_a)
        grad_inputs = grad_inputs_padded[expert_ids, slots]
        return grad_inputs, grad_weight_a, grad_weight_b, None, None, None


_INDEPENDENT_EXPERT_LORA_METAL_SOURCE = r"""
#include <metal_stdlib>
using namespace metal;

constant uint kRank = 8;
constant uint kThreads = 256;

kernel void expert_lora_forward_a(
    device const bfloat* inputs [[buffer(0)]],
    device const bfloat* weight_a [[buffer(1)]],
    device const long* expert_ids [[buffer(2)]],
    device bfloat* low_rank [[buffer(3)]],
    constant uint& row_count [[buffer(4)]],
    constant uint& input_dim [[buffer(5)]],
    uint row [[threadgroup_position_in_grid]],
    uint lane [[thread_index_in_threadgroup]]) {
  if (row >= row_count) return;
  const uint expert = uint(expert_ids[row]);
  threadgroup float partial[kThreads * kRank];
  for (uint rank = 0; rank < kRank; ++rank) {
    float value = 0.0f;
    const uint a_base = (expert * kRank + rank) * input_dim;
    const uint x_base = row * input_dim;
    for (uint column = lane; column < input_dim; column += kThreads) {
      value += float(inputs[x_base + column]) * float(weight_a[a_base + column]);
    }
    partial[lane * kRank + rank] = value;
  }
  threadgroup_barrier(mem_flags::mem_threadgroup);
  for (uint stride = kThreads / 2; stride > 0; stride >>= 1) {
    if (lane < stride) {
      for (uint rank = 0; rank < kRank; ++rank) {
        partial[lane * kRank + rank] +=
            partial[(lane + stride) * kRank + rank];
      }
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
  }
  if (lane == 0) {
    for (uint rank = 0; rank < kRank; ++rank) {
      low_rank[row * kRank + rank] = bfloat(partial[rank]);
    }
  }
}

kernel void expert_lora_forward_b(
    device const bfloat* low_rank [[buffer(0)]],
    device const bfloat* weight_b [[buffer(1)]],
    device const long* expert_ids [[buffer(2)]],
    device bfloat* output [[buffer(3)]],
    constant uint& row_count [[buffer(4)]],
    constant uint& output_dim [[buffer(5)]],
    uint index [[thread_position_in_grid]]) {
  const uint total = row_count * output_dim;
  if (index >= total) return;
  const uint row = index / output_dim;
  const uint column = index - row * output_dim;
  const uint expert = uint(expert_ids[row]);
  const uint b_base = (expert * output_dim + column) * kRank;
  float value = 0.0f;
  for (uint rank = 0; rank < kRank; ++rank) {
    value += float(low_rank[row * kRank + rank])
        * float(weight_b[b_base + rank]);
  }
  output[index] = bfloat(value);
}

kernel void expert_lora_backward_low_rank(
    device const bfloat* grad_output [[buffer(0)]],
    device const bfloat* weight_b [[buffer(1)]],
    device const long* expert_ids [[buffer(2)]],
    device bfloat* grad_low_rank [[buffer(3)]],
    constant uint& row_count [[buffer(4)]],
    constant uint& output_dim [[buffer(5)]],
    uint row [[threadgroup_position_in_grid]],
    uint lane [[thread_index_in_threadgroup]]) {
  if (row >= row_count) return;
  const uint expert = uint(expert_ids[row]);
  threadgroup float partial[kThreads * kRank];
  for (uint rank = 0; rank < kRank; ++rank) {
    float value = 0.0f;
    const uint grad_base = row * output_dim;
    for (uint column = lane; column < output_dim; column += kThreads) {
      const uint b_index =
          (expert * output_dim + column) * kRank + rank;
      value += float(grad_output[grad_base + column])
          * float(weight_b[b_index]);
    }
    partial[lane * kRank + rank] = value;
  }
  threadgroup_barrier(mem_flags::mem_threadgroup);
  for (uint stride = kThreads / 2; stride > 0; stride >>= 1) {
    if (lane < stride) {
      for (uint rank = 0; rank < kRank; ++rank) {
        partial[lane * kRank + rank] +=
            partial[(lane + stride) * kRank + rank];
      }
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
  }
  if (lane == 0) {
    for (uint rank = 0; rank < kRank; ++rank) {
      grad_low_rank[row * kRank + rank] = bfloat(partial[rank]);
    }
  }
}

kernel void expert_lora_backward_inputs(
    device const bfloat* grad_low_rank [[buffer(0)]],
    device const bfloat* weight_a [[buffer(1)]],
    device const long* expert_ids [[buffer(2)]],
    device bfloat* grad_inputs [[buffer(3)]],
    constant uint& row_count [[buffer(4)]],
    constant uint& input_dim [[buffer(5)]],
    uint index [[thread_position_in_grid]]) {
  const uint total = row_count * input_dim;
  if (index >= total) return;
  const uint row = index / input_dim;
  const uint column = index - row * input_dim;
  const uint expert = uint(expert_ids[row]);
  float value = 0.0f;
  for (uint rank = 0; rank < kRank; ++rank) {
    const uint a_index =
        (expert * kRank + rank) * input_dim + column;
    value += float(grad_low_rank[row * kRank + rank])
        * float(weight_a[a_index]);
  }
  grad_inputs[index] = bfloat(value);
}

kernel void expert_lora_backward_a(
    device const bfloat* grad_low_rank [[buffer(0)]],
    device const bfloat* inputs [[buffer(1)]],
    device const long* offsets [[buffer(2)]],
    device bfloat* grad_weight_a [[buffer(3)]],
    constant uint& expert_count [[buffer(4)]],
    constant uint& input_dim [[buffer(5)]],
    uint index [[thread_position_in_grid]]) {
  const uint total = expert_count * kRank * input_dim;
  if (index >= total) return;
  const uint column = index % input_dim;
  const uint rank_expert = index / input_dim;
  const uint rank = rank_expert % kRank;
  const uint expert = rank_expert / kRank;
  const uint start = uint(offsets[expert]);
  const uint stop = uint(offsets[expert + 1]);
  float value = 0.0f;
  for (uint row = start; row < stop; ++row) {
    value += float(grad_low_rank[row * kRank + rank])
        * float(inputs[row * input_dim + column]);
  }
  grad_weight_a[index] = bfloat(value);
}

kernel void expert_lora_backward_b(
    device const bfloat* grad_output [[buffer(0)]],
    device const bfloat* low_rank [[buffer(1)]],
    device const long* offsets [[buffer(2)]],
    device bfloat* grad_weight_b [[buffer(3)]],
    constant uint& expert_count [[buffer(4)]],
    constant uint& output_dim [[buffer(5)]],
    uint index [[thread_position_in_grid]]) {
  const uint total = expert_count * output_dim * kRank;
  if (index >= total) return;
  const uint rank = index % kRank;
  const uint output_expert = index / kRank;
  const uint column = output_expert % output_dim;
  const uint expert = output_expert / output_dim;
  const uint start = uint(offsets[expert]);
  const uint stop = uint(offsets[expert + 1]);
  float value = 0.0f;
  for (uint row = start; row < stop; ++row) {
    value += float(grad_output[row * output_dim + column])
        * float(low_rank[row * kRank + rank]);
  }
  grad_weight_b[index] = bfloat(value);
}
"""


_independent_expert_lora_metal_library = None


def _get_independent_expert_lora_metal_library():
    global _independent_expert_lora_metal_library
    if _independent_expert_lora_metal_library is None:
        _independent_expert_lora_metal_library = torch.mps.compile_shader(
            _INDEPENDENT_EXPERT_LORA_METAL_SOURCE
        )
    return _independent_expert_lora_metal_library


class _MetalIndependentExpertLoRA(torch.autograd.Function):
    """Independent rank-8 expert A/B with direct BF16 Metal forward/backward."""

    @staticmethod
    def forward(
        ctx,
        inputs: torch.Tensor,
        weight_a: torch.Tensor,
        weight_b: torch.Tensor,
        expert_ids: torch.LongTensor,
        counts: torch.LongTensor,
    ) -> torch.Tensor:
        inputs = inputs.contiguous()
        weight_a = weight_a.contiguous()
        weight_b = weight_b.contiguous()
        expert_ids = expert_ids.contiguous()
        expert_count, rank, input_dim = weight_a.shape
        row_count = inputs.shape[0]
        output_dim = weight_b.shape[1]
        if rank != 8:
            raise ValueError("The direct Metal expert LoRA kernel requires rank 8.")
        offsets = torch.cat(
            (
                counts.new_zeros(1, dtype=torch.int64),
                counts.cumsum(0, dtype=torch.int64),
            )
        ).contiguous()
        low_rank = inputs.new_empty((row_count, rank))
        output = inputs.new_empty((row_count, output_dim))
        library = _get_independent_expert_lora_metal_library()
        library.expert_lora_forward_a(
            inputs,
            weight_a,
            expert_ids,
            low_rank,
            row_count,
            input_dim,
            threads=(row_count * 256, 1, 1),
            group_size=(256, 1, 1),
        )
        library.expert_lora_forward_b(
            low_rank,
            weight_b,
            expert_ids,
            output,
            row_count,
            output_dim,
            threads=(row_count * output_dim, 1, 1),
            group_size=(256, 1, 1),
        )
        ctx.save_for_backward(
            inputs,
            low_rank,
            weight_a,
            weight_b,
            expert_ids,
            offsets,
        )
        ctx.dimensions = (expert_count, row_count, input_dim, output_dim)
        return output

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        inputs, low_rank, weight_a, weight_b, expert_ids, offsets = ctx.saved_tensors
        expert_count, row_count, input_dim, output_dim = ctx.dimensions
        grad_output = grad_output.contiguous()
        grad_low_rank = low_rank.new_empty(low_rank.shape)
        grad_inputs = inputs.new_empty(inputs.shape)
        grad_weight_a = weight_a.new_empty(weight_a.shape)
        grad_weight_b = weight_b.new_empty(weight_b.shape)
        library = _get_independent_expert_lora_metal_library()
        library.expert_lora_backward_low_rank(
            grad_output,
            weight_b,
            expert_ids,
            grad_low_rank,
            row_count,
            output_dim,
            threads=(row_count * 256, 1, 1),
            group_size=(256, 1, 1),
        )
        library.expert_lora_backward_inputs(
            grad_low_rank,
            weight_a,
            expert_ids,
            grad_inputs,
            row_count,
            input_dim,
            threads=(row_count * input_dim, 1, 1),
            group_size=(256, 1, 1),
        )
        library.expert_lora_backward_a(
            grad_low_rank,
            inputs,
            offsets,
            grad_weight_a,
            expert_count,
            input_dim,
            threads=(expert_count * 8 * input_dim, 1, 1),
            group_size=(256, 1, 1),
        )
        library.expert_lora_backward_b(
            grad_output,
            low_rank,
            offsets,
            grad_weight_b,
            expert_count,
            output_dim,
            threads=(expert_count * output_dim * 8, 1, 1),
            group_size=(256, 1, 1),
        )
        return grad_inputs, grad_weight_a, grad_weight_b, None, None


def _independent_expert_lora(
    inputs: torch.Tensor,
    weight_a: torch.Tensor,
    weight_b: torch.Tensor,
    expert_ids: torch.LongTensor,
    counts: torch.LongTensor,
) -> torch.Tensor:
    """Dispatch the compact independent-expert operator without sharing A/B."""

    if inputs.shape[0] == 0:
        return inputs.new_empty((0, weight_b.shape[1]))
    if (
        inputs.device.type == "mps"
        and inputs.dtype == torch.bfloat16
        and weight_a.dtype == torch.bfloat16
        and weight_b.dtype == torch.bfloat16
        and weight_a.shape[1] == 8
        and hasattr(torch.mps, "compile_shader")
    ):
        return _MetalIndependentExpertLoRA.apply(
            inputs,
            weight_a,
            weight_b,
            expert_ids,
            counts,
        )
    return _IndependentExpertLoRA.apply(
        inputs,
        weight_a,
        weight_b,
        expert_ids,
        counts,
        int(counts.max().detach().cpu()),
    )


def mps_segmented_experts_forward(
    self: torch.nn.Module,
    hidden_states: torch.Tensor,
    top_k_index: torch.Tensor,
    top_k_weights: torch.Tensor,
) -> torch.Tensor:
    """Run GPU-grouped base experts plus independent Metal LoRA A/B."""
    if hidden_states.device.type != "mps":
        return lora_aware_eager_experts_forward(
            self,
            hidden_states,
            top_k_index,
            top_k_weights,
        )

    num_top_k = top_k_index.size(-1)
    num_tokens, hidden_dim = hidden_states.shape
    sample_weights = top_k_weights.reshape(-1)
    expert_ids = top_k_index.reshape(-1)
    expert_ids_g, permutation = torch.sort(expert_ids)
    selected_hidden_g = hidden_states[permutation // num_top_k]
    sample_weights_g = sample_weights[permutation]

    counts = torch.bincount(
        expert_ids_g,
        minlength=self.num_experts,
    )[: self.num_experts]
    offsets = torch.cumsum(counts, dim=0, dtype=torch.int32)
    gate_up_weights = self.gate_up_proj if self.has_gate else self.up_proj
    gate_up_bias = None
    if self.has_bias:
        packed_bias = self.gate_up_proj_bias if self.has_gate else self.up_proj_bias
        gate_up_bias = packed_bias[expert_ids_g]
    gate_up_g = _grouped_linear(
        selected_hidden_g,
        gate_up_weights,
        offsets,
        bias=gate_up_bias,
        is_transposed=self.is_transposed,
    )
    if hasattr(self, "lora_gate_up_a") and self.training:
        gate_up_update = _independent_expert_lora(
            self.lora_dropout(selected_hidden_g),
            self.lora_gate_up_a,
            self.lora_gate_up_b,
            expert_ids_g,
            counts,
        )
        gate_up_g = gate_up_g + gate_up_update * self.lora_scaling
    elif hasattr(self, "lora_gate_up_a"):
        gate_up_update = _independent_expert_lora(
            selected_hidden_g,
            self.lora_gate_up_a,
            self.lora_gate_up_b,
            expert_ids_g,
            counts,
        )
        gate_up_g = gate_up_g + gate_up_update * self.lora_scaling
    activated_g = (
        self._apply_gate(gate_up_g)
        if self.has_gate
        else self.act_fn(gate_up_g)
    )
    down_bias = self.down_proj_bias[expert_ids_g] if self.has_bias else None
    projected_g = _grouped_linear(
        activated_g,
        self.down_proj,
        offsets,
        bias=down_bias,
        is_transposed=self.is_transposed,
    )
    if hasattr(self, "lora_down_a") and self.training:
        down_update = _independent_expert_lora(
            self.lora_dropout(activated_g),
            self.lora_down_a,
            self.lora_down_b,
            expert_ids_g,
            counts,
        )
        projected_g = projected_g + down_update * self.lora_scaling
    elif hasattr(self, "lora_down_a"):
        down_update = _independent_expert_lora(
            activated_g,
            self.lora_down_a,
            self.lora_down_b,
            expert_ids_g,
            counts,
        )
        projected_g = projected_g + down_update * self.lora_scaling
    weighted_g = projected_g * sample_weights_g[:, None]
    weighted = torch.empty_like(weighted_g)
    weighted.index_copy_(0, permutation, weighted_g)
    return weighted.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype)


_SDPA_MASK_FLOOR = -1.0e9
_LONG_KV_CHUNK = 256
_LONG_KV_THRESHOLD = 2048
_original_sdpa_attention_forward = ALL_ATTENTION_FUNCTIONS.get_interface("sdpa", None)


class _MpsGqaChunkedAttention(torch.autograd.Function):
    """Exact GQA attention that never repeats long KV onto every query head.

    Transformers' SDPA path `repeat_kv`s 2 global heads to 16 for full
    attention.  On an 8k prefix that materializes ~2GB of K and V per layer;
    mlp-moe retains all 5 full layers, and the second denoise still holds the
    first latent graph.  MPS then returns NaN LoRA grads with a finite loss.
    Saving compact GQA tensors and reducing K in chunks keeps the math
    identical while staying inside unified memory.
    """

    @staticmethod
    def forward(
        ctx,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_mask: torch.Tensor | None,
        scale: float,
        chunk_size: int,
    ) -> torch.Tensor:
        query_fp32 = query.float()
        key_fp32 = key.float()
        value_fp32 = value.float()
        batch, query_heads, queries, head_dim = query_fp32.shape
        key_heads = int(key_fp32.shape[1])
        key_len = int(key_fp32.shape[-2])
        repeats = query_heads // key_heads
        packed = batch * query_heads
        query_flat = query_fp32.reshape(packed, queries, head_dim)
        negative_inf = query_fp32.new_full((), _SDPA_MASK_FLOOR)
        running_max = query_fp32.new_full((packed, queries, 1), float("-inf"))
        running_sum = torch.zeros_like(running_max)
        acc = torch.zeros(
            packed, queries, head_dim,
            device=query_fp32.device, dtype=query_fp32.dtype,
        )
        chunk = max(int(chunk_size), 1)

        def _pack_kv(tensor: torch.Tensor, start: int, stop: int) -> torch.Tensor:
            width = stop - start
            packed_kv = tensor[:, :, start:stop, :].unsqueeze(2).expand(
                batch, key_heads, repeats, width, head_dim
            )
            return packed_kv.reshape(packed, width, head_dim)

        for start in range(0, key_len, chunk):
            stop = min(start + chunk, key_len)
            key_flat = _pack_kv(key_fp32, start, stop)
            value_flat = _pack_kv(value_fp32, start, stop)
            scores = torch.bmm(query_flat, key_flat.transpose(1, 2)) * scale
            if attn_mask is not None:
                mask_chunk = attn_mask[..., start:stop]
                if mask_chunk.dtype == torch.bool:
                    scores = scores.masked_fill(
                        ~mask_chunk.expand(batch, query_heads, queries, stop - start).reshape(
                            packed, queries, stop - start
                        ),
                        negative_inf,
                    )
                else:
                    scores = scores + mask_chunk.float().expand(
                        batch, query_heads, queries, stop - start
                    ).reshape(packed, queries, stop - start)
            block_max = scores.amax(dim=-1, keepdim=True)
            new_max = torch.maximum(running_max, block_max)
            weights = torch.exp(scores - new_max)
            rescale = torch.exp(running_max - new_max)
            running_sum = running_sum * rescale + weights.sum(dim=-1, keepdim=True)
            acc = acc * rescale + torch.bmm(weights, value_flat)
            running_max = new_max
        output = acc / running_sum.clamp_min(torch.finfo(running_sum.dtype).tiny)
        output = output.view(batch, query_heads, queries, head_dim)
        ctx.save_for_backward(
            query, key, value,
            attn_mask if attn_mask is not None else query.new_empty(0),
        )
        ctx.scale = float(scale)
        ctx.chunk_size = chunk
        ctx.has_mask = attn_mask is not None
        return output.to(dtype=query.dtype)

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        query, key, value, mask_or_empty = ctx.saved_tensors
        attn_mask = mask_or_empty if ctx.has_mask else None
        scale = ctx.scale
        chunk = ctx.chunk_size
        query_fp32 = query.float()
        key_fp32 = key.float()
        value_fp32 = value.float()
        grad_out = grad_output.float()
        batch, query_heads, queries, head_dim = query_fp32.shape
        key_heads = int(key_fp32.shape[1])
        key_len = int(key_fp32.shape[-2])
        repeats = query_heads // key_heads
        packed = batch * query_heads
        query_flat = query_fp32.reshape(packed, queries, head_dim)
        grad_out_flat = grad_out.reshape(packed, queries, head_dim)
        negative_inf = query_fp32.new_full((), _SDPA_MASK_FLOOR)
        running_max = query_fp32.new_full((packed, queries, 1), float("-inf"))
        running_sum = torch.zeros_like(running_max)

        def _pack_kv(tensor: torch.Tensor, start: int, stop: int) -> torch.Tensor:
            width = stop - start
            packed_kv = tensor[:, :, start:stop, :].unsqueeze(2).expand(
                batch, key_heads, repeats, width, head_dim
            )
            return packed_kv.reshape(packed, width, head_dim)

        def _apply_mask(scores: torch.Tensor, start: int, stop: int) -> torch.Tensor:
            if attn_mask is None:
                return scores
            mask_chunk = attn_mask[..., start:stop]
            width = stop - start
            if mask_chunk.dtype == torch.bool:
                return scores.masked_fill(
                    ~mask_chunk.expand(batch, query_heads, queries, width).reshape(
                        packed, queries, width
                    ),
                    negative_inf,
                )
            return scores + mask_chunk.float().expand(
                batch, query_heads, queries, width
            ).reshape(packed, queries, width)

        for start in range(0, key_len, chunk):
            stop = min(start + chunk, key_len)
            scores = _apply_mask(
                torch.bmm(query_flat, _pack_kv(key_fp32, start, stop).transpose(1, 2)) * scale,
                start,
                stop,
            )
            block_max = scores.amax(dim=-1, keepdim=True)
            new_max = torch.maximum(running_max, block_max)
            running_sum = running_sum * torch.exp(running_max - new_max) + torch.exp(
                scores - new_max
            ).sum(dim=-1, keepdim=True)
            running_max = new_max
        grad_query_flat = torch.zeros_like(query_flat)
        grad_key = torch.zeros_like(key_fp32)
        grad_value = torch.zeros_like(value_fp32)
        for start in range(0, key_len, chunk):
            stop = min(start + chunk, key_len)
            width = stop - start
            key_flat = _pack_kv(key_fp32, start, stop)
            value_flat = _pack_kv(value_fp32, start, stop)
            scores = _apply_mask(
                torch.bmm(query_flat, key_flat.transpose(1, 2)) * scale,
                start,
                stop,
            )
            probs = torch.exp(scores - running_max) / running_sum.clamp_min(
                torch.finfo(running_sum.dtype).tiny
            )
            grad_value_packed = torch.bmm(probs.transpose(1, 2), grad_out_flat)
            grad_value[:, :, start:stop, :] = grad_value_packed.view(
                batch, key_heads, repeats, width, head_dim
            ).sum(dim=2)
            grad_probs = torch.bmm(grad_out_flat, value_flat.transpose(1, 2))
            delta = grad_probs - (grad_probs * probs).sum(dim=-1, keepdim=True)
            grad_scores = probs * delta
            grad_query_flat = grad_query_flat + torch.bmm(grad_scores, key_flat) * scale
            grad_key_packed = torch.bmm(grad_scores.transpose(1, 2), query_flat) * scale
            grad_key[:, :, start:stop, :] = grad_key_packed.view(
                batch, key_heads, repeats, width, head_dim
            ).sum(dim=2)
        return (
            grad_query_flat.view(batch, query_heads, queries, head_dim).to(dtype=query.dtype),
            grad_key.to(dtype=key.dtype),
            grad_value.to(dtype=value.dtype),
            None,
            None,
            None,
        )


def mps_safe_sdpa_attention_forward(
    module: nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    dropout: float = 0.0,
    scaling: float | None = None,
    is_causal: bool | None = None,
    position_bias: torch.Tensor | None = None,
    **kwargs,
) -> tuple[torch.Tensor, None]:
    """Use compact GQA attention for long MPS decoder backwards; else official SDPA."""

    kv_length = int(key.shape[-2])
    use_chunked = (
        query.device.type == "mps"
        and torch.is_grad_enabled()
        and query.requires_grad
        and kv_length >= _LONG_KV_THRESHOLD
        and dropout == 0.0
        and position_bias is None
        and not kwargs.get("output_attentions", False)
    )
    if not use_chunked:
        return _original_sdpa_attention_forward(
            module, query, key, value, attention_mask,
            dropout=dropout, scaling=scaling, is_causal=is_causal,
            position_bias=position_bias, **kwargs,
        )
    scale = 1.0 if scaling is None else float(scaling)
    output = _MpsGqaChunkedAttention.apply(
        query, key, value, attention_mask, scale, _LONG_KV_CHUNK,
    )
    return output.transpose(1, 2).contiguous(), None


def register_mps_backends() -> None:
    """Register the schema25 expert and attention implementations explicitly."""

    ALL_EXPERTS_FUNCTIONS.register("lora_eager", lora_aware_eager_experts_forward)
    ALL_EXPERTS_FUNCTIONS.register("modilify_mk2_lora", inference_experts_forward)
    ALL_EXPERTS_FUNCTIONS.register("mps_grouped_mm", mps_grouped_mm_experts_forward)
    ALL_EXPERTS_FUNCTIONS.register("mps_segmented", mps_segmented_experts_forward)
    ALL_ATTENTION_FUNCTIONS.register("sdpa", mps_safe_sdpa_attention_forward)


__all__ = [
    "attach_expert_lora",
    "grouped_expert_offsets",
    "lora_aware_eager_experts_forward",
    "mps_grouped_mm_experts_forward",
    "mps_safe_sdpa_attention_forward",
    "mps_segmented_experts_forward",
    "register_mps_backends",
]


def inference_experts_forward(self, hidden_states, top_k_index, top_k_weights):
    """Use the reference MPS path locally and the Torch fallback elsewhere."""
    implementation = (mps_segmented_experts_forward if hidden_states.device.type == "mps"
                      else lora_aware_eager_experts_forward)
    return implementation(self, hidden_states, top_k_index, top_k_weights)