DeepInfer/deepinfer/neuron_ops/static_moe.py

"""Static-shape expert parallelism for the Neuron decode graphs (issue #12).

The dispatch -> m_grouped GEMM -> combine path sizes everything by data (how many tokens each rank receives, the
psum layout of each expert), which costs host syncs and recompiles. Here every shape is fixed by the batch size:

  1. pad the local tokens to ``max_num_decoding_tokens`` rows and all-gather tokens, routing and gate weights
     over the EP group (one packed uint8 collective);
  2. run each local expert densely over all gathered rows, the gate weight of non-routed rows being 0;
  3. reduce-scatter the fp32 sum back to the owning rank and add the shared expert.

Decode batches are small, so computing every local expert over the gathered rows costs little next to the
weight traffic, and the step traces into one graph with no data-dependent shapes.
"""

from __future__ import annotations

import os
from typing import TYPE_CHECKING

import deep_gemm
import torch
import torch.distributed._functional_collectives as funcol
from tile_kernels.quant import swiglu_forward_and_per_token_cast

from deepinfer.model_executor.layers.linear import get_fp8_cast_params, get_fp8_casting_granularity
from deepinfer.neuron_ops import nki_gemm

if TYPE_CHECKING:
    from deepinfer.model_executor.models.v4_modules.mlp import DeepseekMoE

# gloo (host mode) has no reduce-scatter.
_HOST = os.environ.get("DEEPINFER_NEURON_DEVICE", "neuron") == "cpu"
_ALIGN = 16


def _as_rows(t: torch.Tensor, rows: int) -> torch.Tensor:
    return t.contiguous().view(rows, -1).view(torch.uint8)


def _pack(parts: list[torch.Tensor], cap: int) -> tuple[torch.Tensor, list[tuple[int, int]]]:
    """Byte rows ``[cap, W]``; each part starts at a 16-byte aligned column. Rows past the tokens are 0xFF so the
    padded routing reads as -1 (int32) / NaN (fp32), i.e. selects no expert."""
    num_tokens = parts[0].shape[0]
    spans, cols, width = [], [], 0
    for part in parts:
        b = _as_rows(part, num_tokens)
        spans.append((width, b.shape[1]))
        pad = -b.shape[1] % _ALIGN
        cols.append(torch.nn.functional.pad(b, (0, pad)) if pad else b)
        width += b.shape[1] + pad
    packed = torch.cat(cols, dim=1)
    if cap > num_tokens:
        packed = torch.cat([packed, packed.new_full((cap - num_tokens, width), 0xFF)], dim=0)
    return packed, spans


def _unpack(packed: torch.Tensor, span: tuple[int, int], like: torch.Tensor) -> torch.Tensor:
    start, nbytes = span
    rows = packed.shape[0]
    return packed[:, start : start + nbytes].contiguous().view(like.dtype).view(rows, *like.shape[1:])


def _all_gather(t: torch.Tensor, group: torch.distributed.ProcessGroup, size: int) -> torch.Tensor:
    if size == 1:
        return t
    return funcol.wait_tensor(funcol.all_gather_tensor(t, 0, group))


def _reduce_scatter(t: torch.Tensor, group: torch.distributed.ProcessGroup, size: int, rank: int) -> torch.Tensor:
    if size == 1:
        return t
    if _HOST:
        n = t.shape[0] // size
        return funcol.wait_tensor(funcol.all_reduce(t, "sum", group)).narrow(0, rank * n, n)
    return funcol.wait_tensor(funcol.reduce_scatter_tensor(t, "sum", 0, group))


