neuron: prompt chunks up to 1024 rows keep only the coarse region splits (TTFT 942 -> 813 ms, step 414 -> 406 ms, tokens identical); prefill grouped MoE without host sync (DEEPINFER_NEURON_MOE_PREFILL_CAPACITY, off: 942 -> 1000 ms)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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"
 

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:

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