sim: report correct kernel latency under enable_data=False
Three coupled fixes so enable_data=False now reports byte-identical kernel-body sim latency (max(t_end) - min(t_start) over op_log) as enable_data=True, and skips setup-write sim events that were pure wall-clock overhead. Before this change, at Case 4 (Cube-SP x PE-SP) decode, S_kv=8K: enable_data=False -> latency_ns = 0 (op_log empty) enable_data=True -> latency_ns = 30.646 us (correct) After: enable_data=False -> latency_ns = 30.646 us (byte-equal) enable_data=True -> latency_ns = 30.646 us (unchanged) engine.py: always create OpLogger. Previously OpLogger was gated on enable_data=True together with MemoryStore, so op_log stayed empty when data mode was off. Decouple: MemoryStore is gated (Phase 2 DataExecutor still needs it) but OpLogger runs unconditionally, receiving memory_store=None when data mode is off. OpLogger already guards its arr.copy() snapshot paths on `if self._memory_store is not None`. pe_cpu.py: always use the ADR-0020 greenlet execution path. The if store is not None: _execute_greenlet(); else: _execute_legacy branch gated the *execution model* on data-mode presence. The legacy command-list path predates ADR-0020 and doesn't route IpcqSendCmd through PE_IPCQ (it goes to PE_SCHEDULER instead), so all 189 Case-4 ipcq_copy fabric transfers were silently no-op'd in that mode. Forcing greenlet gives identical sim behavior in both modes. KernelRunner already guards its store reads on `self._store is not None`. context.py: gate setup MemoryWriteMsg on memory_store presence. Under enable_data=False there is no MemoryStore to populate and no Phase 2 replay, so the per-shard sim events for Q/K/V deploy are pure wall-clock overhead with no effect on reported kernel latency (Yangwook's max-min formula excludes pre-kernel ops). Handle count for Case 4 at S_kv=8K drops 229 -> 37 under enable_data=False. Verified: - tests/attention/test_milestone_gqa_decode_long_ctx_4cases.py: 18/18 pass - Parity probe: enable_data=False/True both report 30.646 us kernel sim time, 1649 op_log records, 189 ipcq_copy events Known follow-up: pe_cpu._execute_legacy is now dead code but left in place; separate cleanup once we confirm no downstream caller. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -117,17 +117,13 @@ class PeCpuComponent(ComponentBase):
|
|||||||
pe_exec_start = env.now
|
pe_exec_start = env.now
|
||||||
scheduler_id = f"{self._pe_prefix}.pe_scheduler"
|
scheduler_id = f"{self._pe_prefix}.pe_scheduler"
|
||||||
|
|
||||||
# Choose execution mode: greenlet (ADR-0020) or legacy command-list
|
# ADR-0020 greenlet execution — always used so IPCQ / ring credits /
|
||||||
|
# fabric-transfer sim events fire regardless of data-mode. KernelRunner
|
||||||
|
# guards its store reads on ``self._store is not None``.
|
||||||
store = getattr(self.ctx, "memory_store", None) if self.ctx else None
|
store = getattr(self.ctx, "memory_store", None) if self.ctx else None
|
||||||
|
|
||||||
if store is not None:
|
|
||||||
composite_results = yield from self._execute_greenlet(
|
composite_results = yield from self._execute_greenlet(
|
||||||
env, kernel_fn, kernel_args, num_programs, scheduler_id, store,
|
env, kernel_fn, kernel_args, num_programs, scheduler_id, store,
|
||||||
)
|
)
|
||||||
else:
|
|
||||||
composite_results = yield from self._execute_legacy(
|
|
||||||
env, kernel_fn, kernel_args, num_programs, scheduler_id,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Record PE-internal execution time
|
# Record PE-internal execution time
|
||||||
txn.result_data["pe_exec_ns"] = env.now - pe_exec_start
|
txn.result_data["pe_exec_ns"] = env.now - pe_exec_start
|
||||||
|
|||||||
@@ -570,8 +570,14 @@ class RuntimeContext:
|
|||||||
h = self.submit(msg)
|
h = self.submit(msg)
|
||||||
self.wait(h)
|
self.wait(h)
|
||||||
|
|
||||||
# Submit MemoryWriteMsg per shard (deploy data to device)
|
# Submit MemoryWriteMsg per shard (deploy data to device). Gated on
|
||||||
if pattern is not None:
|
# memory_store presence: under enable_data=False there is no
|
||||||
|
# MemoryStore to populate and no Phase 2 DataExecutor replay, so the
|
||||||
|
# per-shard sim events are pure wall-clock overhead with no effect
|
||||||
|
# on reported kernel latency (Yangwook's max(t_end) - min(t_start)
|
||||||
|
# formula excludes ops before the kernel starts).
|
||||||
|
store = getattr(self.engine, "_memory_store", None)
|
||||||
|
if pattern is not None and store is not None:
|
||||||
for shard in handle.shards:
|
for shard in handle.shards:
|
||||||
h = self.submit(MemoryWriteMsg(
|
h = self.submit(MemoryWriteMsg(
|
||||||
correlation_id=self.correlation_id,
|
correlation_id=self.correlation_id,
|
||||||
@@ -591,8 +597,7 @@ class RuntimeContext:
|
|||||||
# VA; Phase 2 DataExecutor reads via the addresses captured in
|
# VA; Phase 2 DataExecutor reads via the addresses captured in
|
||||||
# op_log (VA for tl.load). Without this, zero-init tensors are
|
# op_log (VA for tl.load). Without this, zero-init tensors are
|
||||||
# invisible to kernels in Phase 2.
|
# invisible to kernels in Phase 2.
|
||||||
store = getattr(self.engine, "_memory_store", None)
|
if pattern == "zero" and handle.va_base:
|
||||||
if store is not None and pattern == "zero" and handle.va_base:
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
from kernbench.runtime_api.tensor import _numpy_dtype
|
from kernbench.runtime_api.tensor import _numpy_dtype
|
||||||
np_dtype = _numpy_dtype(dtype)
|
np_dtype = _numpy_dtype(dtype)
|
||||||
|
|||||||
@@ -53,13 +53,15 @@ class GraphEngine:
|
|||||||
self._events: dict[str, simpy.Event] = {}
|
self._events: dict[str, simpy.Event] = {}
|
||||||
self._counter = 0
|
self._counter = 0
|
||||||
overrides = component_overrides or {}
|
overrides = component_overrides or {}
|
||||||
# ADR-0020: optional data execution support
|
# ADR-0020: optional data execution support. OpLogger is always
|
||||||
self._op_logger = None
|
# created so op_log-based latency (max(t_end) - min(t_start)) is
|
||||||
|
# available in both modes; MemoryStore is created only when
|
||||||
|
# enable_data=True so Phase 2 DataExecutor replay is gated on it.
|
||||||
self._memory_store = None
|
self._memory_store = None
|
||||||
if enable_data:
|
if enable_data:
|
||||||
from kernbench.sim_engine.memory_store import MemoryStore
|
from kernbench.sim_engine.memory_store import MemoryStore
|
||||||
from kernbench.sim_engine.op_log import OpLogger
|
|
||||||
self._memory_store = MemoryStore()
|
self._memory_store = MemoryStore()
|
||||||
|
from kernbench.sim_engine.op_log import OpLogger
|
||||||
self._op_logger = OpLogger(memory_store=self._memory_store)
|
self._op_logger = OpLogger(memory_store=self._memory_store)
|
||||||
# Cursor for incremental Phase 2 replay (ADR-0020 D6).
|
# Cursor for incremental Phase 2 replay (ADR-0020 D6).
|
||||||
# SimPy env.now is monotonic so newly logged records always sort
|
# SimPy env.now is monotonic so newly logged records always sort
|
||||||
|
|||||||
Reference in New Issue
Block a user