def static_moe_forward(
    moe: DeepseekMoE,
    hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
    topk_idx: torch.Tensor,
    topk_weights: torch.Tensor,
    hidden_states_bf16: torch.Tensor | None = None,
) -> torch.Tensor:
    if hidden_states_bf16 is not None and nki_gemm.has_kernel(moe.gate_up_proj_weight):
        return _nki_moe_forward(moe, hidden_states_bf16, topk_idx, topk_weights)
    group, size, rank = moe.ep_group, moe.ep_size, moe.ep_rank
    num_local = moe.num_local_physical_experts
    x_parts = list(hidden_states) if moe.use_fp8_act else [hidden_states]
    num_tokens = x_parts[0].shape[0]
    cap = max(moe.max_num_decoding_tokens, num_tokens)

    parts = x_parts + [topk_idx.to(torch.int32), topk_weights.to(torch.float32)]
    if size == 1:
        # Nothing to gather. Skipping the byte packing also keeps dtype bitcasts of graph inputs out of the region
        # graph, which neuronx-cc mishandles (InferNonlocalTensors, issue #25).
        g_parts = parts
    else:
        packed, spans = _pack(parts, cap)
        gathered = _all_gather(packed, group, size)  # [size * cap, W]
        g_parts = [_unpack(gathered, span, part) for span, part in zip(spans, parts)]
    g_x = tuple(g_parts[:-2]) if moe.use_fp8_act else g_parts[0]
    g_idx, g_w = g_parts[-2], g_parts[-1]
    num_rows = g_idx.shape[0]

    recipe_a = (1, get_fp8_casting_granularity())
    recipe_b = (1, 32 if moe.use_fp4_weight else get_fp8_casting_granularity())
    swiglu_kwargs = get_fp8_cast_params() if moe.use_fp8_act else {"fmt": "bf16", "num_per_channels": None}
    gate_up_w = moe._weight_and_sf(moe.gate_up_proj_weight)
    down_w = moe._weight_and_sf(moe.down_proj_weight)

    acc = torch.zeros((num_rows, moe.hidden_size), dtype=torch.float32, device=g_idx.device)
    for e in range(num_local):
        hit = g_idx == rank * num_local + e
        w_e = torch.where(hit, g_w, torch.zeros_like(g_w)).sum(-1)
        gate_up = torch.empty((num_rows, moe.gate_up_proj_weight.shape[1]), dtype=torch.bfloat16, device=g_idx.device)
        _expert_gemm(g_x, gate_up_w, gate_up, e, moe.use_fp8_act, recipe_a, recipe_b)
        down_in = swiglu_forward_and_per_token_cast(
            gate_up, **swiglu_kwargs, topk_weights=w_e, clamp_value=moe.swiglu_clamp_value
        )
        down = torch.empty((num_rows, moe.down_proj_weight.shape[1]), dtype=torch.bfloat16, device=g_idx.device)
        _expert_gemm(down_in, down_w, down, e, moe.use_fp8_act, recipe_a, recipe_b)
        acc = acc + down.float()

    out = _reduce_scatter(acc, group, size, rank)[:num_tokens]
    shared = moe.shared_forward(hidden_states)
    if shared is not None:
        out = out + shared.float()
    return out.to(torch.bfloat16)


def _expert_gemm(a, b, d, e: int, fp8: bool, recipe_a, recipe_b) -> None:  # noqa: ANN001
    if fp8:
        deep_gemm.m_grouped_fp8_fp4_gemm_nt_shared_a(a, b, d, e, recipe_a=recipe_a, recipe_b=recipe_b)
    else:
        d.copy_((a.float() @ b[e].float().transpose(0, 1)).to(d.dtype))


def _nki_moe_forward(
    moe: DeepseekMoE, x: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor
) -> torch.Tensor:
    """The same dense static MoE on bf16 activations through the NKI block-quant kernels (vllm-neuron#71 style).

    The kernels quantize activations to fp8 themselves (1x32, power-of-two scales), exactly as the fp8 path does, so
    no fp8 activation or int16 scale crosses a region graph boundary.
    """
    group, size, rank = moe.ep_group, moe.ep_size, moe.ep_rank
    num_local = moe.num_local_physical_experts
    num_tokens = x.shape[0]
    cap = max(moe.max_num_decoding_tokens, num_tokens)
    x_local = x
    # Local expert ids relative to this rank, from a rank *tensor*: the graph stays identical across ranks.
    idx = topk_idx.to(torch.int32)
    w = topk_weights.to(torch.float32)
    if size > 1:
        pad = cap - num_tokens
        x = torch.nn.functional.pad(x, (0, 0, 0, pad))
        idx = torch.nn.functional.pad(idx, (0, 0, 0, pad), value=-1)
        w = torch.nn.functional.pad(w, (0, 0, 0, pad))
        x, idx, w = (_all_gather(t, group, size) for t in (x, idx, w))
    acc = torch.zeros((x.shape[0], moe.hidden_size), dtype=torch.float32, device=x.device)
    local_idx = idx - moe.neuron_rank_tensor() * num_local
    for e in range(num_local):
        w_e = torch.where(local_idx == e, w, torch.zeros_like(w)).sum(-1)
        gate_up = nki_gemm.linear_qdq(x, moe.gate_up_proj_weight, group=e)
        act = swiglu_forward_and_per_token_cast(
            gate_up, fmt="bf16", num_per_channels=None, topk_weights=w_e, clamp_value=moe.swiglu_clamp_value
        )
        acc = acc + nki_gemm.linear_qdq(act, moe.down_proj_weight, group=e).float()
    out = _reduce_scatter(acc, group, size, rank)[:num_tokens]
    if moe.has_shared:
        shared = nki_gemm.linear_qdq(x_local, moe.shared_gate_up_proj.weight)
        shared = swiglu_forward_and_per_token_cast(
            shared, fmt="bf16", num_per_channels=None, clamp_value=moe.swiglu_clamp_value
        )
        out = out + nki_gemm.linear_qdq(shared, moe.shared_down_proj.weight).float()
    return out.to(torch.bfloat16)