Neuron: count MoE capacity overflows under DEEPINFER_NEURON_REGION_TIMING

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-23 16:23:28 +08:00
parent d8f7a2e898
commit 7d95b7b6d6
2 changed files with 9 additions and 1 deletions

View File

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

View File

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