549 lines
22 KiB
Python
549 lines
22 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
# Copyright (c) 2025-2026 SandAI. All Rights Reserved.
|
|
|
|
"""Native multi-head MoE used by MAGI-2 Preview.
|
|
|
|
Adapted from SandAI's Apache-2.0 ``flash_mh_moe`` implementation and modified
|
|
to use vLLM's existing expert-parallel group. MAGI's routing is unusual: each
|
|
of twelve 256-wide hidden-state heads independently selects experts from its
|
|
own 256-expert bank. It is therefore not representable by vLLM's conventional
|
|
whole-token :class:`FusedMoE` primitive.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import os
|
|
from dataclasses import dataclass
|
|
from typing import Literal
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from vllm.triton_utils import HAS_TRITON, tl, triton
|
|
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
from .parallel import Magi2ParallelGroup, ep_dispatch, ep_undispatch, get_magi2_ep_group
|
|
|
|
RoutingScore = Literal["softmax", "sigmoid"]
|
|
|
|
|
|
def swiglu7_pair(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
|
|
"""Released clamped SwiGLU7 expert activation, evaluated in fp32."""
|
|
|
|
dtype = gate.dtype
|
|
gate = gate.float().clamp(max=7.0)
|
|
up = up.float().clamp(min=-7.0, max=7.0)
|
|
return (gate * torch.sigmoid(1.702 * gate) * (up + 1.0)).to(dtype)
|
|
|
|
|
|
def compute_topk_probs_and_indices(
|
|
router_logits: torch.Tensor,
|
|
top_k: int,
|
|
*,
|
|
score_func: RoutingScore = "sigmoid",
|
|
expert_bias: torch.Tensor | None = None,
|
|
route_norm: bool = True,
|
|
norm_eps: float = 1e-12,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Route independently for every ``[head, token]`` pair.
|
|
|
|
The auxiliary-free bias affects expert selection but deliberately does not
|
|
affect the returned routing probability, matching the training recipe.
|
|
"""
|
|
|
|
if router_logits.ndim != 3:
|
|
raise ValueError("router_logits must be [heads,tokens,experts]")
|
|
if not 0 < top_k <= router_logits.shape[-1]:
|
|
raise ValueError("top_k must be in [1, num_experts]")
|
|
if score_func == "sigmoid":
|
|
router_scores = torch.sigmoid(router_logits)
|
|
elif score_func == "softmax":
|
|
router_scores = torch.softmax(router_logits, dim=-1)
|
|
else:
|
|
raise ValueError(f"unsupported routing score function {score_func!r}")
|
|
selection_scores = router_scores
|
|
if expert_bias is not None:
|
|
selection_scores = selection_scores + expert_bias.view(router_logits.shape[0], 1, -1)
|
|
# Keep the reference's default sorted=True behavior. Besides defining the
|
|
# route order for ties, this also fixes the reduction order used by the
|
|
# following L1 normalization.
|
|
topk_indices = torch.topk(selection_scores, top_k, dim=-1).indices
|
|
topk_probs = router_scores.gather(-1, topk_indices)
|
|
if route_norm:
|
|
topk_probs = F.normalize(topk_probs, p=1, dim=-1, eps=norm_eps)
|
|
return topk_probs, topk_indices
|
|
|
|
|
|
def global_sort_routes(
|
|
topk_probs: torch.Tensor,
|
|
topk_indices: torch.Tensor,
|
|
num_experts: int,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""Convert per-head routes into a stable flattened-expert CSR layout."""
|
|
|
|
if topk_probs.shape != topk_indices.shape or topk_indices.ndim != 3:
|
|
raise ValueError("top-k probabilities and indices must have the same [H,S,K] shape")
|
|
heads, sequence, top_k = topk_indices.shape
|
|
device = topk_indices.device
|
|
head_offset = torch.arange(heads, device=device).view(heads, 1, 1) * num_experts
|
|
flattened_experts = (topk_indices + head_offset).reshape(-1)
|
|
flat_probs = topk_probs.reshape(-1)
|
|
flat_tokens = torch.arange(sequence, device=device).view(1, sequence, 1).expand(heads, sequence, top_k).reshape(-1)
|
|
order = flattened_experts.argsort(stable=True)
|
|
gather_ids = flat_tokens[order].to(torch.int32)
|
|
sorted_probs = flat_probs[order].float()
|
|
counts = torch.bincount(flattened_experts, minlength=heads * num_experts)
|
|
offsets = torch.zeros(heads * num_experts + 1, device=device, dtype=torch.long)
|
|
offsets[1:] = counts.cumsum(0)
|
|
return gather_ids, sorted_probs, offsets
|
|
|
|
|
|
def torch_mh_moe_forward(
|
|
x: torch.Tensor,
|
|
gather_ids: torch.Tensor,
|
|
probs: torch.Tensor,
|
|
expert_offsets: torch.Tensor,
|
|
w_gate: torch.Tensor,
|
|
w_up: torch.Tensor,
|
|
w_down: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Small-shape correctness oracle for the fused expert kernel."""
|
|
|
|
if x.ndim != 3:
|
|
raise ValueError("multi-head MoE input must be [tokens,heads,head_dim]")
|
|
output = torch.zeros_like(x)
|
|
experts_per_head = (expert_offsets.numel() - 1) // x.shape[1]
|
|
for flat_expert in range(expert_offsets.numel() - 1):
|
|
begin = int(expert_offsets[flat_expert].item())
|
|
end = int(expert_offsets[flat_expert + 1].item())
|
|
if begin == end:
|
|
continue
|
|
head = flat_expert // experts_per_head
|
|
token_ids = gather_ids[begin:end].long()
|
|
expert_input = x.index_select(0, token_ids)[:, head]
|
|
gate = expert_input @ w_gate[flat_expert]
|
|
up = expert_input @ w_up[flat_expert]
|
|
hidden = swiglu7_pair(gate, up)
|
|
expert_output = hidden @ w_down[flat_expert]
|
|
expert_output = expert_output * probs[begin:end, None].to(expert_output.dtype)
|
|
output[:, head].index_add_(0, token_ids, expert_output)
|
|
return output
|
|
|
|
|
|
# vLLM's no-Triton placeholder exposes ``tl.constexpr`` as ``None``.
|
|
_SWIGLU7_ALPHA = tl.constexpr(1.702) if HAS_TRITON else 1.702
|
|
_SWIGLU7_LIMIT = tl.constexpr(7.0) if HAS_TRITON else 7.0
|
|
_SWIGLU7_BIAS = tl.constexpr(1.0) if HAS_TRITON else 1.0
|
|
|
|
|
|
@triton.jit
|
|
def _swiglu7_kernel(gate, up, out_dtype: tl.constexpr):
|
|
gate_clamped = tl.minimum(gate, _SWIGLU7_LIMIT)
|
|
up_clamped = tl.maximum(tl.minimum(up, _SWIGLU7_LIMIT), -_SWIGLU7_LIMIT)
|
|
sigmoid = tl.sigmoid(_SWIGLU7_ALPHA * gate_clamped)
|
|
swish = gate_clamped * sigmoid
|
|
return (swish * (up_clamped + _SWIGLU7_BIAS)).to(out_dtype)
|
|
|
|
|
|
@triton.jit
|
|
def _binary_search_expert(
|
|
cumulative_tiles,
|
|
tile_id,
|
|
num_experts: tl.constexpr,
|
|
log2_num_experts: tl.constexpr,
|
|
):
|
|
lo = 0
|
|
hi = num_experts
|
|
for _ in tl.static_range(0, log2_num_experts + 1):
|
|
mid = (lo + hi + 1) // 2
|
|
below = tl.load(cumulative_tiles + mid) <= tile_id
|
|
lo = tl.where(below, mid, lo)
|
|
hi = tl.where(below, hi, mid - 1)
|
|
return lo
|
|
|
|
|
|
@triton.jit
|
|
def _mh_moe_kernel(
|
|
x_ptr,
|
|
wg_ptr,
|
|
wu_ptr,
|
|
wd_ptr,
|
|
y_ptr,
|
|
gather_ids_ptr,
|
|
probs_ptr,
|
|
expert_offsets_ptr,
|
|
cumulative_tiles_ptr,
|
|
stride_x_s,
|
|
stride_x_h,
|
|
stride_x_dh,
|
|
stride_wg_e,
|
|
stride_wg_dh,
|
|
stride_wg_de,
|
|
stride_wu_e,
|
|
stride_wu_dh,
|
|
stride_wu_de,
|
|
stride_wd_e,
|
|
stride_wd_de,
|
|
stride_wd_dh,
|
|
stride_y_s,
|
|
stride_y_h,
|
|
stride_y_dh,
|
|
d_head: tl.constexpr,
|
|
d_expert: tl.constexpr,
|
|
num_heads: tl.constexpr,
|
|
num_flat_experts: tl.constexpr,
|
|
log2_num_experts: tl.constexpr,
|
|
block_t: tl.constexpr,
|
|
block_dh: tl.constexpr,
|
|
block_de: tl.constexpr,
|
|
acc_dtype: tl.constexpr,
|
|
deterministic: tl.constexpr = False,
|
|
):
|
|
tile_id = tl.program_id(0)
|
|
total_tiles = tl.load(cumulative_tiles_ptr + num_flat_experts)
|
|
if tile_id >= total_tiles:
|
|
return
|
|
|
|
expert = _binary_search_expert(cumulative_tiles_ptr, tile_id, num_flat_experts, log2_num_experts)
|
|
expert_i64 = expert.to(tl.int64)
|
|
head = expert // (num_flat_experts // num_heads)
|
|
tile_in_expert = tile_id - tl.load(cumulative_tiles_ptr + expert)
|
|
token_start = tl.load(expert_offsets_ptr + expert) + tile_in_expert * block_t
|
|
expert_end = tl.load(expert_offsets_ptr + expert + 1)
|
|
count = tl.minimum(token_start + block_t, expert_end) - token_start
|
|
|
|
dh_block_offsets = tl.arange(0, block_dh)
|
|
de_block_offsets = tl.arange(0, block_de)
|
|
token_offsets = tl.arange(0, block_t)
|
|
dh_offsets = tl.arange(0, d_head)
|
|
|
|
token_positions = token_start + token_offsets
|
|
token_mask = token_offsets < count
|
|
# Token indices fit in int32, but multiplying a large packed-batch index by
|
|
# the hidden-width stride does not. Promote before computing element offsets.
|
|
gather_ids = tl.load(gather_ids_ptr + token_positions, mask=token_mask, other=0).to(tl.int64)
|
|
probabilities = tl.load(probs_ptr + token_positions, mask=token_mask, other=0.0)
|
|
x_base = gather_ids * stride_x_s + head * stride_x_h
|
|
output_acc = tl.zeros([block_t, d_head], dtype=acc_dtype)
|
|
|
|
for de_start in tl.range(0, d_expert, block_de):
|
|
de_offsets = de_start + de_block_offsets
|
|
gate_acc = tl.zeros([block_t, block_de], dtype=acc_dtype)
|
|
up_acc = tl.zeros([block_t, block_de], dtype=acc_dtype)
|
|
for dh_start in tl.static_range(0, d_head, block_dh):
|
|
local_dh = dh_start + dh_block_offsets
|
|
x_block = tl.load(
|
|
x_ptr + x_base[:, None] + local_dh[None, :] * stride_x_dh,
|
|
mask=token_mask[:, None],
|
|
other=0.0,
|
|
)
|
|
wg = tl.load(
|
|
wg_ptr
|
|
+ expert_i64 * stride_wg_e
|
|
+ local_dh[:, None] * stride_wg_dh
|
|
+ de_offsets[None, :] * stride_wg_de
|
|
)
|
|
wu = tl.load(
|
|
wu_ptr
|
|
+ expert_i64 * stride_wu_e
|
|
+ local_dh[:, None] * stride_wu_dh
|
|
+ de_offsets[None, :] * stride_wu_de
|
|
)
|
|
gate_acc += tl.dot(x_block, wg)
|
|
up_acc += tl.dot(x_block, wu)
|
|
hidden = _swiglu7_kernel(gate_acc, up_acc, wd_ptr.dtype.element_ty)
|
|
down = tl.load(
|
|
wd_ptr + expert_i64 * stride_wd_e + de_offsets[:, None] * stride_wd_de + dh_offsets[None, :] * stride_wd_dh
|
|
)
|
|
output_acc += tl.dot(hidden, down)
|
|
output_acc = output_acc * probabilities[:, None]
|
|
|
|
if deterministic:
|
|
output_ptrs = y_ptr + token_positions[:, None] * stride_y_s + dh_offsets[None, :] * stride_y_dh
|
|
tl.store(output_ptrs, output_acc.to(y_ptr.dtype.element_ty), mask=token_mask[:, None])
|
|
else:
|
|
output_base = gather_ids * stride_y_s + head * stride_y_h
|
|
output_ptrs = y_ptr + output_base[:, None] + dh_offsets[None, :] * stride_y_dh
|
|
tl.atomic_add(output_ptrs, output_acc.to(y_ptr.dtype.element_ty), mask=token_mask[:, None])
|
|
|
|
|
|
def _deterministic_scatter(
|
|
sorted_output: torch.Tensor,
|
|
reference: torch.Tensor,
|
|
gather_ids: torch.Tensor,
|
|
expert_offsets: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
num_flat_experts = expert_offsets.numel() - 1
|
|
experts_per_head = num_flat_experts // reference.shape[1]
|
|
expert_lengths = torch.diff(expert_offsets)
|
|
head_values = torch.arange(num_flat_experts, device=gather_ids.device) // experts_per_head
|
|
head_ids = torch.repeat_interleave(head_values, expert_lengths)
|
|
scatter_ids = gather_ids.long() * reference.shape[1] + head_ids.long()
|
|
output = torch.zeros_like(reference).view(-1, reference.shape[-1])
|
|
output.scatter_add_(0, scatter_ids[:, None].expand_as(sorted_output), sorted_output.to(output.dtype))
|
|
return output.view_as(reference)
|
|
|
|
|
|
def _select_block_config() -> tuple[int, int, int, int, int]:
|
|
"""Return the reference kernel config, capped for pre-Blackwell GPUs."""
|
|
|
|
capability = current_omni_platform.get_device_capability()
|
|
if capability is not None and capability.major >= 10: # Blackwell
|
|
return (128, 64, 32, 2, 8)
|
|
# BLOCK_T=128 needs 122,880 bytes of shared memory and is not safe on the
|
|
# qualified L20X path. This is the reference kernel's portable config.
|
|
return (64, 64, 32, 2, 4)
|
|
|
|
|
|
def triton_mh_moe_forward(
|
|
x: torch.Tensor,
|
|
gather_ids: torch.Tensor,
|
|
probs: torch.Tensor,
|
|
expert_offsets: torch.Tensor,
|
|
w_gate: torch.Tensor,
|
|
w_up: torch.Tensor,
|
|
w_down: torch.Tensor,
|
|
*,
|
|
deterministic: bool = False,
|
|
) -> torch.Tensor:
|
|
"""Fused gather/expert/scatter kernel for released MAGI dimensions."""
|
|
|
|
routed_tokens = gather_ids.numel()
|
|
if routed_tokens == 0:
|
|
return torch.zeros_like(x)
|
|
d_head, d_expert = x.shape[-1], w_down.shape[1]
|
|
block_t, block_dh, block_de, num_stages, num_warps = _select_block_config()
|
|
if d_head % block_dh or d_expert % block_de:
|
|
return torch_mh_moe_forward(x, gather_ids, probs, expert_offsets, w_gate, w_up, w_down)
|
|
|
|
if deterministic:
|
|
output = torch.empty((routed_tokens, 1, d_head), device=x.device, dtype=x.dtype)
|
|
else:
|
|
output = torch.zeros_like(x)
|
|
num_flat_experts = expert_offsets.numel() - 1
|
|
expert_tiles = (torch.diff(expert_offsets) + block_t - 1) // block_t
|
|
cumulative_tiles = torch.cat(
|
|
(torch.zeros(1, dtype=torch.int32, device=x.device), expert_tiles.cumsum(0, dtype=torch.int32))
|
|
)
|
|
# Match the reference launch bound. Empty/excess programs return after
|
|
# comparing against ``cumulative_tiles[-1]`` inside the kernel.
|
|
grid = ((routed_tokens + block_t - 1) // block_t + num_flat_experts,)
|
|
log2_experts = max(1, math.ceil(math.log2(max(num_flat_experts, 1) + 1)))
|
|
_mh_moe_kernel[grid](
|
|
x,
|
|
w_gate,
|
|
w_up,
|
|
w_down,
|
|
output,
|
|
gather_ids,
|
|
probs,
|
|
expert_offsets,
|
|
cumulative_tiles,
|
|
x.stride(0),
|
|
x.stride(1),
|
|
x.stride(2),
|
|
w_gate.stride(0),
|
|
w_gate.stride(1),
|
|
w_gate.stride(2),
|
|
w_up.stride(0),
|
|
w_up.stride(1),
|
|
w_up.stride(2),
|
|
w_down.stride(0),
|
|
w_down.stride(1),
|
|
w_down.stride(2),
|
|
output.stride(0),
|
|
output.stride(1),
|
|
output.stride(2),
|
|
d_head,
|
|
d_expert,
|
|
x.shape[1],
|
|
num_flat_experts,
|
|
log2_experts,
|
|
block_t,
|
|
block_dh,
|
|
block_de,
|
|
tl.float32,
|
|
deterministic,
|
|
num_stages=num_stages,
|
|
num_warps=num_warps,
|
|
)
|
|
if deterministic:
|
|
return _deterministic_scatter(output.view(routed_tokens, d_head), x, gather_ids, expert_offsets)
|
|
return output
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Magi2MultiHeadMoEConfig:
|
|
hidden_size: int
|
|
num_heads: int
|
|
num_experts: int
|
|
top_k: int
|
|
expert_intermediate_size: int
|
|
params_dtype: torch.dtype
|
|
score_func: RoutingScore = "sigmoid"
|
|
route_norm: bool = True
|
|
route_scale: float = 1.0
|
|
|
|
|
|
class Magi2MultiHeadMoE(nn.Module):
|
|
"""Checkpoint-compatible MAGI-2 head-routed expert layer."""
|
|
|
|
_EP_SHARDED_PARAMETER_NAMES = frozenset(
|
|
{"gate", "W_gate", "W_up", "W_down", "router.expert_bias", "router.expert_bias_ema"}
|
|
)
|
|
|
|
def __init__(
|
|
self,
|
|
config: Magi2MultiHeadMoEConfig,
|
|
*,
|
|
ep_group: Magi2ParallelGroup | None = None,
|
|
) -> None:
|
|
super().__init__()
|
|
if config.hidden_size % config.num_heads:
|
|
raise ValueError("hidden_size must be divisible by the number of MoE heads")
|
|
self.config = config
|
|
self.num_heads = config.num_heads
|
|
self.num_experts = config.num_experts
|
|
self.top_k = config.top_k
|
|
self.d_head = config.hidden_size // config.num_heads
|
|
self.d_expert = config.expert_intermediate_size
|
|
self.ep_group = ep_group or get_magi2_ep_group()
|
|
self.padded_num_heads = math.ceil(self.num_heads / self.ep_group.world_size) * self.ep_group.world_size
|
|
self.local_num_heads = self.padded_num_heads // self.ep_group.world_size
|
|
self.local_flatten_num_experts = self.local_num_heads * self.num_experts
|
|
self.ep_pad_heads = self.padded_num_heads - self.num_heads
|
|
self.local_head_start = self.ep_group.rank * self.local_num_heads
|
|
self.has_real_moe_heads = self.local_head_start < self.num_heads
|
|
|
|
self.gate = nn.Parameter(torch.empty(self.local_flatten_num_experts, self.d_head, dtype=torch.float32))
|
|
self.W_gate = nn.Parameter(
|
|
torch.empty(self.local_flatten_num_experts, self.d_head, self.d_expert, dtype=config.params_dtype)
|
|
)
|
|
self.W_up = nn.Parameter(
|
|
torch.empty(self.local_flatten_num_experts, self.d_head, self.d_expert, dtype=config.params_dtype)
|
|
)
|
|
self.W_down = nn.Parameter(
|
|
torch.empty(self.local_flatten_num_experts, self.d_expert, self.d_head, dtype=config.params_dtype)
|
|
)
|
|
self.router = nn.Module()
|
|
# Both tensors are released checkpoint entries. Non-trainable
|
|
# Parameters let the DLO mmap path bind them on a meta-constructed
|
|
# model; persistent buffers are intentionally not mmap-loaded by the
|
|
# generic backend.
|
|
self.router.expert_bias = nn.Parameter(
|
|
torch.zeros(self.local_flatten_num_experts, dtype=torch.float32),
|
|
requires_grad=False,
|
|
)
|
|
self.router.expert_bias_ema = nn.Parameter(
|
|
torch.zeros(self.local_flatten_num_experts, dtype=torch.float32),
|
|
requires_grad=False,
|
|
)
|
|
|
|
for name in self._EP_SHARDED_PARAMETER_NAMES:
|
|
target: nn.Module | Magi2MultiHeadMoE = self
|
|
parts = name.split(".")
|
|
for part in parts[:-1]:
|
|
target = getattr(target, part)
|
|
parameter = getattr(target, parts[-1])
|
|
parameter.mmap_weight_transform = self.ep_slice
|
|
|
|
def ep_slice(self, checkpoint_tensor: torch.Tensor) -> torch.Tensor:
|
|
"""Slice flattened ``(head,expert)`` checkpoint rows for this rank."""
|
|
|
|
if checkpoint_tensor.shape[0] == self.local_flatten_num_experts:
|
|
return checkpoint_tensor
|
|
start = self.local_head_start * self.num_experts
|
|
end = min(start + self.local_flatten_num_experts, checkpoint_tensor.shape[0])
|
|
if start >= checkpoint_tensor.shape[0]:
|
|
return torch.zeros(
|
|
(self.local_flatten_num_experts, *checkpoint_tensor.shape[1:]),
|
|
dtype=checkpoint_tensor.dtype,
|
|
device=checkpoint_tensor.device,
|
|
)
|
|
local = checkpoint_tensor[start:end]
|
|
if local.shape[0] < self.local_flatten_num_experts:
|
|
# Uneven EP/head partitions require materialized zero padding;
|
|
# divisible production layouts keep the mmap-backed slice above.
|
|
padding = torch.zeros(
|
|
(self.local_flatten_num_experts - local.shape[0], *local.shape[1:]),
|
|
dtype=local.dtype,
|
|
device=local.device,
|
|
)
|
|
local = torch.cat((local, padding), dim=0)
|
|
return local
|
|
|
|
def _route(self, x_heads: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
gate = self.gate.view(self.local_num_heads, self.num_experts, self.d_head).float()
|
|
logits = torch.einsum("shd,hed->hse", x_heads.float(), gate)
|
|
bias_source = (os.environ.get("MAGI2_ROUTER_BIAS_SOURCE") or "ema").strip().lower()
|
|
bias_tensor = self.router.expert_bias if bias_source == "main" else self.router.expert_bias_ema
|
|
bias = bias_tensor.view(self.local_num_heads, self.num_experts)
|
|
probs, indices = compute_topk_probs_and_indices(
|
|
logits,
|
|
self.top_k,
|
|
score_func=self.config.score_func,
|
|
expert_bias=bias,
|
|
route_norm=self.config.route_norm,
|
|
)
|
|
return probs * self.config.route_scale, indices
|
|
|
|
def _local_forward(self, x_heads: torch.Tensor) -> torch.Tensor:
|
|
probabilities, indices = self._route(x_heads)
|
|
gather_ids, sorted_probs, offsets = global_sort_routes(probabilities, indices, self.num_experts)
|
|
if x_heads.is_cuda:
|
|
return triton_mh_moe_forward(
|
|
x_heads,
|
|
gather_ids,
|
|
sorted_probs,
|
|
offsets,
|
|
self.W_gate,
|
|
self.W_up,
|
|
self.W_down,
|
|
deterministic=os.environ.get("MAGI2_DETERMINISTIC", "0") == "1",
|
|
)
|
|
return torch_mh_moe_forward(x_heads, gather_ids, sorted_probs, offsets, self.W_gate, self.W_up, self.W_down)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
if self.ep_group.world_size > 1 and self.ep_group.replicated_sequence:
|
|
# TP column-parallel ``split_linear`` already emits exactly this
|
|
# rank's contiguous MoE-head slice. Compute it once and leave it
|
|
# sharded for the row-parallel ``merge_linear``; no token dispatch
|
|
# or head all-gather belongs on the true TP path.
|
|
local_hidden_size = self.local_num_heads * self.d_head
|
|
if x.shape[-1] != local_hidden_size:
|
|
raise ValueError(f"TP-local MAGI MoE input has width {x.shape[-1]}, expected {local_hidden_size}")
|
|
local = x.view(-1, self.local_num_heads, self.d_head)
|
|
output = self._local_forward(local) if self.has_real_moe_heads else torch.zeros_like(local)
|
|
return output.reshape(-1, local_hidden_size)
|
|
|
|
x_heads = x.view(-1, self.num_heads, self.d_head)
|
|
if self.ep_pad_heads:
|
|
padding = x_heads.new_zeros((x_heads.shape[0], self.ep_pad_heads, self.d_head))
|
|
x_heads = torch.cat((x_heads, padding), dim=1)
|
|
sequence_split_sizes: list[int] | None = None
|
|
if self.ep_group.world_size > 1:
|
|
local_size = torch.tensor([x_heads.shape[0]], dtype=torch.int64, device=x_heads.device)
|
|
gathered_sizes = [torch.empty_like(local_size) for _ in range(self.ep_group.world_size)]
|
|
torch.distributed.all_gather(gathered_sizes, local_size, group=self.ep_group.group)
|
|
sequence_split_sizes = [int(size.item()) for size in gathered_sizes]
|
|
x_heads = ep_dispatch(x_heads, self.ep_group, sequence_split_sizes)
|
|
output = self._local_forward(x_heads) if self.has_real_moe_heads else torch.zeros_like(x_heads)
|
|
if self.ep_group.world_size > 1:
|
|
output = ep_undispatch(output, self.ep_group, sequence_split_sizes)
|
|
if self.ep_pad_heads:
|
|
output = output[:, : self.num_heads]
|
|
return output.reshape(-1, self.num_heads * self.d_head)
|
|
|
|
|
|
__all__ = [
|
|
"Magi2MultiHeadMoE",
|
|
"Magi2MultiHeadMoEConfig",
|
|
"compute_topk_probs_and_indices",
|
|
"global_sort_routes",
|
|
"torch_mh_moe_forward",
|
|
"triton_mh_moe_forward",
|
|
]
|