# Copyright (c) 2026 SandAI. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import importlib import os from dataclasses import dataclass from enum import Enum from typing import Any, Callable, List, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange, repeat from inference.common import Modality, VarlenHandler, is_hopper_arch from inference.infra.parallelism import ulysses_scheduler from inference.utils import env_is_true from torch import Tensor from torch.nn import Parameter from .stage_profile import stage_profile try: from magi_compiler import magi_compile as _magi_compile from magi_compiler.api import magi_register_custom_op as _magi_register_custom_op from magi_compiler.config import CompileConfig as _CompileConfig except Exception: _magi_compile = None _magi_register_custom_op = None _CompileConfig = Any CompileConfig = _CompileConfig if _magi_compile is None or env_is_true("DAVINCI_DISABLE_MAGI_COMPILE") or env_is_true("MAGI_DISABLE_COMPILE"): def magi_compile(*args, **kwargs): def _decorator(obj): return obj return _decorator def magi_register_custom_op(*args, **kwargs): def _decorator(obj): return obj return _decorator else: magi_compile = _magi_compile magi_register_custom_op = _magi_register_custom_op @dataclass class FFAHandler: q_ranges: torch.Tensor k_ranges: torch.Tensor max_seqlen_q: int max_seqlen_k: int attn_type_map: torch.Tensor softmax_scale: float # Define the MLP activation type class MLPActivationType(Enum): """Enumeration of supported activation functions for MLP""" SWIGLU7 = "swiglu7" GELU7 = "gelu7" def swiglu7(x, alpha: float = 1.702, limit: float = 7.0, out_dtype: Optional[torch.dtype] = None): out_dtype = x.dtype if out_dtype is None else out_dtype compute_dtype = x.dtype if x.dtype in (torch.bfloat16, torch.float16) else torch.float32 x = x.to(compute_dtype) x_glu, x_linear = x[..., ::2], x[..., 1::2] # Clamp the input values x_glu = x_glu.clamp(min=None, max=limit) x_linear = x_linear.clamp(min=-limit, max=limit) out_glu = x_glu * torch.sigmoid(alpha * x_glu) # Note we add an extra bias of 1 to the linear layer (from GPT-OSS) return (out_glu * (x_linear + 1)).to(out_dtype) def gelu7(x, alpha: float = 1.702, limit: float = 7.0, out_dtype: Optional[torch.dtype] = None): out_dtype = x.dtype if out_dtype is None else out_dtype compute_dtype = x.dtype if x.dtype in (torch.bfloat16, torch.float16) else torch.float32 x = x.to(compute_dtype) x_glu = x # Clamp the input values x_glu = x_glu.clamp(min=None, max=limit) out_glu = x_glu * torch.sigmoid(alpha * x_glu) # Note we add an extra bias of 1 to the linear layer return out_glu.to(out_dtype) def create_activation_func(activation_type: MLPActivationType) -> Callable: match activation_type: case MLPActivationType.SWIGLU7: return swiglu7 case MLPActivationType.GELU7: return gelu7 case _: raise ValueError(f"Unknown activation type: {activation_type}") class ModalityDispatcher: permuted_modality_mapping: torch.Tensor group_size: torch.Tensor group_size_cpu: list[int] num_modalities: int def __init__(self, modality_mapping: torch.Tensor, num_modalities: int): """ Initialize dispatcher. This runs once during object construction and precomputes all mappings. """ self.modality_mapping = modality_mapping self.num_modalities = num_modalities self.permuted_modality_mapping = self._precompute_permute_mapping(modality_mapping) self.group_size = torch.bincount(self.permuted_modality_mapping, minlength=num_modalities).to(torch.int32) self.group_size_cpu: list[int] = [int(x) for x in self.group_size.to("cpu").tolist()] def _precompute_permute_mapping(self, modality_mapping): # 1. Compute forward and inverse permutation mappings. # argsort is an efficient O(N log N) operation. self.permute_mapping = torch.argsort(modality_mapping) self.inv_permute_mapping = torch.argsort(self.permute_mapping) # 2. Compute group size for each modality. # bincount is highly efficient for counting. permuted_modality_mapping = modality_mapping[self.permute_mapping] return permuted_modality_mapping def dispatch(self, x: torch.Tensor) -> list[torch.Tensor]: grouped_tensors = torch.split(x, self.group_size_cpu, dim=0) return list(grouped_tensors) def undispatch(self, *processed_groups: list[torch.Tensor]) -> torch.Tensor: return torch.cat(processed_groups, dim=0) @staticmethod def permute(x: torch.Tensor, permute_mapping: torch.Tensor) -> torch.Tensor: """Apply forward permutation to tensor.""" return x[permute_mapping] @staticmethod def inv_permute(x: torch.Tensor, inv_permute_mapping: torch.Tensor) -> torch.Tensor: """Apply inverse permutation to tensor.""" return x[inv_permute_mapping] def freq_bands( num_bands: int, temperature: float = 10000.0, step: int = 2, device: Optional[torch.device] = None ) -> torch.Tensor: exp = torch.arange(0, num_bands, step, dtype=torch.int64, device=device).to(torch.float32) / num_bands bands = 1.0 / (temperature**exp) return bands def rotate_half(x, interleaved=False): if not interleaved: x1, x2 = x.chunk(2, dim=-1) out = x.clone() half = out.shape[-1] // 2 out[..., :half] = -x2 out[..., half:] = x1 return out else: out = x.clone() out[..., ::2] = -x[..., 1::2] out[..., 1::2] = x[..., ::2] return out def apply_rotary_emb_torch(x, cos, sin, interleaved=False): """ x: (batch_size, seqlen, nheads, headdim) cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2) """ ro_dim = cos.shape[-1] * 2 assert ro_dim <= x.shape[-1] x_ro = x[..., :ro_dim] x_pass = x[..., ro_dim:] if not interleaved: x1, x2 = x_ro.chunk(2, dim=-1) if cos.ndim == x1.ndim - 1: cos = cos.unsqueeze(-2) sin = sin.unsqueeze(-2) elif cos.ndim == x1.ndim - 2: cos = cos.unsqueeze(0).unsqueeze(-2) sin = sin.unsqueeze(0).unsqueeze(-2) elif cos.ndim < x1.ndim: for _ in range(x1.ndim - cos.ndim): cos = cos.unsqueeze(0) sin = sin.unsqueeze(0) out = x.clone() half = ro_dim // 2 out1 = out[..., :half] out2 = out[..., half:ro_dim] out1.copy_(x1) out1.mul_(cos) out1.addcmul_(x2, sin, value=-1.0) out2.copy_(x2) out2.mul_(cos) out2.addcmul_(x1, sin, value=1.0) return out else: x_ = x_ro.view(*x_ro.shape[:-1], ro_dim // 2, 2) x1 = x_[..., 0] x2 = x_[..., 1] if cos.ndim == x1.ndim - 1: cos = cos.unsqueeze(-2) sin = sin.unsqueeze(-2) elif cos.ndim == x1.ndim - 2: cos = cos.unsqueeze(0).unsqueeze(-2) sin = sin.unsqueeze(0).unsqueeze(-2) elif cos.ndim < x1.ndim: for _ in range(x1.ndim - cos.ndim): cos = cos.unsqueeze(0) sin = sin.unsqueeze(0) out = x.clone() out_ro_ = out[..., :ro_dim].view(*x_.shape) out0 = out_ro_[..., 0] out1 = out_ro_[..., 1] out0.copy_(x1) out0.mul_(cos) out0.addcmul_(x2, sin, value=-1.0) out1.copy_(x2) out1.mul_(cos) out1.addcmul_(x1, sin, value=1.0) return out class ElementWiseFourierEmbed(nn.Module): def __init__( self, dim: int, max_res: int = 224, temperature: float = 10000.0, in_pixels: bool = True, linear_bands: bool = False, learnable: bool = False, device: torch.device = torch.device("cpu"), dtype: torch.dtype = torch.float32, ): """ Args: dim: Output feature dimension, total channels, must be divisible by 6 max_res: Max pixel-frequency resolution for pixel-domain bands temperature: Temperature in inverse-frequency mode in_pixels: True -> pixel-frequency bands, False -> inverse-frequency bands linear_bands: Whether pixel-frequency bands are linearly spaced learnable: Whether frequency bands are trainable """ super().__init__() self.dim = dim self.in_pixels = in_pixels self.learnable = learnable self.temperature = temperature self.max_res = max_res self.linear_bands = linear_bands self.device = device self.dtype = dtype # Make frequency bands trainable or register as buffer bands = self.get_default_bands() if self.learnable: self.bands = nn.Parameter(bands) else: self.register_buffer("bands", bands) def forward(self, coords: torch.Tensor) -> torch.Tensor: """ Args: coords: [L,9], column order (time, row, col, T, H, W, ref_T, ref_H, ref_W) Returns: emb: [L, dim] element-wise Fourier embedding """ # Use slicing instead of unbind + stack to reduce intermediates coords_xyz = coords[:, :3] # [L,3] -> (t, h, w) sizes = coords[:, 3:6] # [L,3] -> (T, H, W) refs = coords[:, 6:9] # [L,3] -> (ref_T, ref_H, ref_W) # Compute scale factors scales = (refs - 1) / (sizes - 1) # [L,3] # NOTE: if both ref and size are 1, scale is fixed to 1; otherwise invalid scales[(refs == 1) & (sizes == 1)] = 1 assert not scales.isnan().any(), "scales has nan" assert not scales.isinf().any(), "scales has inf" # Center alignment: apply to h,w only (not time) centers = (sizes - 1) / 2 # [L,3] centers[:, 0] = 0 # Do not center the time dimension coords_xyz = coords_xyz - centers # [L,3] # Project to frequency bands in one shot: [L,3,B] proj = coords_xyz.unsqueeze(-1) * scales.unsqueeze(-1) * self.bands # Compute sin & cos and concatenate sin_proj = proj.sin() # [L,3,B] cos_proj = proj.cos() return torch.cat((sin_proj, cos_proj), dim=1).flatten(1) def reset_parameters(self): bands = self.get_default_bands() self.bands.copy_(bands) def get_default_bands(self): if self.in_pixels: raise NotImplementedError("in_pixels are not implemented yet") else: bands = freq_bands(self.dim // 8, temperature=self.temperature, step=1, device=self.device).to(self.dtype) return bands class MultiModalityRMSNorm(nn.Module): __constants__ = ["dim", "eps", "num_modality"] dim: int eps: float num_modality: int def __init__(self, dim: int, eps: float = 1e-6, device: torch.device | None = None, num_modality: int = 1): super().__init__() self.dim = dim self.eps = eps self.num_modality = num_modality self.weight = torch.nn.Parameter(torch.zeros(dim * num_modality, device=device, dtype=torch.float32)) if num_modality > 1: self.forward = self.forward_multi_experts else: self.forward = self.forward_single_expert self.reset_parameters() def reset_parameters(self): nn.init.zeros_(self.weight) def rms(self, x: torch.Tensor) -> torch.Tensor: t, original_dtype = x.float(), x.dtype t = t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps) return t def forward_multi_experts(self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher) -> torch.Tensor: original_dtype = x.dtype t = self.rms(x) weight_chunked = self.weight.chunk(self.num_modality, dim=0) outs: list[torch.Tensor] = [] start = 0 for i in range(self.num_modality): size = modality_dispatcher.group_size_cpu[i] if size <= 0: continue seg = (t[start : start + size] * (weight_chunked[i] + 1)).to(original_dtype) outs.append(seg) start += size return torch.cat(outs, dim=0) def forward_single_expert(self, x: torch.Tensor, modality_dispatcher: Optional[ModalityDispatcher] = None) -> torch.Tensor: original_dtype = x.dtype chunk_str = os.getenv("MAGI_RMSNORM_CHUNK", "").strip() try: chunk = int(chunk_str) if chunk_str else 0 except ValueError: chunk = 0 if chunk <= 0 and x.is_cuda and x.dim() == 2: if x.numel() >= 20_000_000: chunk = 4096 if chunk > 0 and x.dim() == 2 and x.shape[0] > chunk: w = (self.weight + 1).to(torch.float32) out = torch.empty_like(x, dtype=original_dtype) n = int(x.shape[0]) for s in range(0, n, int(chunk)): e = min(n, s + int(chunk)) xs = x[s:e].to(torch.float32) xs = xs * torch.rsqrt(torch.mean(xs * xs, dim=-1, keepdim=True) + self.eps) out[s:e] = (xs * w).to(original_dtype) return out t = x.float() t = t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps) return (t * (self.weight + 1)).to(original_dtype) class _BF16ComputeLinear(torch.autograd.Function): @staticmethod def forward( ctx, input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], output_dtype: Optional[torch.dtype], compute_dtype: torch.dtype = torch.bfloat16, ): ctx.compute_dtype = compute_dtype ctx.input_dtype = input.dtype ctx.weight_dtype = weight.dtype ctx.bias_dtype = bias.dtype if bias is not None else None ctx.has_bias = bias is not None ctx.input_shape = tuple(input.shape) input_cast = input.to(compute_dtype) ctx.save_for_backward(input_cast, weight) weight_cast = weight.to(compute_dtype) output = torch.matmul(input_cast, weight_cast.t()) if bias is not None: output = output + bias.to(compute_dtype) return output.to(output_dtype) @staticmethod def backward(ctx, grad_output: torch.Tensor): input_cast, weight = ctx.saved_tensors compute_dtype: torch.dtype = ctx.compute_dtype input_dtype: torch.dtype = ctx.input_dtype input_shape = ctx.input_shape grad_output_f32 = grad_output.to(torch.float32) input_f32 = input_cast.to(torch.float32) weight_f32 = weight.to(torch.float32) grad_output_2d = grad_output_f32.reshape(-1, grad_output_f32.shape[-1]) input_2d = input_f32.reshape(-1, input_f32.shape[-1]) grad_input = grad_output_2d.matmul(weight_f32).reshape(input_shape).to(input_dtype) grad_weight = grad_output_2d.t().matmul(input_2d) grad_bias = grad_output_2d.sum(dim=0) if ctx.has_bias else None return grad_input, grad_weight, grad_bias, None, None class BaseLinear(nn.Module): __constants__ = ["in_features", "out_features", "num_layers", "num_experts"] in_features: int out_features: int num_layers_for_initialization: int num_experts: int weight: Tensor def __init__( self, in_features, out_features, num_layers_for_initialization, num_experts, bias=True, device=None, dtype=None ): super().__init__() factory_kwargs = {"device": device, "dtype": torch.bfloat16} self.in_features = in_features self.out_features = out_features self.num_layers_for_initialization = num_layers_for_initialization self.num_experts = num_experts self.use_bias = bias self.weight = Parameter(torch.empty((out_features * num_experts, in_features), **factory_kwargs)) if bias: self.bias = Parameter(torch.empty(out_features * num_experts, **factory_kwargs)) else: self.register_parameter("bias", None) def forward( self, input: torch.Tensor, output_dtype: Optional[torch.dtype] = None, modality_dispatcher: Optional[ModalityDispatcher] = None, ) -> torch.Tensor: output_dtype = input.dtype if output_dtype is None else output_dtype return _BF16ComputeLinear.apply(input, self.weight, self.bias, output_dtype, torch.bfloat16) class NativeMoELinear(BaseLinear): def forward( self, input: torch.Tensor, output_dtype: Optional[torch.dtype] = None, modality_dispatcher: Optional[ModalityDispatcher] = None, ) -> torch.Tensor: output_dtype = input.dtype if output_dtype is None else output_dtype input_list = modality_dispatcher.dispatch(input) # type: ignore weight_chunked = self.weight.chunk(self.num_experts, dim=0) if self.bias is not None: bias_chunked = self.bias.chunk(self.num_experts, dim=0) for i in range(self.num_experts): input_list[i] = _BF16ComputeLinear.apply( input_list[i], weight_chunked[i], bias_chunked[i] if self.bias is not None else None, output_dtype, torch.bfloat16, ) return modality_dispatcher.undispatch(*input_list) # type: ignore def create_linear( in_features, out_features, num_layers=1, num_experts=1, bias=True, device=None, dtype=None ) -> BaseLinear | NativeMoELinear: if num_experts == 1: return BaseLinear(in_features, out_features, num_layers, num_experts, bias, device, dtype) else: return NativeMoELinear(in_features, out_features, num_layers, num_experts, bias, device, dtype) HAS_MAGI_ATTENTION = importlib.util.find_spec("magi_attention") is not None HAS_FA3 = importlib.util.find_spec("flash_attn_interface") is not None @magi_register_custom_op(name="infra::flash_attn_func", is_subgraph_boundary=True) def flash_attn_func(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: if query.dim() != 4 or key.dim() != 4 or value.dim() != 4: raise ValueError( f"flash_attn_func expects (B, S, H, D). Got q={tuple(query.shape)} k={tuple(key.shape)} v={tuple(value.shape)}" ) if HAS_FA3 and is_hopper_arch(): try: from flash_attn_interface import flash_attn_func as fa3_flash_attn_func return fa3_flash_attn_func(query, key, value) except ModuleNotFoundError: pass try: from flash_attn.flash_attn_interface import flash_attn_func as fa2_flash_attn_func return fa2_flash_attn_func(query, key, value) except ModuleNotFoundError: q = query.permute(0, 2, 1, 3).contiguous() k = key.permute(0, 2, 1, 3).contiguous() v = value.permute(0, 2, 1, 3).contiguous() if k.shape[1] != q.shape[1]: if q.shape[1] % k.shape[1] != 0: raise RuntimeError( f"GQA fallback expects num_heads_q % num_heads_kv == 0. Got num_heads_q={q.shape[1]} num_heads_kv={k.shape[1]}" ) repeat_factor = q.shape[1] // k.shape[1] k = k.repeat_interleave(repeat_factor, dim=1) v = v.repeat_interleave(repeat_factor, dim=1) out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False) return out.permute(0, 2, 1, 3).contiguous() def _split_q_range_with_no_overlap( q_ranges: torch.Tensor, k_ranges: torch.Tensor ) -> Tuple[List[List[int]], List[List[List[int]]]]: range_boundary = torch.unique(q_ranges, sorted=True).tolist() candidates = [[start, end, []] for start, end in zip(range_boundary[:-1], range_boundary[1:])] q_ranges = q_ranges.tolist() k_ranges = k_ranges.tolist() for q_range, k_range in zip(q_ranges, k_ranges): q_start, q_end = q_range for q_range_cand in candidates: if q_start <= q_range_cand[0] and q_range_cand[1] <= q_end: q_range_cand[2].append(k_range) q_ranges_out = [] k_ranges_out = [] for q_range_cand in candidates: if len(q_range_cand[2]) > 0: q_ranges_out.append(q_range_cand[0:2]) k_ranges_out.append(q_range_cand[2]) return q_ranges_out, k_ranges_out def _flash_attn_with_correction( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: List[List[int]], k_range_list: List[List[List[int]]] ): output = torch.zeros_like(query) output_lse = torch.zeros((query.shape[0], query.shape[1]), dtype=torch.float32, device=query.device) from flash_attn.flash_attn_interface import flash_attn_func for q_range, k_ranges in zip(q_ranges, k_range_list): q_start, q_end = q_range qo_out, qo_lse = None, None for k_range in k_ranges: k_start, k_end = k_range cur_qo_out, cur_qo_lse, _ = flash_attn_func( query[q_start:q_end].unsqueeze(0), key[k_start:k_end].unsqueeze(0), value[k_start:k_end].unsqueeze(0), return_attn_probs=True, ) cur_qo_out, cur_qo_lse = cur_qo_out.squeeze(0), cur_qo_lse.squeeze(0) if qo_out is None: qo_out = cur_qo_out qo_lse = cur_qo_lse else: qo_lse[qo_lse == torch.inf] = -torch.inf cur_qo_lse[cur_qo_lse == torch.inf] = -torch.inf max_lse = torch.max(qo_lse, cur_qo_lse) qo_se, cur_qo_se = torch.exp(qo_lse - max_lse), torch.exp(cur_qo_lse - max_lse) sum_se = qo_se + cur_qo_se qo_scale, cur_qo_scale = qo_se / sum_se, cur_qo_se / sum_se qo_out = qo_out * qo_scale.permute(1, 0).unsqueeze(-1) + cur_qo_out * cur_qo_scale.permute(1, 0).unsqueeze(-1) qo_lse = torch.log(sum_se) + max_lse output[q_start:q_end] = qo_out output_lse[q_start:q_end, :] = qo_lse.permute(1, 0) return output, output_lse def _custom_flex_flash_attn_func( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor, **kwargs ): q_ranges, k_range_list = _split_q_range_with_no_overlap(q_ranges, k_ranges) return _flash_attn_with_correction(query, key, value, q_ranges, k_range_list) def _flex_flash_attn_func_infer_output_meta( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: output = torch.empty_like(query) output_lse = torch.empty((query.shape[0], query.shape[1]), dtype=torch.float32, device=query.device) return output, output_lse @magi_register_custom_op( name="infra::flex_flash_attn_func", mutates_args=(), infer_output_meta_fn=_flex_flash_attn_func_infer_output_meta, is_subgraph_boundary=True, ) def flex_flash_attn_func( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: if HAS_MAGI_ATTENTION and is_hopper_arch(): from magi_attention.api import flex_flash_attn_func as magi_flex_flash_attn_func return magi_flex_flash_attn_func(query, key, value, q_ranges, k_ranges) else: return _custom_flex_flash_attn_func(query, key, value, q_ranges, k_ranges) def _attention_with_cp_infer_output_meta(q: torch.Tensor, *args, **kwargs) -> torch.Tensor: return torch.empty_like(q, dtype=torch.bfloat16).squeeze(0) @magi_register_custom_op( name="infra::flash_attn_with_cp", mutates_args=(), infer_output_meta_fn=_attention_with_cp_infer_output_meta, is_subgraph_boundary=True, ) def flash_attn_with_cp(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cp_split_sizes: List[int]) -> torch.Tensor: q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16) from inference.infra.distributed import get_cp_group, get_cp_rank, get_cp_world_size from inference.infra.parallelism.all_to_all_primitive import batch_scatter_head_gather_seqlen, scatter_seqlen_gather_head from inference.utils import print_rank_0 # One-shot probe so we can tell all_to_all hang from flash_attn hang. _probe = not hasattr(flash_attn_with_cp, "_probed") if _probe and get_cp_rank() == 0: print_rank_0( f"flash_attn_with_cp: enter q={tuple(q.shape)} split={cp_split_sizes} cp={get_cp_world_size()}" ) if get_cp_world_size() > 1: q, k, v = batch_scatter_head_gather_seqlen([q.squeeze(0), k.squeeze(0), v.squeeze(0)], cp_split_sizes, get_cp_group()) q = q.unsqueeze(0) k = k.unsqueeze(0) v = v.unsqueeze(0) if _probe: torch.cuda.synchronize() if get_cp_rank() == 0: print_rank_0(f"flash_attn_with_cp: after all_to_all q={tuple(q.shape)}") if env_is_true("DAVINCI_DISABLE_MAGI_COMPILE") or env_is_true("MAGI_DISABLE_COMPILE"): self_attn_out = flash_attn_func(q, k, v).squeeze(0) else: self_attn_out = torch.ops.infra.flash_attn_func(q, k, v).squeeze(0) if _probe: torch.cuda.synchronize() if get_cp_rank() == 0: print_rank_0(f"flash_attn_with_cp: after flash_attn out={tuple(self_attn_out.shape)}") if get_cp_world_size() > 1: self_attn_out = scatter_seqlen_gather_head(self_attn_out, cp_split_sizes, get_cp_group(), async_op=False) self_attn_out = rearrange(self_attn_out, "(cp sq) hn hd -> sq (cp hn) hd", cp=get_cp_world_size()) if _probe: torch.cuda.synchronize() if get_cp_rank() == 0: print_rank_0(f"flash_attn_with_cp: exit out={tuple(self_attn_out.shape)}") flash_attn_with_cp._probed = True return self_attn_out @magi_register_custom_op( name="infra::flex_flash_attn_with_cp", mutates_args=(), infer_output_meta_fn=_attention_with_cp_infer_output_meta, is_subgraph_boundary=True, ) def flex_flash_attn_with_cp( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor, cp_split_sizes: List[int], ) -> torch.Tensor: q, k, v = q.to(torch.bfloat16).squeeze(0), k.to(torch.bfloat16).squeeze(0), v.to(torch.bfloat16).squeeze(0) from inference.infra.distributed import get_cp_group, get_cp_world_size from inference.infra.parallelism.all_to_all_primitive import batch_scatter_head_gather_seqlen, scatter_seqlen_gather_head if get_cp_world_size() > 1: q, k, v = batch_scatter_head_gather_seqlen([q, k, v], cp_split_sizes, get_cp_group()) if env_is_true("DAVINCI_DISABLE_MAGI_COMPILE") or env_is_true("MAGI_DISABLE_COMPILE"): out, _ = flex_flash_attn_func(q, k, v, q_ranges=q_ranges, k_ranges=k_ranges) else: out, _ = torch.ops.infra.flex_flash_attn_func(q, k, v, q_ranges=q_ranges, k_ranges=k_ranges) if get_cp_world_size() > 1: out = scatter_seqlen_gather_head(out, cp_split_sizes, get_cp_group(), async_op=False) out = rearrange(out, "(cp sq) hn hd -> sq (cp hn) hd", cp=get_cp_world_size()) return out @dataclass class AttentionConfig: hidden_size: int num_heads_q: int num_heads_kv: int head_dim: int params_dtype: torch.dtype checkpoint_qk_layernorm_rope: bool num_modality: int num_layers: int use_local_attn: bool = False enable_attn_gating: bool = False class Attention(torch.nn.Module): config: AttentionConfig def __init__(self, config: AttentionConfig): super().__init__() self.config = config self.pre_norm = MultiModalityRMSNorm(config.hidden_size, eps=1e-6, num_modality=config.num_modality) self.gating_size = config.num_heads_q if config.enable_attn_gating else 0 self.linear_qkv = create_linear( config.hidden_size, config.num_heads_q * config.head_dim + config.num_heads_kv * config.head_dim * 2 + self.gating_size, num_experts=config.num_modality, bias=False, dtype=config.params_dtype, num_layers=config.num_layers, ) self.linear_proj = create_linear( config.num_heads_q * config.head_dim, config.hidden_size, bias=False, num_experts=config.num_modality, dtype=config.params_dtype, num_layers=config.num_layers, ) self.q_norm = MultiModalityRMSNorm(config.head_dim, num_modality=config.num_modality) self.k_norm = MultiModalityRMSNorm(config.head_dim, num_modality=config.num_modality) self.q_size = config.num_heads_q * config.head_dim self.kv_size = config.num_heads_kv * config.head_dim def reset_parameters(self): if hasattr(self.linear_proj, "reset_parameters_output_layer"): self.linear_proj.reset_parameters_output_layer() def forward( self, hidden_states: torch.Tensor, rope: torch.Tensor, permute_mapping: torch.Tensor, inv_permute_mapping: torch.Tensor, varlen_handler: VarlenHandler, local_attn_handler: FFAHandler, modality_dispatcher: ModalityDispatcher, cp_split_sizes: List[int], ) -> torch.Tensor: hidden_states = self.pre_norm(hidden_states, modality_dispatcher=modality_dispatcher).to(torch.bfloat16) qkv: torch.Tensor = self.linear_qkv(hidden_states, modality_dispatcher=modality_dispatcher).to(hidden_states.dtype) q, k, v, g = torch.split(qkv, [self.q_size, self.kv_size, self.kv_size, self.gating_size], dim=1) q = q.view(-1, self.config.num_heads_q, self.config.head_dim) k = k.view(-1, self.config.num_heads_kv, self.config.head_dim) v = v.view(-1, self.config.num_heads_kv, self.config.head_dim) g = g.view(k.shape[0], self.config.num_heads_q, -1) q = self.q_norm(q, modality_dispatcher=modality_dispatcher) k = self.k_norm(k, modality_dispatcher=modality_dispatcher) q = ModalityDispatcher.inv_permute(q, inv_permute_mapping).unsqueeze(0) k = ModalityDispatcher.inv_permute(k, inv_permute_mapping).unsqueeze(0) v = ModalityDispatcher.inv_permute(v, inv_permute_mapping).unsqueeze(0) sin_emb, cos_emb = rope.tensor_split(2, -1) q = apply_rotary_emb_torch(q, cos_emb, sin_emb) k = apply_rotary_emb_torch(k, cos_emb, sin_emb) if self.config.use_local_attn: self_attn_out = flex_flash_attn_with_cp( q, k, v, local_attn_handler.q_ranges, local_attn_handler.k_ranges, cp_split_sizes ) else: self_attn_out = flash_attn_with_cp(q, k, v, cp_split_sizes) self_attn_out = ModalityDispatcher.permute(self_attn_out, permute_mapping) if self.config.enable_attn_gating: self_attn_out = self_attn_out * torch.sigmoid(g) self_attn_out = self_attn_out.view(-1, self.config.num_heads_q * self.config.head_dim).to(torch.bfloat16) out = self.linear_proj(self_attn_out, modality_dispatcher=modality_dispatcher) return out @dataclass class MLPConfig: hidden_size: int intermediate_size: int activation_type: MLPActivationType params_dtype: torch.dtype num_modality: int = 1 num_layers: int = 1 gated_act: bool = False class MLP(torch.nn.Module): config: MLPConfig def __init__(self, config: MLPConfig): super().__init__() num_experts = config.num_modality self.pre_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=config.num_modality) intermediate_size_up = config.intermediate_size * 2 if config.gated_act else config.intermediate_size self.up_gate_proj = create_linear( config.hidden_size, intermediate_size_up, bias=False, dtype=config.params_dtype, num_layers=config.num_layers, num_experts=num_experts, ) self.down_proj = create_linear( config.intermediate_size, config.hidden_size, bias=False, dtype=config.params_dtype, num_layers=config.num_layers, num_experts=num_experts, ) self.activation_func = create_activation_func(config.activation_type) def forward(self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher) -> torch.Tensor: x = self.pre_norm(x, modality_dispatcher=modality_dispatcher).to(torch.bfloat16) x = self.up_gate_proj(x, modality_dispatcher=modality_dispatcher).to(torch.bfloat16) x = self.activation_func(x).to(torch.bfloat16) x = self.down_proj(x, modality_dispatcher=modality_dispatcher).to(torch.bfloat16) return x def extra_repr(self) -> str: return f"{self.up_gate_proj.weight.shape=}, {self.down_proj.weight.shape=}" class ExpertMLP(torch.nn.Module): """Single dense expert FFN (SwiGLU7 / GELU7), used inside SparseMoEMLP.""" def __init__( self, *, hidden_size: int, intermediate_size: int, activation_type: MLPActivationType, gated_act: bool, params_dtype: torch.dtype, ): super().__init__() self.hidden_size = int(hidden_size) self.intermediate_size = int(intermediate_size) self.gated_act = bool(gated_act) intermediate_size_up = intermediate_size * 2 if gated_act else intermediate_size self.up_gate_proj = nn.Linear(hidden_size, intermediate_size_up, bias=False, dtype=params_dtype) self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False, dtype=params_dtype) self.activation_func = create_activation_func(activation_type) def forward_with_intermediate(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: projected = self.up_gate_proj(x.to(torch.bfloat16)).to(torch.bfloat16) z = self.activation_func(projected).to(torch.bfloat16) y = self.down_proj(z).to(torch.bfloat16) return y, z def forward(self, x: torch.Tensor) -> torch.Tensor: y, _ = self.forward_with_intermediate(x) return y @dataclass class SparseMoEConfig: hidden_size: int intermediate_size: int activation_type: MLPActivationType params_dtype: torch.dtype num_experts: int = 3 top_k: int = 1 gated_act: bool = True router_noise: float = 0.0 num_modality: int = 1 layer_idx: int = -1 def token_sparse_moe_mlp_layers( sparse_moe_layers: Optional[List[int]], mm_layers: Optional[List[int]] = None, convert_mm_mlp: bool = False, ) -> List[int]: """MLP layers that use token-level SparseMoEMLP (learned top-k router). When ``convert_mm_mlp`` is true and ``sparse_moe_layers`` is non-empty, mm-layer MLPs are converted too: modality-hard NativeMoELinear is replaced by per-token top-k routing. Attention on mm_layers stays per-modality. Empty ``sparse_moe_layers`` leaves the MagiHuman MLP unchanged. """ sparse = [int(x) for x in (sparse_moe_layers or [])] if not sparse: return [] if not convert_mm_mlp: return sorted(set(sparse)) mm = [int(x) for x in (mm_layers or [])] return sorted(set(sparse) | set(mm)) def _env_optional_int(name: str) -> Optional[int]: raw = os.environ.get(name, "").strip() if not raw: return None return int(raw) class SparseMoEMLP(torch.nn.Module): """ Token-level sparse MoE MLP. Each token independently scores E experts with Linear(hidden -> E), softmax, then top-k (default k=1). Only selected experts run on that token; outputs are gate-weighted. This is Mixtral/Switch-style routing, not MagiHuman modality-hard NativeMoELinear. - Router: Linear(hidden -> E), top-k softmax gates - Experts: E independent ExpertMLP - Aux loss: Switch-Transformer style load balancing on top-1 assignment - pre_norm may keep per-modality RMSNorm affine on mm_layers (not routing) """ def __init__(self, config: SparseMoEConfig): super().__init__() if config.num_experts < 1: raise ValueError(f"num_experts must be >= 1, got {config.num_experts}") if config.top_k < 1 or config.top_k > config.num_experts: raise ValueError(f"top_k must be in [1, num_experts], got top_k={config.top_k}, E={config.num_experts}") self.num_experts = int(config.num_experts) self.top_k = int(config.top_k) self.router_noise = float(config.router_noise) self.layer_idx = int(getattr(config, "layer_idx", -1)) self.pre_norm = MultiModalityRMSNorm( config.hidden_size, num_modality=int(getattr(config, "num_modality", 1) or 1) ) self.router = nn.Linear(config.hidden_size, self.num_experts, bias=False, dtype=torch.float32) self.experts = nn.ModuleList( [ ExpertMLP( hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, activation_type=config.activation_type, gated_act=config.gated_act, params_dtype=config.params_dtype, ) for _ in range(self.num_experts) ] ) self._last_aux_loss: Optional[torch.Tensor] = None self._last_expert_frac: Optional[torch.Tensor] = None self._last_router_entropy: Optional[torch.Tensor] = None self._last_mod_counts: Optional[torch.Tensor] = None self._mid_observer = None self.diffusion_timestep = None self.diffusion_step_index = None self.diffusion_num_steps = None self.reset_router_parameters() def set_mid_observer(self, observer) -> None: """Optional routed-token observer: observer(expert_id, modality, gate, z, y).""" self._mid_observer = observer def reset_router_parameters(self): nn.init.normal_(self.router.weight, mean=0.0, std=0.01) def _compute_aux_loss(self, probs: torch.Tensor, top1_idx: torch.Tensor) -> torch.Tensor: # Switch Transformer: aux = E * sum_i (f_i * P_i) # f_i: fraction of tokens whose top-1 expert is i # P_i: mean router probability for expert i num_tokens = probs.shape[0] if num_tokens == 0: zeros = probs.new_zeros(()) self._last_expert_frac = probs.new_zeros((self.num_experts,)) self._last_router_entropy = zeros return zeros one_hot = F.one_hot(top1_idx, num_classes=self.num_experts).to(probs.dtype) f = one_hot.mean(dim=0) p = probs.mean(dim=0) eps = probs.new_tensor(1e-9) self._last_expert_frac = f.detach() self._last_router_entropy = (-(f.clamp_min(eps) * f.clamp_min(eps).log()).sum()).detach() return self.num_experts * torch.sum(f * p) def forward(self, x: torch.Tensor, modality_dispatcher: Optional[ModalityDispatcher] = None) -> torch.Tensor: _sp = stage_profile() _router_tok = _sp.start(f"layer.{self.layer_idx}.moe.router") x_norm = self.pre_norm(x, modality_dispatcher=modality_dispatcher).to(torch.bfloat16) logits = self.router(x_norm.float()) if self.training and self.router_noise > 0: logits = logits + torch.randn_like(logits) * self.router_noise # Infer-only ablations. Training ignores these env vars. if not self.training: dropped = _env_optional_int("MAGI_MOE_DROP_EXPERT") if dropped is not None: if dropped < 0 or dropped >= self.num_experts: raise ValueError( f"MAGI_MOE_DROP_EXPERT={dropped} out of range [0, {self.num_experts})" ) logits = logits.clone() logits[:, dropped] = float("-inf") probs = F.softmax(logits, dim=-1) topk_vals, topk_idx = torch.topk(probs, k=self.top_k, dim=-1) topk_vals = topk_vals / topk_vals.sum(dim=-1, keepdim=True).clamp_min(1e-9) if not self.training: forced = _env_optional_int("MAGI_MOE_FORCE_EXPERT") if forced is not None: if forced < 0 or forced >= self.num_experts: raise ValueError( f"MAGI_MOE_FORCE_EXPERT={forced} out of range [0, {self.num_experts})" ) n = probs.shape[0] topk_idx = torch.full( (n, self.top_k), int(forced), device=probs.device, dtype=torch.long ) topk_vals = torch.zeros((n, self.top_k), device=probs.device, dtype=probs.dtype) topk_vals[:, 0] = 1.0 if self.training: self._last_aux_loss = self._compute_aux_loss(probs, topk_idx[:, 0]) else: self._last_aux_loss = None self._last_expert_frac = None self._last_router_entropy = None # Observe-only: does not change logits, top-k, or expert compute. if env_is_true("MAGI_MOE_TRACE") and modality_dispatcher is not None: top1 = topk_idx[:, 0] mods = modality_dispatcher.permuted_modality_mapping if mods is not None and mods.shape[0] == top1.shape[0]: idx = top1.long() * 3 + mods.long().clamp(0, 2) self._last_mod_counts = torch.bincount(idx, minlength=self.num_experts * 3).reshape( self.num_experts, 3 ).detach() else: self._last_mod_counts = None else: self._last_mod_counts = None _sp.stop(_router_tok) _experts_tok = _sp.start(f"layer.{self.layer_idx}.moe.experts") mode = self._grouped_mode() if mode in ("bmm", "concat"): out = self._experts_grouped(x_norm, topk_idx, topk_vals, mode) _sp.stop(_experts_tok) return out out = torch.zeros_like(x_norm) # Loop experts; gather tokens that selected this expert in any of top-k slots. for expert_id, expert in enumerate(self.experts): # [N, K] -> which slots match this expert match = topk_idx == expert_id if not match.any(): continue token_gate = (topk_vals * match.to(topk_vals.dtype)).sum(dim=-1) # [N] token_mask = token_gate > 0 if not token_mask.any(): continue idx = token_mask.nonzero(as_tuple=False).squeeze(-1) x_expert = x_norm.index_select(0, idx) if self._mid_observer is None: y = expert(x_expert) else: if modality_dispatcher is None: raise RuntimeError("mid observer requires modality_dispatcher") y, z = expert.forward_with_intermediate(x_expert) modalities = modality_dispatcher.permuted_modality_mapping.index_select(0, idx) gate_obs = token_gate.index_select(0, idx) self._mid_observer( expert_id, modalities.detach(), gate_obs.detach(), z.detach(), y.detach(), ) gate = token_gate.index_select(0, idx).unsqueeze(-1).to(y.dtype) out.index_add_(0, idx, y * gate) _sp.stop(_experts_tok) return out def extra_repr(self) -> str: return ( f"layer_idx={self.layer_idx}, num_experts={self.num_experts}, top_k={self.top_k}, " f"router_noise={self.router_noise}, num_modality={self.pre_norm.num_modality}" ) # --- grouped-GEMM expert execution (probe) --------------------------- def _grouped_mode(self) -> str: """Return 'off' / 'bmm' / 'concat'. Groups the per-expert python-loop GEMMs into one batched GEMM when experts share identical shapes and routing is top-1. Controlled by env ``MAGI_MOE_GROUPED`` (off|bmm|concat). An instance override ``_moe_grouped_override`` wins for tests/benchmarks. Default (``auto``) is: use ``bmm`` during inference when it applies, and always fall back to the original loop for training / collector (``_mid_observer``) / non-uniform experts. Set ``off`` to disable. """ if self.training or self._mid_observer is not None: return "off" if self.top_k != 1 or len(self.experts) < 2: return "off" up = [e.up_gate_proj.out_features for e in self.experts] dn = [e.down_proj.in_features for e in self.experts] if len(set(up)) != 1 or len(set(dn)) != 1: return "off" override = getattr(self, "_moe_grouped_override", None) mode = override if override is not None else os.getenv("MAGI_MOE_GROUPED", "auto").strip().lower() if mode in ("0", "off", "false"): return "off" if mode in ("2", "concat"): return "concat" if mode in ("1", "bmm", "auto"): return "bmm" return "bmm" def _experts_grouped(self, x_norm: torch.Tensor, topk_idx: torch.Tensor, topk_vals: torch.Tensor, mode: str) -> torch.Tensor: """Numerically-equivalent grouped-GEMM replacement of the expert loop. Both modes preserve the reference semantics: token t is transformed by exactly its top-1 expert and weighted by the routing gate. - ``bmm``: batch-GEMM over experts (no FLOP overcompute; pads each expert's tokens to the same M for a single cublas batched launch). - ``concat``: one giant GEMM over all tokens x all experts (max kernel size, E x FLOP overcompute; non-selected slices are zeroed). """ e0 = self.experts[0] ids = topk_idx[:, 0] g = topk_vals[:, 0] n_tok = x_norm.shape[0] ecount = self.num_experts dtype = x_norm.dtype def _bmm(): order = ids.argsort(stable=True) ids_s = ids[order] x_s = x_norm[order] g_s = g.to(x_norm.dtype)[order] counts = torch.bincount(ids_s, minlength=ecount) offsets = torch.cumsum(counts, 0) - counts col = torch.arange(n_tok, device=x_norm.device) - offsets[ids_s] m_max = int(counts.max().item()) if m_max == 0: return torch.zeros_like(x_norm) x_pad = x_norm.new_zeros(ecount, m_max, x_norm.shape[1]) g_pad = x_norm.new_zeros(ecount, m_max) x_pad[ids_s, col] = x_s g_pad[ids_s, col] = g_s w_up = torch.stack([e.up_gate_proj.weight.to(dtype) for e in self.experts], dim=0) z = torch.bmm(x_pad, w_up.transpose(-2, -1)).to(dtype) z = e0.activation_func(z) if z.dtype != dtype: z = z.to(dtype) w_dn = torch.stack([e.down_proj.weight.to(dtype) for e in self.experts], dim=0) y = torch.bmm(z, w_dn.transpose(-2, -1)).to(dtype) y = y * g_pad.unsqueeze(-1) y_s = y[ids_s, col] y_orig = y_s.new_empty(n_tok, y_s.shape[1]) y_orig[order] = y_s return y_orig def _concat(): dd = e0.down_proj.in_features w_up = torch.cat([e.up_gate_proj.weight.to(dtype) for e in self.experts], dim=0) # [E*du, H] z = (x_norm @ w_up.t()).to(dtype) # [N, E*du] z = e0.activation_func(z) # [N, E*dd] if z.dtype != dtype: z = z.to(dtype) # owned = per-token slice of its expert (block ids*dd .. ids*dd+dd) owned = z.view(n_tok, ecount, dd).gather( 1, ids.view(n_tok, 1, 1).expand(n_tok, 1, dd) ).squeeze(1) # [N, dd] col_idx = (ids * dd).unsqueeze(1) + torch.arange(dd, device=x_norm.device).unsqueeze(0) z_masked = z.new_zeros(n_tok, ecount * dd) z_masked.scatter_(1, col_idx, owned) # zero other experts w_dn = torch.cat([e.down_proj.weight.to(dtype).t() for e in self.experts], dim=0) # [E*dd, H] y = (z_masked @ w_dn).to(dtype) * g.to(dtype).unsqueeze(-1) # [N, H] return y if mode == "concat": return _concat() return _bmm() def _as_timestep_float(timestep) -> float: if torch.is_tensor(timestep): return float(timestep.detach().float().reshape(-1)[0].item()) return float(timestep) def set_sparse_moe_diffusion_timestep( module: torch.nn.Module, timestep, step_index: Optional[int] = None, num_steps: Optional[int] = None, ) -> None: """Write the current diffusion step onto every SparseMoEMLP. Pruning collectors read these fields; inference is unchanged when unused. """ t = _as_timestep_float(timestep) idx = None if step_index is None else int(step_index) n = None if num_steps is None else int(num_steps) for m in module.modules(): if isinstance(m, SparseMoEMLP): m.diffusion_timestep = t m.diffusion_step_index = idx m.diffusion_num_steps = n def collect_moe_aux_loss(module: torch.nn.Module) -> Optional[torch.Tensor]: """Average aux loss over all SparseMoEMLP submodules that ran in the last forward.""" losses: List[torch.Tensor] = [] for m in module.modules(): if isinstance(m, SparseMoEMLP) and m._last_aux_loss is not None: losses.append(m._last_aux_loss) m._last_aux_loss = None if not losses: return None return torch.stack([x.float() for x in losses]).mean() def collect_moe_routing_stats(module: torch.nn.Module) -> Optional[dict]: """ Aggregate token-routing stats across SparseMoEMLP layers from the last forward. Returns mean expert fractions, mean router entropy, and load-imbalance ratio. """ fracs: List[torch.Tensor] = [] ents: List[torch.Tensor] = [] for m in module.modules(): if isinstance(m, SparseMoEMLP) and m._last_expert_frac is not None: fracs.append(m._last_expert_frac.float()) if m._last_router_entropy is not None: ents.append(m._last_router_entropy.float()) m._last_expert_frac = None m._last_router_entropy = None if not fracs: return None mean_frac = torch.stack(fracs, dim=0).mean(dim=0) mean_ent = torch.stack(ents).mean() if ents else mean_frac.new_zeros(()) eps = 1e-9 imbalance = float((mean_frac.max() / mean_frac.clamp_min(eps).min()).item()) return { "expert_frac": [float(x) for x in mean_frac.detach().cpu().tolist()], "router_entropy": float(mean_ent.detach().item()), "load_imbalance": imbalance, } def remap_dense_mlp_state_to_sparse_moe( dense_state: dict, *, sparse_moe_layers: List[int], num_experts: int, expert_noise_std: float = 1e-3, mm_layers: Optional[List[int]] = None, packed_modality_experts: int = 3, convert_mm_mlp: bool = False, ) -> dict: """ Remap dense / modality-hard MLP keys into SparseMoEMLP expert[0..E-1]. Dense middle MLP: block.layers.{i}.mlp.up_gate_proj.weight / down_proj.weight -> copy into every expert; expert 1..E-1 get Gaussian noise mm_layers NativeMoELinear (packed ``out * 3, in``): chunk along out dim into 3 modality experts, copy 1:1 into SparseMoEMLP experts (video/audio/text weights become the 3 routed experts; router is left absent so it random-inits). Extra experts beyond 3 copy expert 0 plus noise. pre_norm.weight is kept (shape already matches SparseMoEMLP when mm_layers use num_modality=3). router.weight is left absent. """ target_layers = token_sparse_moe_mlp_layers( sparse_moe_layers, mm_layers, convert_mm_mlp=convert_mm_mlp ) mm_set = set(int(x) for x in (mm_layers or [])) n_packed = int(packed_modality_experts) out = dict(dense_state) for layer_idx in target_layers: prefix = f"block.layers.{layer_idx}.mlp." if f"{prefix}experts.0.up_gate_proj.weight" in dense_state: continue up_key = prefix + "up_gate_proj.weight" down_key = prefix + "down_proj.weight" if up_key not in dense_state or down_key not in dense_state: continue up_w = dense_state[up_key] down_w = dense_state[down_key] out.pop(up_key, None) out.pop(down_key, None) if layer_idx in mm_set: if up_w.shape[0] % n_packed != 0 or down_w.shape[0] % n_packed != 0: raise ValueError( f"mm layer {layer_idx} MLP is not packed into {n_packed} experts: " f"up={tuple(up_w.shape)} down={tuple(down_w.shape)}" ) up_chunks = up_w.chunk(n_packed, dim=0) down_chunks = down_w.chunk(n_packed, dim=0) for e in range(num_experts): src = min(e, n_packed - 1) up_e = up_chunks[src].clone() down_e = down_chunks[src].clone() if e >= n_packed and expert_noise_std > 0: up_e = up_e + torch.randn_like(up_e) * expert_noise_std down_e = down_e + torch.randn_like(down_e) * expert_noise_std out[f"{prefix}experts.{e}.up_gate_proj.weight"] = up_e out[f"{prefix}experts.{e}.down_proj.weight"] = down_e else: for e in range(num_experts): up_e = up_w.clone() down_e = down_w.clone() if e > 0 and expert_noise_std > 0: up_e = up_e + torch.randn_like(up_e) * expert_noise_std down_e = down_e + torch.randn_like(down_e) * expert_noise_std out[f"{prefix}experts.{e}.up_gate_proj.weight"] = up_e out[f"{prefix}experts.{e}.down_proj.weight"] = down_e return out def load_dense_weights_into_sparse_moe_model( model: torch.nn.Module, dense_state: dict, *, sparse_moe_layers: List[int], num_experts: int, expert_noise_std: float = 1e-3, mm_layers: Optional[List[int]] = None, convert_mm_mlp: bool = False, ) -> Tuple[list, list]: """Load a dense DiT state_dict into a Sparse-MoE DiT (strict=False after remap). Keys whose tensor shapes do not match the target model (e.g. audio Linear 64→128 when adapting OmniVAE) are dropped before load_state_dict; PyTorch treats size mismatches as hard errors even with strict=False. """ remapped = remap_dense_mlp_state_to_sparse_moe( dense_state, sparse_moe_layers=sparse_moe_layers, num_experts=num_experts, expert_noise_std=expert_noise_std, mm_layers=mm_layers, convert_mm_mlp=convert_mm_mlp, ) model_sd = model.state_dict() filtered = {} shape_skipped = [] for key, value in remapped.items(): if key in model_sd and hasattr(value, "shape") and tuple(value.shape) != tuple(model_sd[key].shape): shape_skipped.append(key) continue filtered[key] = value if shape_skipped: # Attached for callers that want to log OmniVAE / channel-change skips. model._omnivae_shape_skipped_keys = list(shape_skipped) # type: ignore[attr-defined] missing, unexpected = model.load_state_dict(filtered, strict=False) unexpected = list(unexpected) + [f"{k} (shape_mismatch_skipped)" for k in shape_skipped] return list(missing), unexpected @dataclass class AdapterConfig: hidden_size: int num_attention_heads: int text_in_channels: int video_in_channels: int audio_in_channels: int params_dtype: torch.dtype class Adapter(torch.nn.Module): config: AdapterConfig def __init__(self, config: AdapterConfig): super().__init__() self.config = config self.video_embedder = nn.Linear(config.video_in_channels, config.hidden_size, bias=True, dtype=torch.float32) self.text_embedder = nn.Linear(config.text_in_channels, config.hidden_size, bias=True, dtype=torch.float32) self.audio_embedder = nn.Linear(config.audio_in_channels, config.hidden_size, bias=True, dtype=torch.float32) self.rope = ElementWiseFourierEmbed(config.hidden_size // config.num_attention_heads, in_pixels=False, learnable=False) def forward( self, x: torch.Tensor, coords_mapping: torch.Tensor, video_mask: torch.Tensor, audio_mask: torch.Tensor, text_mask: torch.Tensor, ): rope = self.rope(coords_mapping) if torch.is_autocast_enabled() and x.is_cuda: out_dtype = torch.get_autocast_gpu_dtype() else: out_dtype = x.dtype output_x = torch.zeros(x.shape[0], self.config.hidden_size, device=x.device, dtype=out_dtype) output_x[text_mask] = self.text_embedder(x[text_mask, : self.config.text_in_channels]).to(output_x.dtype) output_x[audio_mask] = self.audio_embedder(x[audio_mask, : self.config.audio_in_channels]).to(output_x.dtype) output_x[video_mask] = self.video_embedder(x[video_mask, : self.config.video_in_channels]).to(output_x.dtype) return output_x, rope class TransFormerLayer(torch.nn.Module): def __init__(self, config: Any, layer_idx: int): super().__init__() self.layer_idx = int(layer_idx) sparse_moe_layers = [int(x) for x in getattr(config, "sparse_moe_layers", []) or []] mm_layers = [int(x) for x in getattr(config, "mm_layers", []) or []] convert_mm_mlp = bool(getattr(config, "token_sparse_moe_convert_mm_mlp", False)) mlp_moe_layers = set( token_sparse_moe_mlp_layers(sparse_moe_layers, mm_layers, convert_mm_mlp=convert_mm_mlp) ) # Token-level SparseMoEMLP on listed middle MLPs, and mm-layer MLPs when convert_mm_mlp. # Attention on mm_layers stays per-modality NativeMoELinear. self.use_sparse_moe = layer_idx in mlp_moe_layers num_modality = 3 if layer_idx in set(mm_layers) else 1 use_local_attn = layer_idx in config.local_attn_layers self.post_norm = layer_idx in config.post_norm_layers attention_config = AttentionConfig( hidden_size=config.hidden_size, num_heads_q=config.num_heads_q, num_heads_kv=config.num_heads_kv, head_dim=config.head_dim, params_dtype=config.params_dtype, checkpoint_qk_layernorm_rope=config.checkpoint_qk_layernorm_rope, num_modality=num_modality, num_layers=config.num_layers, use_local_attn=use_local_attn, enable_attn_gating=config.enable_attn_gating, ) self.attention: Attention = Attention(attention_config) activation_type = MLPActivationType.GELU7 if layer_idx in config.gelu7_layers else MLPActivationType.SWIGLU7 if activation_type == MLPActivationType.SWIGLU7: gated_act = True intermediate_size = int(config.hidden_size * 4 * 2 / 3) // 4 * 4 else: gated_act = False intermediate_size = config.hidden_size * 4 if self.use_sparse_moe: # Keep the per-expert MLP structure aligned with the original dense MLP, # while routing behavior is controlled by the training config. sparse_cfg = SparseMoEConfig( hidden_size=config.hidden_size, intermediate_size=intermediate_size, activation_type=activation_type, params_dtype=config.params_dtype, num_experts=int(getattr(config, "num_routed_experts", 3)), top_k=int(getattr(config, "moe_top_k", 1)), gated_act=gated_act, router_noise=float(getattr(config, "moe_router_noise", 0.0)), num_modality=num_modality, layer_idx=int(layer_idx), ) self.mlp = SparseMoEMLP(sparse_cfg) else: mlp_config = MLPConfig( hidden_size=config.hidden_size, intermediate_size=intermediate_size, activation_type=activation_type, params_dtype=config.params_dtype, num_modality=num_modality, num_layers=config.num_layers, gated_act=gated_act, ) self.mlp = MLP(mlp_config) if self.post_norm: self.attn_post_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=num_modality) self.mlp_post_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=num_modality) def forward( self, hidden_states: torch.Tensor, rope: torch.Tensor, permute_mapping: torch.Tensor, inv_permute_mapping: torch.Tensor, varlen_handler: VarlenHandler, local_attn_handler: FFAHandler, modality_dispatcher: ModalityDispatcher, cp_split_sizes: List[int], ) -> torch.Tensor: _sp = stage_profile() _attn_tok = _sp.start(f"layer.{self.layer_idx}.attn") attn_out = self.attention( hidden_states, rope, permute_mapping, inv_permute_mapping, varlen_handler, local_attn_handler, modality_dispatcher, cp_split_sizes, ) if self.post_norm: attn_out = self.attn_post_norm(attn_out, modality_dispatcher=modality_dispatcher) _sp.stop(_attn_tok) hidden_states = hidden_states + attn_out _mlp_tok = _sp.start(f"layer.{self.layer_idx}.mlp.{'moe' if self.use_sparse_moe else 'dense'}") mlp_out = self.mlp(hidden_states, modality_dispatcher) if self.post_norm: mlp_out = self.mlp_post_norm(mlp_out, modality_dispatcher=modality_dispatcher) _sp.stop(_mlp_tok) hidden_states = hidden_states + mlp_out return hidden_states is_base_model = True def config_patch(compile_config: CompileConfig) -> CompileConfig: global is_base_model if is_base_model: is_base_model = False else: # Fully offload SR model for memory-constrained GPU compile_config.offload_config.gpu_resident_weight_ratio = 0.0 return compile_config @magi_compile(config_patch=config_patch) class TransformerBlock(torch.nn.Module): def __init__(self, model_config: Any): super().__init__() self.layers: list[TransFormerLayer] = nn.ModuleList() for layer_idx in range(model_config.num_layers): self.layers.append(TransFormerLayer(model_config, layer_idx)) def forward( self, x: torch.Tensor, rope: torch.Tensor, permute_mapping: torch.Tensor, inv_permute_mapping: torch.Tensor, varlen_handler: VarlenHandler, local_attn_handler: FFAHandler, modality_dispatcher: ModalityDispatcher, cp_split_sizes: List[int], ) -> torch.Tensor: import torch.utils.checkpoint as checkpoint _sp = stage_profile() _blk_tok = _sp.start("block.total") use_checkpointing = torch.is_grad_enabled() and os.getenv("MAGI_GRADIENT_CHECKPOINTING", "1") == "1" def create_custom_forward(module): def custom_forward(x_in, rope_in, permute_mapping_in, inv_permute_mapping_in): return module( x_in, rope_in, permute_mapping_in, inv_permute_mapping_in, varlen_handler, local_attn_handler, modality_dispatcher, cp_split_sizes, ) return custom_forward for _, layer in enumerate(self.layers): if use_checkpointing: x = checkpoint.checkpoint( create_custom_forward(layer), x, rope, permute_mapping, inv_permute_mapping, use_reentrant=False ) else: x = layer( x, rope, permute_mapping, inv_permute_mapping, varlen_handler, local_attn_handler, modality_dispatcher, cp_split_sizes, ) _sp.stop(_blk_tok) return x @dataclass class TransformerConfig: hidden_size: int video_in_channels: int audio_in_channels: int text_in_channels: int params_dtype: torch.dtype post_process_dtype: torch.dtype class DiTModel(torch.nn.Module): config: TransformerConfig def __init__(self, model_config: Any): super().__init__() self.config = TransformerConfig( hidden_size=model_config.hidden_size, video_in_channels=model_config.video_in_channels, audio_in_channels=model_config.audio_in_channels, text_in_channels=model_config.text_in_channels, params_dtype=model_config.params_dtype, post_process_dtype=torch.float32, ) adapter_config = AdapterConfig( hidden_size=model_config.hidden_size, num_attention_heads=model_config.num_heads_q, text_in_channels=model_config.text_in_channels, video_in_channels=model_config.video_in_channels, audio_in_channels=model_config.audio_in_channels, params_dtype=torch.float32, ) self.adapter: Adapter = Adapter(adapter_config) self.block: TransformerBlock = TransformerBlock(model_config=model_config) self.final_norm_video = MultiModalityRMSNorm(self.config.hidden_size) self.final_norm_audio = MultiModalityRMSNorm(self.config.hidden_size) self.final_linear_video = nn.Linear( self.config.hidden_size, self.config.video_in_channels, bias=False, dtype=torch.float32 ) self.final_linear_audio = nn.Linear( self.config.hidden_size, self.config.audio_in_channels, bias=False, dtype=torch.float32 ) def set_diffusion_timestep(self, timestep, step_index=None, num_steps=None) -> None: set_sparse_moe_diffusion_timestep(self, timestep, step_index=step_index, num_steps=num_steps) def forward( self, x: torch.Tensor, coords_mapping: torch.Tensor, modality_mapping: torch.Tensor, varlen_handler: VarlenHandler, local_attn_handler: FFAHandler, ): _sp = stage_profile() _model_tok = _sp.start("model.total") x = ulysses_scheduler().dispatch(x) coords_mapping = ulysses_scheduler().dispatch(coords_mapping) modality_mapping = ulysses_scheduler().dispatch(modality_mapping) cp_split_sizes = ulysses_scheduler().cp_split_sizes modality_dispatcher = ModalityDispatcher(modality_mapping, 3) permute_mapping, inv_permute_mapping = modality_dispatcher.permute_mapping, modality_dispatcher.inv_permute_mapping video_mask = modality_mapping == Modality.VIDEO audio_mask = modality_mapping == Modality.AUDIO text_mask = modality_mapping == Modality.TEXT x, rope = self.adapter(x, coords_mapping, video_mask, audio_mask, text_mask) x = x.to(self.config.params_dtype) x = ModalityDispatcher.permute(x, permute_mapping) x = self.block( x, rope, permute_mapping=permute_mapping, inv_permute_mapping=inv_permute_mapping, varlen_handler=varlen_handler, local_attn_handler=local_attn_handler, modality_dispatcher=modality_dispatcher, cp_split_sizes=cp_split_sizes, ) x = ModalityDispatcher.inv_permute(x, inv_permute_mapping) x_video = x[video_mask].to(self.final_norm_video.weight.dtype) x_video = self.final_norm_video(x_video) x_video = self.final_linear_video(x_video) x_audio = x[audio_mask].to(self.final_norm_audio.weight.dtype) x_audio = self.final_norm_audio(x_audio) x_audio = self.final_linear_audio(x_audio) x_out = torch.zeros( x.shape[0], max(self.config.video_in_channels, self.config.audio_in_channels), device=x.device, dtype=x.dtype ) x_out[video_mask, : self.config.video_in_channels] = x_video x_out[audio_mask, : self.config.audio_in_channels] = x_audio x_out = ulysses_scheduler().undispatch(x_out) _sp.stop(_model_tok) return x_out