From 13e93abf24a364ea1c9859e01ed2e5c5632c5de7 Mon Sep 17 00:00:00 2001 From: BANADDA MUBARAKA <83466862+BANADDA@users.noreply.github.com> Date: Wed, 30 Sep 2026 13:00:46 +0300 Subject: [PATCH] envelope: profile at the declared maximum input, fitted by the engine's own tokenizer --- microtensor/envelope/profiler.py | 65 +++++++++++++++++++++++++---- microtensor/harness/engines/gguf.py | 12 +++++- microtensor/harness/engines/onnx.py | 5 +++ 3 files changed, 74 insertions(+), 8 deletions(-) diff --git a/microtensor/envelope/profiler.py b/microtensor/envelope/profiler.py index 9431efb..1cf2a7a 100644 --- a/microtensor/envelope/profiler.py +++ b/microtensor/envelope/profiler.py @@ -2,7 +2,7 @@ import contextlib import time -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path from typing import Any @@ -16,7 +16,7 @@ from microtensor.core.tracks import DECIDE, GENERATE, HardwareClass, get_class, get_track from microtensor.envelope.device import DeviceProfile, conforms, detect from microtensor.envelope.latency import Distribution, summarise -from microtensor.envelope.probe import max_input_decision, max_input_prompt +from microtensor.envelope.probe import declared_tokens, max_input_decision, max_input_prompt from microtensor.envelope.sampler import ResidentSampler from microtensor.harness.contract import Engine, Request, Response, answer from microtensor.harness.registry import engine_for, load_builtin @@ -106,22 +106,73 @@ def plan_payload(plan: ProfilePlan) -> dict[str, Any]: } +def _with_words(plan: ProfilePlan, words: list[str], count: int) -> ProfilePlan: + text = " ".join(words[:count]) + if plan.mode != DECIDE: + return replace(plan, prompt=text) + decision = dict(plan.max_input.get("decision") or {}) + decision["context"] = text + return replace(plan, prompt=text, max_input={**plan.max_input, "decision": decision}) + + +def fit(engine: Engine, plan: ProfilePlan) -> tuple[ProfilePlan, int]: + measure = getattr(engine, "input_tokens", None) + declared = declared_tokens(plan.max_input) + if not callable(measure) or declared <= 0: + return plan, 0 + budget = declared if plan.mode == DECIDE else declared - plan.max_output_tokens + if budget <= 0: + return plan, 0 + words = plan.prompt.split() + while words and len(words) < budget * 2: + words = words + words + if not words: + return plan, 0 + + def size(count: int) -> int: + return int(measure(_request(_with_words(plan, words, count), 0))) + + low, high = 1, len(words) + if size(low) > budget: + return _with_words(plan, words, low), size(low) + while low < high: + middle = (low + high + 1) // 2 + if size(middle) <= budget: + low = middle + else: + high = middle - 1 + fitted = _with_words(plan, words, low) + return fitted, size(low) + + def _cold_start( engine: Engine, artifact: Path, manifest: LoadManifest, plan: ProfilePlan -) -> tuple[float, Response]: +) -> tuple[float, Response, ProfilePlan, int]: started = time.perf_counter() try: engine.load(artifact, manifest) except Exception as exc: raise ProfileError(f"artifact failed to load: {exc}") from exc + loaded_ms = (time.perf_counter() - started) * 1000.0 + + plan, fed = fit(engine, plan) + started = time.perf_counter() response = answer(engine, _request(plan, 0)) - elapsed_ms = (time.perf_counter() - started) * 1000.0 + elapsed_ms = loaded_ms + (time.perf_counter() - started) * 1000.0 if not response.ok: raise ProfileError(f"artifact produced no output on a cold cache: {response.error}") if not _produced(response, plan): raise ProfileError("artifact produced no output on a cold cache") - return elapsed_ms, response + return elapsed_ms, response, plan, fed + + +def _input_at_peak(plan: ProfilePlan, fed: int) -> dict[str, Any]: + found = {key: value for key, value in plan.max_input.items() if key != "decision"} + if fed: + found["fed_tokens"] = fed + found["declared_tokens"] = declared_tokens(plan.max_input) + return found def profile( @@ -140,7 +191,7 @@ def profile( sampler = ResidentSampler(plan.sample_interval_ms) sampler.start() try: - cold_start_ms, _ = _cold_start(engine, artifact, manifest, plan) + cold_start_ms, _, plan, fed = _cold_start(engine, artifact, manifest, plan) sampler.mark() ttfts: list[float] = [] @@ -175,7 +226,7 @@ def profile( envelope = MeasuredEnvelope( size_bytes=size_bytes, peak_rss_bytes=sampler.peak_bytes, - input_at_peak=dict(plan.max_input), + input_at_peak=_input_at_peak(plan, fed), ttft_p50_ms=round(latency.p50), ttft_p95_ms=round(latency.p95), tokens_per_second=sum(throughputs) / len(throughputs) if throughputs else 0.0, diff --git a/microtensor/harness/engines/gguf.py b/microtensor/harness/engines/gguf.py index 15ca905..df887d4 100644 --- a/microtensor/harness/engines/gguf.py +++ b/microtensor/harness/engines/gguf.py @@ -7,7 +7,7 @@ from typing import Any, Final from microtensor.core.protocol import ArtifactFormat, LoadManifest -from microtensor.core.tracks import DECISION_PROMPT_VERSION, Decoding +from microtensor.core.tracks import DECIDE, DECISION_PROMPT_VERSION, Decoding from microtensor.harness import decision_prompt from microtensor.harness.contract import ( EngineError, @@ -313,6 +313,16 @@ def _rewind(self, shared: int) -> None: raise EngineError("cache rollback refused") self._model.n_tokens = shared + def input_tokens(self, request: Request) -> int: + if request.mode == DECIDE: + context, questions = decision_prompt.parse(request.inputs.get("decision")) + return max(len(self._decision_tokens(context, question)) for question in questions) + text = request.prompt + if request.chat: + rendered = self._render_chat([{"role": "user", "content": request.prompt}]) + text = rendered if rendered is not None else request.prompt + return len(self._model.tokenize(text.encode("utf-8"), add_bos=True, special=True)) + def _answer_scores(self, question: decision_prompt.Question) -> list[float]: logits = _llama().llama_get_logits_ith(self._model._ctx.ctx, -1) return [float(logits[self._answer_ids[text]]) for text in question.answers] diff --git a/microtensor/harness/engines/onnx.py b/microtensor/harness/engines/onnx.py index 513a86e..e44cac4 100644 --- a/microtensor/harness/engines/onnx.py +++ b/microtensor/harness/engines/onnx.py @@ -166,6 +166,11 @@ def _declared_tokens(self) -> int: return value return 0 + def input_tokens(self, request: Request) -> int: + if self._tokenizer is None: + return 0 + return len(self._tokenizer.encode(request.prompt).ids) + def generate(self, request: Request) -> Response: if self._session is None or self._tokenizer is None: return Response.failed(request.task_ref, "engine was asked to generate before load")