"""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 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)