Neuron DSpark: 6-row steps keep only coarse region splits (8 graphs a layer instead of 12)

neuronx-cc segfaults (LegalizeSundaAccess) when the 6-row projections are merged
with the indexer, so that boundary stays; the other fine splits are prompt-only.
bench28 (512 in, 128 out, TP=64, nextn=5): TPOT 277.4 -> 272.6 ms, 3.85 tokens
a step.

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

deepinfer/model_executor/models/decoder_layer.py

@@ -170,7 +170,7 @@ class DeepseekDecoderLayer(HCLayerMixin):
         hidden_states, hidden_states_bf16, residual = self._attn_half(
             hidden_states, kv_cache, input_metadata, residual, engram_input, cross_layer_runtime, dflash_input
         )
-        neuron_regions.split()
+        neuron_regions.split(coarse=True)
         return self.mlp(hidden_states, hidden_states_bf16, input_metadata), residual
 
     def _attn_half(

deepinfer/model_executor/models/v4_modules/attn.py

@@ -896,7 +896,8 @@ class V4Attention(V4AttentionBase):
             neuron_regions.split()
             q_b, index_q_b = self._project_q_b(q_a, input_metadata.positions)
             del q_a
-            neuron_regions.split()
+            # Coarse: the projections merged with the indexer crash neuronx-cc at 6 rows (LegalizeSundaAccess).
+            neuron_regions.split(coarse=True)
 
             kv_ca_index = self._compute_kv_ca_index(
                 wait_projected_kv,
@@ -955,7 +956,7 @@ class V4Attention(V4AttentionBase):
             else:
                 kv_swa_index = input_metadata.kv_swa_index  # [total_q, window_size]
 
-        neuron_regions.split()
+        neuron_regions.split(coarse=True)
         if input_metadata.is_prompt:
             context = self._prefill_attention(
                 q_b,
@@ -976,7 +977,7 @@ class V4Attention(V4AttentionBase):
                 input_metadata,
             )
         del q_b, kv_ca_index
-        neuron_regions.split()
+        neuron_regions.split(coarse=True)
 
         return self._wv_and_o_proj(context, input_metadata.positions)
 

deepinfer/neuron_ops/regions.py

@@ -119,20 +119,24 @@ def _boundary() -> None:
 
 #: Rows of the step being run (set by NeuronGraphRunner.replay): a prefill chunk has more than one.
 _STEP_ROWS = [1]
-#: Steps up to this many rows keep whole-layer graphs (DEEPINFER_NEURON_SPLIT_MIN_ROWS). Whole 6-row layers (DSpark
-#: verification) crash neuronx-cc on the full model, so they are split like prompt chunks by default.
-_SPLIT_MIN_ROWS = int(os.environ.get("DEEPINFER_NEURON_SPLIT_MIN_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"))
 
 
 def set_step_rows(rows: int) -> None:
     _STEP_ROWS[0] = rows
 
 
-def split() -> None:
+def split(coarse: bool = False) -> None:
     """Cut the current region graph here (only while compiling; free in eager). Always for prefill chunks under
     DEEPINFER_NEURON_REGION_SPLIT_PREFILL: neuronx-cc runs the split 128-row attention ~2x faster (62 vs 122 ms a
-    layer, issue #26), while decode's 1-row graphs are faster whole."""
+    layer, issue #26), while decode's 1-row graphs are faster whole. Multi-row steps up to _SPLIT_MIN_ROWS rows cut
+    only at the ``coarse`` boundaries."""
     if not torch.compiler.is_compiling():
         return
-    if envs.DEEPINFER_NEURON_REGION_SPLIT or (envs.DEEPINFER_NEURON_REGION_SPLIT_PREFILL and _STEP_ROWS[0] > _SPLIT_MIN_ROWS):
+    rows = _STEP_ROWS[0]
+    multi_row = envs.DEEPINFER_NEURON_REGION_SPLIT_PREFILL and rows > 1
+    if envs.DEEPINFER_NEURON_REGION_SPLIT or (multi_row and (coarse or rows > _SPLIT_MIN_ROWS)):
         _boundary()

deepinfer/neuron_ops/static_moe.py

@@ -251,7 +251,7 @@ def _nki_moe_forward(
         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()
+    regions.split(coarse=True)
     local_idx = idx - moe.neuron_rank_tensor() * num_local
     overflow = None
     if capacity and capacity < x.shape[0]:
@@ -265,10 +265,10 @@ def _nki_moe_forward(
                 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()
+    regions.split(coarse=True)
     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()
+    regions.split(coarse=True)
     if moe.has_shared:
         shared = nki_gemm.linear_qdq(x_local, moe.shared_gate_up_proj.weight)
         shared = swiglu_forward_and_per_token_cast(