mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 14:57:09 +08:00
d8f7a2e898
- serve/: the minimal OpenAI-compatible server of deepseek-ai/DeepInfer#147, with a Neuron worker setup (basic mode env, per-rank NeuronCore, fixed KV pool via --num-kv-blocks, no MTP lookahead) and a /v1/completions endpoint (prompt text or token ids, no chat template): what vllm bench serve --backend openai drives, as vllm-neuron#71's run_bench.py does. - examples/neuron/bench_neuron.py: the same TTFT/TPOT matrix offline. - static MoE: one packed fp32 all-gather instead of three (each 64-rank collective is ~2 ms of fixed latency a layer); a capacity-grouped expert path for prefill chunks (DEEPINFER_NEURON_MOE_CAPACITY, dense rerun on overflow; not yet validated on device). - Sub-region timings at split() boundaries under DEEPINFER_NEURON_REGION_TIMING. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
235 lines
12 KiB
Python
235 lines
12 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
|
|
|
|
|
|
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,
|
|
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)]
|
|
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 _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)
|
|
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))
|
|
# 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)
|
|
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()
|
|
regions.split()
|
|
out = _reduce_scatter(acc, group, size, rank)[:num_tokens]
|
|
regions.split()
|
|
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()
|
|
if capacity:
|
|
return out.to(torch.bfloat16), overflow if overflow is not None else torch.zeros((), dtype=torch.bool, device=out.device)
|
|
return out.to(torch.bfloat16)
|