diff --git a/deepinfer/model_executor/models/decoder_layer.py b/deepinfer/model_executor/models/decoder_layer.py index cb11437..3556839 100644 --- a/deepinfer/model_executor/models/decoder_layer.py +++ b/deepinfer/model_executor/models/decoder_layer.py @@ -146,7 +146,10 @@ class DeepseekDecoderLayer(HCLayerMixin): out, overflow = neuron_regions.entry(self.mlp, "forward_capacity")( hidden_states, hidden_states_bf16, input_metadata ) - hidden_states = mlp(hidden_states, hidden_states_bf16, input_metadata) if bool(overflow) else out + overflowed = bool(overflow) + if envs.DEEPINFER_NEURON_REGION_TIMING: + neuron_regions.count("moe_capacity_overflow" if overflowed else "moe_capacity_ok") + hidden_states = mlp(hidden_states, hidden_states_bf16, input_metadata) if overflowed else out else: hidden_states = mlp( hidden_states, diff --git a/deepinfer/neuron_ops/regions.py b/deepinfer/neuron_ops/regions.py index a0c8520..cd5bc6d 100644 --- a/deepinfer/neuron_ops/regions.py +++ b/deepinfer/neuron_ops/regions.py @@ -45,6 +45,11 @@ _SYNC_EVERY = int(os.environ.get("DEEPINFER_NEURON_SYNC_EVERY", "0")) _CALLS = [0] +def count(label: str) -> None: + """Count an event in the per-step timings (DEEPINFER_NEURON_REGION_TIMING).""" + _TIMES[label].append(0.0) + + def take_timings() -> dict[str, list[float]]: """The region timings recorded since the last call, then reset.""" out = dict(_TIMES)