diff --git a/apps/mvp-node/tinygrad_worker.py b/apps/mvp-node/tinygrad_worker.py index 89298ec..28a5810 100755 --- a/apps/mvp-node/tinygrad_worker.py +++ b/apps/mvp-node/tinygrad_worker.py @@ -1061,6 +1061,7 @@ def ring_readable(cmd: dict[str, Any]) -> None: materialized = materialize_object(payload, sequence, flags) materialized.update( object_id=object_id, + edge_id=ring["edge_id"], sequence=sequence, extent=extent, flags=flags, @@ -1127,6 +1128,11 @@ def execute_step(cmd: dict[str, Any]) -> None: final_stage = bool(cmd.get("final_stage")) input_kind = obj.get("kind") input_extent = int(obj.get("extent", 0)) + validation_ready = time.monotonic() + input_prepare_ms = 0 + forward_ms = 0 + realize_ms = 0 + payload_pack_ms = 0 execution_backend = "pipeline_stage" if not isinstance(model, PipelineStageTinygradModel): execution_backend = "full_transformer" @@ -1136,29 +1142,50 @@ def execute_step(cmd: dict[str, Any]) -> None: fatal("FullTransformerInputUnsupported", step_id=int(cmd["step_id"]), kind=input_kind) if int(obj.get("flags", 0)) & FLAG_BEGIN_SEQUENCE and hasattr(model, "forward_jit"): model.forward_jit.reset() + forward_started = time.monotonic() token_array = model(obj["tensor"], int(obj.get("start_pos", 0))).realize().numpy().reshape(-1) + forward_ready = time.monotonic() + forward_ms = int((forward_ready - forward_started) * 1000) token = int(token_array[0]) + payload_started = time.monotonic() payload = struct.pack("