coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-28 14:57:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Neuron: optional whole-layer region graph (DEEPINFER_NEURON_REGION_LAYER)

Browse Source 
Attention half and MoE as one region graph per decoder layer: half the NEFF
launches a step (the attention lane all-reduce did not move the decode step,
447.8 ms, so per-launch overhead is the next suspect). Host tiny cp=8 matches.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-24 09:51:01 +08:00
parent e5663a0a1d
commit 3a8ef041db
2 changed files with 16 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
deepinfer/environ.py
View File

@ -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).

14
deepinfer/model_executor/models/decoder_layer.py
View File

@ -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,