Add --verify-data CLI flag, Tensor.data property, parallel DataExecutor
- CLI: --verify-data flag enables Phase 2 data verification (ADR-0020) - Tensor.data: returns actual numpy values (verify-data) or zeros placeholder - Tensor.__repr__: shows value summary or data=N/A (placeholder) - DataExecutor: ThreadPoolExecutor for same-timestamp parallel op execution - BenchResult.engine: exposes op_log/memory_store for Phase 2 access Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -62,6 +62,7 @@ def run_bench(
|
||||
correlation_id=correlation_id,
|
||||
trace=None,
|
||||
traces=collected_traces,
|
||||
engine=engine,
|
||||
)
|
||||
|
||||
if completion_policy == CompletionPolicy.LAST_SUBMITTED:
|
||||
@@ -69,7 +70,7 @@ def run_bench(
|
||||
completion, trace = engine.get_completion(last)
|
||||
return BenchResult(
|
||||
completion=completion, correlation_id=correlation_id,
|
||||
trace=trace, traces=collected_traces,
|
||||
trace=trace, traces=collected_traces, engine=engine,
|
||||
)
|
||||
|
||||
if completion_policy == CompletionPolicy.ALL_OK_FAIL_FAST:
|
||||
@@ -80,11 +81,11 @@ def run_bench(
|
||||
if not c.ok:
|
||||
return BenchResult(
|
||||
completion=c, correlation_id=correlation_id,
|
||||
trace=last_trace, traces=collected_traces,
|
||||
trace=last_trace, traces=collected_traces, engine=engine,
|
||||
)
|
||||
return BenchResult(
|
||||
completion=Completion(ok=True), correlation_id=correlation_id,
|
||||
trace=last_trace, traces=collected_traces,
|
||||
trace=last_trace, traces=collected_traces, engine=engine,
|
||||
)
|
||||
|
||||
# LAST_COMPLETED placeholder (needs engine support for timing). Fall back.
|
||||
@@ -92,5 +93,5 @@ def run_bench(
|
||||
completion, trace = engine.get_completion(last)
|
||||
return BenchResult(
|
||||
completion=completion, correlation_id=correlation_id,
|
||||
trace=trace, traces=collected_traces,
|
||||
trace=trace, traces=collected_traces, engine=engine,
|
||||
)
|
||||
|
||||
@@ -314,6 +314,7 @@ class RuntimeContext:
|
||||
t._handle = handle
|
||||
import weakref
|
||||
t._ctx_ref = weakref.ref(self)
|
||||
t._memory_store = getattr(self.engine, "_memory_store", None)
|
||||
self._tensors.append(weakref.ref(t))
|
||||
|
||||
# Install VA→PA mappings via fabric MmuMapMsg
|
||||
|
||||
@@ -5,6 +5,8 @@ import weakref
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal
|
||||
|
||||
import numpy as np
|
||||
|
||||
from kernbench.policy.address.allocator import PEMemAllocator
|
||||
from kernbench.policy.placement.dp import DPPolicy, ShardSpec
|
||||
from kernbench.runtime_api.kernel import TensorArg, TensorArgShard
|
||||
@@ -50,6 +52,20 @@ def dtype_itemsize(dtype: str) -> int:
|
||||
return _DTYPE_ITEMSIZE[dtype]
|
||||
|
||||
|
||||
_NUMPY_DTYPE = {
|
||||
"f16": np.float16, "fp16": np.float16, "float16": np.float16,
|
||||
"f32": np.float32, "fp32": np.float32, "float32": np.float32,
|
||||
"bf16": np.float16,
|
||||
"i8": np.int8, "int8": np.int8,
|
||||
"i16": np.int16, "int16": np.int16,
|
||||
"i32": np.int32, "int32": np.int32,
|
||||
}
|
||||
|
||||
|
||||
def _numpy_dtype(dtype: str) -> np.dtype:
|
||||
return np.dtype(_NUMPY_DTYPE.get(dtype, np.float16))
|
||||
|
||||
|
||||
def deploy_tensor(
|
||||
*,
|
||||
name: str,
|
||||
@@ -129,6 +145,7 @@ class Tensor:
|
||||
self._dp_metadata: DPMetadata | None = None
|
||||
self._handle: TensorHandle | None = None
|
||||
self._ctx_ref: weakref.ref | None = None # set by RuntimeContext
|
||||
self._memory_store = None # set by RuntimeContext when enable_data=True
|
||||
|
||||
def __del__(self) -> None:
|
||||
if self._ctx_ref is None or self._handle is None:
|
||||
@@ -137,6 +154,28 @@ class Tensor:
|
||||
if ctx is not None:
|
||||
ctx._free_tensor(self)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
parts = [f"tensor(name={self.name}, shape={self.shape}, dtype={self.dtype}"]
|
||||
if self._memory_store is not None and self._handle is not None:
|
||||
arr = self.data
|
||||
parts.append(f", mean={float(arr.mean()):.4g}, norm={float(np.linalg.norm(arr)):.4g}")
|
||||
else:
|
||||
parts.append(", data=N/A (placeholder)")
|
||||
parts.append(")")
|
||||
return "".join(parts)
|
||||
|
||||
@property
|
||||
def data(self) -> np.ndarray:
|
||||
"""Tensor data as numpy array. Returns actual values when enable_data=True,
|
||||
zeros placeholder otherwise (like an uninitialized tensor)."""
|
||||
if self._memory_store is not None and self._handle is not None:
|
||||
shard = self._handle.shards[0]
|
||||
try:
|
||||
return self._memory_store.read("hbm", shard.pa, shape=self.shape, dtype=self.dtype)
|
||||
except KeyError:
|
||||
pass
|
||||
return np.zeros(self.shape, dtype=_numpy_dtype(self.dtype))
|
||||
|
||||
@property
|
||||
def itemsize(self) -> int:
|
||||
return dtype_itemsize(self.dtype)
|
||||
|
||||
@@ -12,6 +12,7 @@ class BenchResult:
|
||||
correlation_id: str
|
||||
trace: Trace | None = None
|
||||
traces: list[dict] | None = None
|
||||
engine: object | None = None # GraphEngine ref for Phase 2 data access
|
||||
|
||||
def summary_text(self) -> str:
|
||||
if self.completion.ok:
|
||||
|
||||
Reference in New Issue
Block a user