"""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)