mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 14:57:09 +08:00
369 lines
19 KiB
Python
369 lines
19 KiB
Python
"""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, regions
|
|
|
|
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
|
|
|
|
|
|
|
|
#: DEEPINFER_NEURON_GATED_EXPERTS: skip unrouted experts inside the fp4 kernels at decode sizes (dynamic loop on a
|
|
#: per-expert flag, blockquant.fp4_linear_2core_gated).
|
|
_GATED_EXPERTS = os.environ.get("DEEPINFER_NEURON_GATED_EXPERTS", "1") == "1"
|
|
_GATED_MAX_ROWS = 16
|
|
_FLAG_FORCE = os.environ.get("DEEPINFER_NEURON_MOE_FLAG_FORCE", "")
|
|
#: DEEPINFER_NEURON_SHARED_SLICE: each rank computes a 128-column slice of the shared expert's intermediate dimension
|
|
#: and adds it to the routed experts' partial sum before the MoE all-reduce, instead of every rank computing the whole
|
|
#: shared expert after it (35 M fp8 weights a layer, dequantized on every rank).
|
|
_SHARED_SLICE = os.environ.get("DEEPINFER_NEURON_SHARED_SLICE", "1") == "1"
|
|
_SHARED_COLS = 128
|
|
|
|
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 tokens_replicated(moe: DeepseekMoE) -> bool:
|
|
"""Every EP rank holds the same tokens: the EP group is the CP group, whose members run one replicated batch
|
|
(Neuron, issue #28). The MoE then needs no all-gather, and an all-reduce instead of a reduce-scatter."""
|
|
from deepinfer.model_executor.parallel_utils.parallel_state import PG
|
|
|
|
return moe.ep_size > 1 and PG.context_parallel_group().size() == moe.ep_size
|
|
|
|
|
|
#: Cores of one Trainium2 device (LNC=2): ranks r // 4 share a device.
|
|
_DEVICE_CORES = 4
|
|
_HIER: dict = {}
|
|
|
|
|
|
def hierarchical_groups(group: torch.distributed.ProcessGroup): # noqa: ANN201
|
|
"""(intra-device group, cross-device group) for ``group`` = the world, or None. Created once, on every rank in the
|
|
same order (call at module init, not inside a graph). Measured at 64 ranks, [1, 5120] fp32: one world all-reduce
|
|
3.37 ms, 4-core device all-reduce then a 16-device all-reduce (one core each) 2.01 ms (issue #27)."""
|
|
import torch.distributed as dist
|
|
|
|
from deepinfer.model_executor.parallel_utils.parallel_state import _split_group
|
|
|
|
world = dist.get_world_size()
|
|
if os.environ.get("DEEPINFER_NEURON_HIER_ALLREDUCE", "1") != "1" or group.size() != world or world % _DEVICE_CORES or world <= _DEVICE_CORES:
|
|
return None
|
|
if world not in _HIER:
|
|
from deepinfer import PLATFORM
|
|
|
|
backend = PLATFORM.distributed_backend
|
|
local = _split_group([list(range(r, r + _DEVICE_CORES)) for r in range(0, world, _DEVICE_CORES)], backend=backend)
|
|
cross = _split_group([list(range(c, world, _DEVICE_CORES)) for c in range(_DEVICE_CORES)], backend=backend)
|
|
_HIER[world] = (local, cross)
|
|
return _HIER[world]
|
|
|
|
|
|
def _all_reduce(t: torch.Tensor, group: torch.distributed.ProcessGroup, hier=None) -> torch.Tensor: # noqa: ANN001
|
|
if hier is not None:
|
|
local, cross = hier
|
|
t = funcol.wait_tensor(funcol.all_reduce(t, "sum", local))
|
|
return funcol.wait_tensor(funcol.all_reduce(t, "sum", cross))
|
|
return funcol.wait_tensor(funcol.all_reduce(t, "sum", 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 prepare_shared_slices(model: torch.nn.Module) -> int:
|
|
"""Register this rank's shared-expert slice on every MoE (after nki_gemm.prepare_model). ``I // 128`` ranks own a
|
|
slice; the others get zero weights of the same shape, so every rank traces the same graph and their slice adds
|
|
exact zeros. Returns the number of MoE layers sliced."""
|
|
if not _SHARED_SLICE:
|
|
return 0
|
|
count = 0
|
|
for moe in model.modules():
|
|
gu = getattr(moe, "shared_gate_up_proj", None)
|
|
dn = getattr(moe, "shared_down_proj", None)
|
|
if not getattr(moe, "has_shared", False) or gu is None or dn is None:
|
|
continue
|
|
gw, dw = gu.weight, dn.weight
|
|
if getattr(gw, "_nki_scale", None) is None or getattr(dw, "_nki_scale", None) is None or gw.dtype != torch.float8_e4m3fn:
|
|
continue
|
|
size, rank = moe.ep_size, moe.ep_rank
|
|
inter = gw.shape[0] // 2
|
|
slices = inter // _SHARED_COLS
|
|
if size == 1 or inter % _SHARED_COLS or slices > size:
|
|
continue
|
|
c, b = _SHARED_COLS, _SHARED_COLS // 32
|
|
# On the host: strided slices of device tensors are not safe eagerly on a NeuronCore (one-time, at load).
|
|
device = gw.device
|
|
g_w, g_s, d_w, d_s = (t.cpu() for t in (gw._nki_weight, gw._nki_scale, dw._nki_weight, dw._nki_scale))
|
|
if rank < slices:
|
|
lo = rank * c
|
|
gu_w = torch.cat([g_w[lo : lo + c], g_w[inter + lo : inter + lo + c]])
|
|
gu_s = torch.cat([g_s[lo // 32 : lo // 32 + b], g_s[(inter + lo) // 32 : (inter + lo) // 32 + b]])
|
|
dn_w, dn_s = d_w[:, lo : lo + c], d_s[:, lo // 32 : lo // 32 + b]
|
|
else:
|
|
gu_w = torch.zeros((2 * c, g_w.shape[1]), dtype=g_w.dtype, device=g_w.device)
|
|
gu_s = torch.full((2 * b, g_s.shape[1]), 127, dtype=g_s.dtype, device=g_s.device) # e8m0 127 = 1.0
|
|
dn_w = torch.zeros((d_w.shape[0], c), dtype=d_w.dtype, device=d_w.device)
|
|
dn_s = torch.full((d_s.shape[0], b), 127, dtype=d_s.dtype, device=d_s.device)
|
|
for name, t in (("gu_w", gu_w), ("gu_s", gu_s), ("dn_w", dn_w), ("dn_s", dn_s)):
|
|
moe.register_buffer(f"_neuron_shared_{name}", t.contiguous().to(device), persistent=False)
|
|
count += 1
|
|
return count
|
|
|
|
|
|
def _shared_slice(moe: DeepseekMoE, x: torch.Tensor) -> torch.Tensor:
|
|
gate_up = nki_gemm.linear_qdq_raw(x, moe._neuron_shared_gu_w, moe._neuron_shared_gu_s)
|
|
act = swiglu_forward_and_per_token_cast(gate_up, fmt="bf16", num_per_channels=None, clamp_value=moe.swiglu_clamp_value)
|
|
return nki_gemm.linear_qdq_raw(act, moe._neuron_shared_dn_w, moe._neuron_shared_dn_s)
|
|
|
|
|
|
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,
|
|
valid: torch.Tensor | None = None,
|
|
capacity: int = 0,
|
|
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
|
"""``valid`` ([tokens] bool, the decode metadata's valid_token_mask): padding rows of a graph batch select no
|
|
expert, so they cost nothing past the dense compute and never reach a real row's output."""
|
|
if valid is not None:
|
|
topk_idx = torch.where(valid.view(-1, 1), topk_idx, torch.full_like(topk_idx, -1))
|
|
topk_weights = torch.where(valid.view(-1, 1), topk_weights, torch.zeros_like(topk_weights))
|
|
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, capacity)
|
|
assert not capacity, "the grouped expert path needs the NKI kernels"
|
|
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)]
|
|
replicated = tokens_replicated(moe)
|
|
if size == 1 or replicated:
|
|
# 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()
|
|
|
|
hier = getattr(moe, "_neuron_hier_ar", None)
|
|
out = _all_reduce(acc, group, hier) if replicated else _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 _grouped_experts(moe: DeepseekMoE, x: torch.Tensor, local_idx: torch.Tensor, w: torch.Tensor, capacity: int):
|
|
"""The local experts over at most ``capacity`` routed rows each instead of all ``rows`` (issue #26: a prefill
|
|
chunk gathers 64 ranks x 128 rows, and each row selects only 6 of the 384 experts). Returns the fp32
|
|
``[rows, hidden]`` sum and a bool scalar that is True when some expert had more routed rows than ``capacity``
|
|
(the result is then incomplete; the caller reruns the dense path)."""
|
|
rows = x.shape[0]
|
|
x_pad = torch.cat([x, x.new_zeros((1, x.shape[1]))], dim=0) # row `rows` is the empty slot
|
|
acc = torch.zeros((rows + 1, moe.hidden_size), dtype=torch.float32, device=x.device)
|
|
row_ids = torch.arange(rows, device=x.device, dtype=torch.int32)
|
|
overflow = torch.zeros((), dtype=torch.bool, device=x.device)
|
|
for e in range(moe.num_local_physical_experts):
|
|
hit = local_idx == e
|
|
w_e = torch.where(hit, w, torch.zeros_like(w)).sum(-1)
|
|
row_hit = hit.any(dim=-1)
|
|
# Slot of each routed row (in row order); rows past the capacity and unrouted rows go to the spare slot.
|
|
slot = torch.cumsum(row_hit.to(torch.int32), dim=0) - 1
|
|
dest = torch.where(row_hit & (slot < capacity), slot, torch.full_like(slot, capacity))
|
|
table = torch.full((capacity + 1,), rows, dtype=torch.int32, device=x.device).scatter(0, dest, row_ids)[:capacity]
|
|
overflow = overflow | (row_hit.sum() > capacity)
|
|
xs = x_pad.index_select(0, table)
|
|
ws = torch.cat([w_e, w_e.new_zeros(1)]).index_select(0, table)
|
|
gate_up = nki_gemm.linear_qdq(xs, moe.gate_up_proj_weight, group=e)
|
|
act = swiglu_forward_and_per_token_cast(
|
|
gate_up, fmt="bf16", num_per_channels=None, topk_weights=ws, clamp_value=moe.swiglu_clamp_value
|
|
)
|
|
acc = acc.index_add(0, table, nki_gemm.linear_qdq(act, moe.down_proj_weight, group=e).float())
|
|
return acc[:rows], overflow
|
|
|
|
|
|
def _nki_moe_forward(
|
|
moe: DeepseekMoE,
|
|
x: torch.Tensor,
|
|
topk_idx: torch.Tensor,
|
|
topk_weights: torch.Tensor,
|
|
capacity: int = 0,
|
|
) -> torch.Tensor | tuple[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)
|
|
replicated = tokens_replicated(moe)
|
|
if size > 1 and not replicated:
|
|
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))
|
|
# One collective instead of three (each 64-rank collective costs ~2 ms of fixed latency a layer, issue #27):
|
|
# fp32 holds the bf16 activations and the expert ids (< 2**24) exactly, and needs no dtype bitcast.
|
|
topk = idx.shape[1]
|
|
packed = torch.cat([x.float(), idx.float(), w], dim=1)
|
|
packed = _all_gather(packed, group, size)
|
|
x = packed[:, : x.shape[1]].to(torch.bfloat16)
|
|
idx = packed[:, x.shape[1] : x.shape[1] + topk].to(torch.int32)
|
|
w = packed[:, x.shape[1] + topk :]
|
|
regions.split()
|
|
local_idx = idx - moe.neuron_rank_tensor() * num_local
|
|
overflow = None
|
|
if capacity and capacity < x.shape[0]:
|
|
acc, overflow = _grouped_experts(moe, x, local_idx, w, capacity)
|
|
else:
|
|
acc = torch.zeros((x.shape[0], moe.hidden_size), dtype=torch.float32, device=x.device)
|
|
# Decode steps: an expert no row routes to is skipped inside its kernels (a 6-row step routes to about one of
|
|
# a rank's 6 experts); prompt chunks route to all of them anyway.
|
|
gated = _GATED_EXPERTS and x.shape[0] <= _GATED_MAX_ROWS
|
|
for e in range(num_local):
|
|
w_e = torch.where(local_idx == e, w, torch.zeros_like(w)).sum(-1)
|
|
flag = (local_idx == e).any().to(torch.int32).reshape(1, 1) if gated else None
|
|
if gated and _FLAG_FORCE in ("0", "1"): # timing diagnostic only: wrong results
|
|
flag = torch.full((1, 1), int(_FLAG_FORCE), dtype=torch.int32, device=x.device)
|
|
gate_up = nki_gemm.linear_qdq_gated(x, moe.gate_up_proj_weight, e, flag) if gated else None
|
|
if gate_up is None:
|
|
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
|
|
)
|
|
down = nki_gemm.linear_qdq_gated(act, moe.down_proj_weight, e, flag) if gated else None
|
|
if down is None:
|
|
down = nki_gemm.linear_qdq(act, moe.down_proj_weight, group=e)
|
|
acc = acc + down.float()
|
|
sliced = moe.has_shared and getattr(moe, "_neuron_shared_gu_w", None) is not None
|
|
if sliced:
|
|
acc = acc + _shared_slice(moe, x).float()
|
|
regions.split()
|
|
hier = getattr(moe, "_neuron_hier_ar", None)
|
|
flag_reduced = None
|
|
if capacity and replicated and size > 1:
|
|
# The ranks' overflow flags ride the output's all-reduce as one extra row (a separate 64-rank all-reduce
|
|
# cost ~3.4 ms a layer): every rank then sees the sum, so all of them take the same path.
|
|
flag = (overflow if overflow is not None else torch.zeros((), dtype=torch.bool, device=acc.device)).to(torch.float32)
|
|
acc = torch.cat([acc, flag.reshape(1, 1).expand(1, acc.shape[1])], dim=0)
|
|
out = _all_reduce(acc, group, hier)
|
|
out, flag_reduced = out[:-1], out[-1, 0]
|
|
else:
|
|
out = _all_reduce(acc, group, hier) if replicated else _reduce_scatter(acc, group, size, rank)[:num_tokens]
|
|
regions.split()
|
|
if moe.has_shared and not sliced:
|
|
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()
|
|
if capacity and flag_reduced is not None:
|
|
return out.to(torch.bfloat16), flag_reduced > 0
|
|
if capacity:
|
|
flag = (overflow if overflow is not None else torch.zeros((), dtype=torch.bool, device=out.device)).to(torch.float32)
|
|
if size > 1:
|
|
# Every rank must take the same path: the dense rerun has collectives of its own (a rank-local decision
|
|
# left the others waiting in its all-gather: "Collective operation pending").
|
|
flag = funcol.wait_tensor(funcol.all_reduce(flag.reshape(1), "max", group)).reshape(())
|
|
return out.to(torch.bfloat16), flag > 0
|
|
return out.to(torch.bfloat16)
|