From 976f9fe4d290dd00895e403c7f743dbb71b2bd1f Mon Sep 17 00:00:00 2001 From: whn09 Date: Sun, 25 Oct 2026 00:31:45 +0800 Subject: [PATCH] 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 --- deepinfer/model_executor/models/decoder_layer.py | 2 +- .../model_executor/models/v4_modules/attn.py | 7 ++++--- deepinfer/neuron_ops/regions.py | 16 ++++++++++------ deepinfer/neuron_ops/static_moe.py | 6 +++--- 4 files changed, 18 insertions(+), 13 deletions(-) diff --git a/deepinfer/model_executor/models/decoder_layer.py b/deepinfer/model_executor/models/decoder_layer.py index e85c368..3ad5af9 100644 --- a/deepinfer/model_executor/models/decoder_layer.py +++ b/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( diff --git a/deepinfer/model_executor/models/v4_modules/attn.py b/deepinfer/model_executor/models/v4_modules/attn.py index afd9ee9..98a73a4 100644 --- a/deepinfer/model_executor/models/v4_modules/attn.py +++ b/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) diff --git a/deepinfer/neuron_ops/regions.py b/deepinfer/neuron_ops/regions.py index 72c226c..fb049a2 100644 --- a/deepinfer/neuron_ops/regions.py +++ b/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() diff --git a/deepinfer/neuron_ops/static_moe.py b/deepinfer/neuron_ops/static_moe.py index 85de881..d184ef6 100644 --- a/deepinfer/neuron_ops/static_moe.py +++ b/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(