diff --git a/deepinfer/neuron_ops/nki_gemm.py b/deepinfer/neuron_ops/nki_gemm.py index fd214c4..62ef83d 100644 --- a/deepinfer/neuron_ops/nki_gemm.py +++ b/deepinfer/neuron_ops/nki_gemm.py @@ -144,6 +144,11 @@ def linear_qdq_gated(x: torch.Tensor, w: torch.Tensor, group: int, flag: torch.T return _kernel("fp4_linear_2core_gated", 2)(x, weight, scale, flag)[: x.shape[0]] +def prefill_bucket() -> int: + """DEEPINFER_NEURON_PREFILL_BUCKET: the row count of a prompt-chunk step (0: none).""" + return int(os.environ.get("DEEPINFER_NEURON_PREFILL_BUCKET", "0")) + + def has_kernel(w: torch.Tensor | None) -> bool: return ENABLED and w is not None and getattr(w, "_nki_scale", None) is not None and w.device.type == "neuron" diff --git a/deepinfer/neuron_ops/regions.py b/deepinfer/neuron_ops/regions.py index 08507c0..5ef0897 100644 --- a/deepinfer/neuron_ops/regions.py +++ b/deepinfer/neuron_ops/regions.py @@ -137,8 +137,9 @@ def _boundary() -> None: _STEP_ROWS = [1] #: Multi-row steps up to this many rows (DSpark verification: 1 + drafts) keep only the coarse splits #: (DEEPINFER_NEURON_SPLIT_MIN_ROWS). Whole 6-row layers crash neuronx-cc on the full model, and the 12 fine -#: splits a layer cost ~1 ms each in launches. -_SPLIT_MIN_ROWS = int(os.environ.get("DEEPINFER_NEURON_SPLIT_MIN_ROWS", "16")) +#: splits a layer cost ~1 ms each in launches. A 512-row prompt chunk too, since its attention went dense: the fine +#: splits cost TTFT 942 -> 813 ms (issue #29; they paid off when the chunk was 128 rows on the gathered kernel). +_SPLIT_MIN_ROWS = int(os.environ.get("DEEPINFER_NEURON_SPLIT_MIN_ROWS", "1024")) def set_step_rows(rows: int) -> None: diff --git a/deepinfer/neuron_ops/static_moe.py b/deepinfer/neuron_ops/static_moe.py index 3d163ad..b6c857c 100644 --- a/deepinfer/neuron_ops/static_moe.py +++ b/deepinfer/neuron_ops/static_moe.py @@ -44,6 +44,11 @@ _FLAG_FORCE = os.environ.get("DEEPINFER_NEURON_MOE_FLAG_FORCE", "") #: 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 +#: DEEPINFER_NEURON_MOE_PREFILL_CAPACITY: rows a prompt chunk's local expert computes (its first routed rows); an +#: expert routed more rows reruns over the whole chunk inside its gated kernels instead, so the result is exact and +#: no host sync is needed (DEEPINFER_NEURON_MOE_CAPACITY's overflow check is one a layer). 0 = dense (default: at 64 the +#: per-expert routing, gathers and four kernel calls cost more than the rows saved, TTFT 942 -> 1000 ms). +_PREFILL_CAPACITY = int(os.environ.get("DEEPINFER_NEURON_MOE_PREFILL_CAPACITY", "0")) def _as_rows(t: torch.Tensor, rows: int) -> torch.Tensor: return t.contiguous().view(rows, -1).view(torch.uint8) @@ -246,6 +251,50 @@ def _expert_gemm(a, b, d, e: int, fp8: bool, recipe_a, recipe_b) -> None: # noq d.copy_((a.float() @ b[e].float().transpose(0, 1)).to(d.dtype)) +def _gated_supported(moe: DeepseekMoE) -> bool: + """Both expert weights take the gated fp4 kernel (fp4, two-core column split).""" + return all( + w.dtype == torch.int8 and nki_gemm.TWO_CORE and w._nki_weight.shape[1] % 256 == 0 + for w in (moe.gate_up_proj_weight, moe.down_proj_weight) + if getattr(w, "_nki_weight", None) is not None + ) and getattr(moe.gate_up_proj_weight, "_nki_weight", None) is not None + + +def _grouped_experts_exact(moe: DeepseekMoE, x: torch.Tensor, local_idx: torch.Tensor, w: torch.Tensor, capacity: int): + """The local experts' fp32 ``[rows, hidden]`` sum, each expert over its first ``capacity`` routed rows; an expert + routed more rows contributes nothing from those and instead runs over all rows in its gated kernels (flag set on + the device). Same arithmetic per row as the dense loop, no host sync.""" + 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) + 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 = 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] + over = row_hit.sum() > capacity + keep = (~over).to(w_e.dtype) + xs = x_pad.index_select(0, table) + ws = torch.cat([w_e, w_e.new_zeros(1)]).index_select(0, table) * keep + 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()) + # Overflow: the whole chunk through the gated kernels (zeros, nearly free, when the flag is 0). + flag = over.to(torch.int32).reshape(1, 1) + gate_up = nki_gemm.linear_qdq_gated(x, moe.gate_up_proj_weight, e, flag) + 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) + acc = torch.cat([acc[:rows] + down.float(), acc[rows:]], dim=0) + return acc[:rows] + + 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 @@ -314,6 +363,14 @@ def _nki_moe_forward( overflow = None if capacity and capacity < x.shape[0]: acc, overflow = _grouped_experts(moe, x, local_idx, w, capacity) + elif ( + not capacity + and replicated + and _PREFILL_CAPACITY + and _PREFILL_CAPACITY < x.shape[0] == nki_gemm.prefill_bucket() + and _gated_supported(moe) + ): + acc = _grouped_experts_exact(moe, x, local_idx, w, _PREFILL_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