diff --git a/deepinfer/environ.py b/deepinfer/environ.py index cbae526..9c6fc9e 100644 --- a/deepinfer/environ.py +++ b/deepinfer/environ.py @@ -114,6 +114,8 @@ class envs: DEEPINFER_NEURON_GRAPH_MODE = EnvVar.str(default="regions") # Neuron only: cut the region graphs further at the stage boundaries marked by neuron_ops.regions.split(). DEEPINFER_NEURON_REGION_SPLIT = EnvVar.bool(default=False) + # Neuron only: one region graph per decoder layer (attention half + MoE) instead of two. + DEEPINFER_NEURON_REGION_LAYER = EnvVar.bool(default=False) # Neuron only: split the region graphs of prefill chunks (steps of more than one row) regardless. DEEPINFER_NEURON_REGION_SPLIT_PREFILL = EnvVar.bool(default=True) # Neuron only: synchronize around every region graph call and log per-kind times each decode step (profiling). diff --git a/deepinfer/model_executor/models/decoder_layer.py b/deepinfer/model_executor/models/decoder_layer.py index 3556839..e85c368 100644 --- a/deepinfer/model_executor/models/decoder_layer.py +++ b/deepinfer/model_executor/models/decoder_layer.py @@ -115,6 +115,12 @@ class DeepseekDecoderLayer(HCLayerMixin): # A region graph takes the fetched Engram rows as a tensor input rather than the per-step fetch hook. engram_input = embed_hook() if use_regions and embed_hook is not None else embed_hook + if use_regions and envs.DEEPINFER_NEURON_REGION_LAYER: + # One region graph for the whole layer (attention half + MoE): half the NEFF launches a step. + return neuron_regions.entry(self, "_layer_body")( + hidden_states, kv_cache, input_metadata, residual, engram_input, cross_layer_runtime, dflash_input + ) + second_layer, second_last_layer, same_layers = get_enable_timer_layer(self.layer_idx) attn_timer_scale = len(same_layers) with ( @@ -159,6 +165,14 @@ class DeepseekDecoderLayer(HCLayerMixin): return hidden_states, residual + def _layer_body(self, hidden_states, kv_cache, input_metadata, residual, engram_input, cross_layer_runtime, dflash_input): # noqa: ANN001, ANN202 + """The layer as one Neuron region graph (DEEPINFER_NEURON_REGION_LAYER).""" + 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() + return self.mlp(hidden_states, hidden_states_bf16, input_metadata), residual + def _attn_half( self, hidden_states: torch.Tensor,