diff --git a/deepinfer/model_executor/model_loader.py b/deepinfer/model_executor/model_loader.py index 81a4997..4ebbc5d 100644 --- a/deepinfer/model_executor/model_loader.py +++ b/deepinfer/model_executor/model_loader.py @@ -85,6 +85,9 @@ def get_model( from deepinfer.neuron_ops import nki_gemm logger.info(f"Neuron: {nki_gemm.prepare_model(model)} weights take the NKI block-quant GEMMs") + from deepinfer.neuron_ops import static_moe + + logger.info(f"Neuron: {static_moe.prepare_shared_slices(model)} MoE layers compute a slice of the shared expert") from deepinfer.neuron_ops import nki_hc logger.info(f"Neuron: {nki_hc.prepare_model(model)} mHC fn matrices split into bf16 hi + lo") diff --git a/deepinfer/neuron_ops/nki_gemm.py b/deepinfer/neuron_ops/nki_gemm.py index cc3716b..7b197c1 100644 --- a/deepinfer/neuron_ops/nki_gemm.py +++ b/deepinfer/neuron_ops/nki_gemm.py @@ -120,6 +120,12 @@ def linear_qdq(x: torch.Tensor, w: torch.Tensor, group: int | None = None) -> to return _linear_kernel(name, weight.shape[0])(x, weight, scale)[: x.shape[0]] +def linear_qdq_raw(x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """``linear_qdq`` for a bare fp8 weight (uint8 bytes ``[n, k]`` + kernel scale), e.g. a rank's shared-expert + slice (static_moe.prepare_shared_slices).""" + return _linear_kernel("fp8_linear", weight.shape[0])(x, weight, scale)[: x.shape[0]] + + def linear_qdq_gated(x: torch.Tensor, w: torch.Tensor, group: int, flag: torch.Tensor) -> torch.Tensor | None: """``linear_qdq`` for fp4 expert ``group`` that the kernel skips (zeros) unless ``flag`` (int32 ``[1, 1]``) is 1. None when the gated kernel does not apply (not fp4, or columns that do not split across the two cores).""" diff --git a/deepinfer/neuron_ops/regions.py b/deepinfer/neuron_ops/regions.py index fb049a2..08507c0 100644 --- a/deepinfer/neuron_ops/regions.py +++ b/deepinfer/neuron_ops/regions.py @@ -79,6 +79,12 @@ def entry(module: torch.nn.Module, name: str) -> Callable[..., Any]: def fn(mod: torch.nn.Module, *args: Any, **kwargs: Any) -> Any: with hide_cuda(): + if _HOST_TIMING: + started = time.perf_counter() + out = compiled(mod, *args, **kwargs) + _HOST[label][0] += 1 + _HOST[label][1] += time.perf_counter() - started + return out if not envs.DEEPINFER_NEURON_REGION_TIMING: out = compiled(mod, *args, **kwargs) if _SYNC_EVERY: @@ -103,6 +109,16 @@ def entry(module: torch.nn.Module, name: str) -> Callable[..., Any]: _MARK: list = [0.0, 0, ""] +#: DEEPINFER_NEURON_HOST_TIMING: host seconds spent inside each region kind's compiled call (no synchronize: launch, +#: guard evaluation, resume frames, and any wait for the execution queue), per step (take_host_timings). +_HOST_TIMING = os.environ.get("DEEPINFER_NEURON_HOST_TIMING") == "1" +_HOST: dict[str, list] = defaultdict(lambda: [0, 0.0]) + + +def take_host_timings() -> dict[str, list]: + out = {k: list(v) for k, v in _HOST.items()} + _HOST.clear() + return out @torch._dynamo.disable diff --git a/deepinfer/neuron_ops/static_moe.py b/deepinfer/neuron_ops/static_moe.py index 4201b46..025e0f9 100644 --- a/deepinfer/neuron_ops/static_moe.py +++ b/deepinfer/neuron_ops/static_moe.py @@ -39,6 +39,11 @@ _ALIGN = 16 _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) @@ -124,6 +129,52 @@ def _reduce_scatter(t: torch.Tensor, group: torch.distributed.ProcessGroup, size 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], @@ -283,11 +334,14 @@ def _nki_moe_forward( 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) out = _all_reduce(acc, group, hier) if replicated else _reduce_scatter(acc, group, size, rank)[:num_tokens] regions.split() - if moe.has_shared: + 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 diff --git a/deepinfer/worker/neuron_graph_runner.py b/deepinfer/worker/neuron_graph_runner.py index 388c035..24212a6 100644 --- a/deepinfer/worker/neuron_graph_runner.py +++ b/deepinfer/worker/neuron_graph_runner.py @@ -199,6 +199,21 @@ class NeuronGraphRunner(CudaGraphRunner): reqs = -(-batch_size // (1 + envs.DEEPINFER_USE_NEXTN_DECODE)) nki_attention.set_table_bounds(reqs, getattr(input_metadata_device, "neuron_ctx_limit", None)) if not envs.DEEPINFER_NEURON_REGION_TIMING: + if regions._HOST_TIMING: + started = time.perf_counter() + with hide_cuda(): + out = self.compiled[batch_size](input_metadata_device) + self._host_steps = getattr(self, "_host_steps", 0) + 1 + if self._host_steps % 20 == 0 and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0): + host = regions.take_host_timings() + inside = sum(v[1] for v in host.values()) + logger.info( + f"Neuron host bs={batch_size}: forward {1e3 * (time.perf_counter() - started):.1f} ms, in regions " + f"{1e3 * inside:.1f} ms; " + ", ".join(f"{k} {v[0]}x {1e3 * v[1]:.1f} ms" for k, v in sorted(host.items(), key=lambda kv: -kv[1][1])) + ) + else: + regions.take_host_timings() + return out with hide_cuda(): return self.compiled[batch_size](input_metadata_device) torch.neuron.synchronize()