#!/usr/bin/env python3 from __future__ import annotations import hashlib import json import mmap import linecache import os import sys import threading import time import struct import traceback import shutil import urllib.parse import urllib.request from pathlib import Path from typing import Any Tensor: Any = None dtypes: Any = None model: Any = None tokenizer: Any = None role: dict[str, Any] = {} loaded: dict[str, Any] = {} arena: mmap.mmap | None = None rings: dict[int, dict[str, Any]] = {} device_objects: dict[int, dict[str, Any]] = {} next_handle = 42 HEADER_LEN = 40 WORKER_GENERATION = 1 BENCHMARK_SCHEMA = 1 _benchmark_start = time.monotonic() _benchmark_seq = 0 class CpuLineSampler: def __init__( self, *, phase: str, request_id: int | None, model_id: str | None, interval_secs: float, ) -> None: self.phase = phase self.request_id = request_id self.model_id = model_id self.interval_secs = interval_secs self.target_thread_id = threading.get_ident() self.samples: dict[tuple[str, int, str], int] = {} self.wall_start = time.perf_counter() self.process_cpu_start = time.process_time() self._running = True self._thread = threading.Thread(target=self._run, name="cpu-line-sampler", daemon=True) self._thread.start() def _run(self) -> None: while self._running: frame = sys._current_frames().get(self.target_thread_id) if frame is not None: code = frame.f_code key = (code.co_filename, frame.f_lineno, code.co_name) self.samples[key] = self.samples.get(key, 0) + 1 time.sleep(self.interval_secs) def stop(self) -> None: self._running = False self._thread.join(timeout=max(0.25, self.interval_secs * 4.0)) wall_elapsed_ms = (time.perf_counter() - self.wall_start) * 1000.0 process_cpu_elapsed_ms = (time.process_time() - self.process_cpu_start) * 1000.0 total_samples = sum(self.samples.values()) top = [] for (filename, line, function), count in sorted( self.samples.items(), key=lambda item: item[1], reverse=True )[:32]: top.append( { "file": filename, "line": line, "function": function, "source": linecache.getline(filename, line).strip(), "samples": count, "percent": round((count * 100.0 / total_samples), 2) if total_samples else 0.0, } ) control( type="CpuLineProfileSummary", phase=self.phase, request_id=self.request_id, model_id=self.model_id, interval_ms=round(self.interval_secs * 1000.0, 3), wall_elapsed_ms=round(wall_elapsed_ms, 3), process_cpu_elapsed_ms=round(process_cpu_elapsed_ms, 3), process_cpu_over_wall=round(process_cpu_elapsed_ms / wall_elapsed_ms, 4) if wall_elapsed_ms > 0.0 else 0.0, total_samples=total_samples, top=top, ) def start_cpu_line_sampler( *, phase: str, request_id: int | None, model_id: str | None, ) -> CpuLineSampler | None: raw = os.environ.get("MVP_CPU_LINE_PROFILE") if not env_flag("MVP_CPU_LINE_PROFILE", False): control( type="CpuLineProfileSkipped", phase=phase, request_id=request_id, model_id=model_id, env_value=raw, ) return None interval_ms = float(os.environ.get("MVP_CPU_LINE_PROFILE_INTERVAL_MS", "2")) interval_secs = max(0.0005, interval_ms / 1000.0) control( type="CpuLineProfileStarted", phase=phase, request_id=request_id, model_id=model_id, interval_ms=round(interval_secs * 1000.0, 3), ) return CpuLineSampler( phase=phase, request_id=request_id, model_id=model_id, interval_secs=interval_secs, ) def stop_cpu_line_sampler(sampler: CpuLineSampler | None) -> None: if sampler is not None: sampler.stop() def env_flag(name: str, default: bool = True) -> bool: raw = os.environ.get(name) if raw is None: return default return raw.strip().lower() not in {"0", "false", "no", "off"} def benchmark_stamp() -> dict[str, Any]: global _benchmark_seq _benchmark_seq += 1 return { "schema": BENCHMARK_SCHEMA, "component": "tinygrad-worker", "pid": os.getpid(), "seq": _benchmark_seq, "wall_unix_ms": time.time_ns() // 1_000_000, "mono_ms": int((time.monotonic() - _benchmark_start) * 1000), } def env_int(name: str) -> int | None: raw = os.environ.get(name) if raw is None: return None try: return int(raw) except ValueError: return None def control(**event: Any) -> None: event.setdefault("benchmark", benchmark_stamp()) if (run_id := env_int("MVP_RUN_ID")) is not None: event.setdefault("run_id", run_id) if (node_id := env_int("MVP_LOGICAL_NODE_ID")) is not None: event.setdefault("node_id", node_id) if (stage_index := env_int("MVP_STAGE_INDEX")) is not None: event.setdefault("stage_index", stage_index) print(json.dumps(event, separators=(",", ":")), flush=True) def log(message: str) -> None: print(f"mvp_tinygrad_worker: {message}", file=sys.stderr, flush=True) def fatal(reason: str, **fields: Any) -> None: control(type="WorkerFatal", reason=reason, **fields) raise SystemExit(1) def configure_tinygrad_cuda_compiler(device: str) -> None: if device.split(":", 1)[0].upper() != "CUDA": return if os.environ.get("CUDA_PTX") or os.environ.get("CUDA_CC"): return if shutil.which("nvcc") is not None: return os.environ["CUDA_PTX"] = "1" control(type="TinygradCudaCompilerSelected", requested_device=device, compiler="PTX", reason="nvcc_not_found") def select_tinygrad_device(device: str) -> str: device_kind = device.split(":", 1)[0].upper() if device_kind == "CPU" and ":" not in device and shutil.which("clang") is None: selected = "CPU:X86" os.environ["DEV"] = selected control(type="TinygradCpuCompilerSelected", requested_device=device, selected_device=selected, compiler="X86", reason="clang_not_found") return selected os.environ["DEV"] = device configure_tinygrad_cuda_compiler(device) return device def initialize(cmd: dict[str, Any]) -> None: global Tensor, dtypes, arena if int(cmd.get("helper_abi_version", 1)) != 1: fatal("UnsupportedHelperAbi", helper_abi_version=cmd.get("helper_abi_version")) requested_device = str(cmd.get("backend", {}).get("device") or os.environ.get("DEV") or "CUDA") device = select_tinygrad_device(requested_device) arena_fd = os.environ.get("MVP_ARENA_FD") if arena_fd is not None: arena_bytes = int(os.environ.get("MVP_ARENA_BYTES", "0") or "0") if arena_bytes > 0: arena = mmap.mmap(int(arena_fd), arena_bytes) started = time.monotonic() control(type="TinygradImportStarted", requested_device=device, env_DEV=os.environ.get("DEV")) from tinygrad import Tensor as TinyTensor, dtypes as tiny_dtypes control(type="TinygradImportReady", requested_device=device, env_DEV=os.environ.get("DEV")) Tensor = TinyTensor dtypes = tiny_dtypes control(type="TinygradDeviceProbeStarted", requested_device=device) value = Tensor([1], dtype=dtypes.int32).realize().numpy().tolist() control(type="TinygradDeviceProbeReady", requested_device=device, probe_result=value) control( type="WorkerReady", pid=os.getpid(), backend={"requested_device": device, "env_DEV": os.environ.get("DEV"), "tinygrad_device": device}, cuda_probe=value, elapsed_ms=int((time.monotonic() - started) * 1000), ) def configure_role(cmd: dict[str, Any]) -> None: config = cmd.get("config", {}) role.clear() role.update( role_id=int(cmd.get("role_id", 1)), run_id=int(config.get("run_id", 1)), stage_index=int(config.get("stage_index", 0)), layer_start=int(config.get("layer_start", 0)), layer_end_exclusive=int(config.get("layer_end_exclusive", 0)), ) control(type="RoleConfigured", role_id=role["role_id"], stage_index=role["stage_index"]) def cache_root() -> Path: raw = os.environ.get("MVP_MODEL_CACHE_DIR", "").strip() root = Path(raw).expanduser() if raw else Path.home() / ".cache" / "mvp-node" root.mkdir(parents=True, exist_ok=True) return root def hf_url(repo: str, file: str, revision: str | None) -> str: encoded_file = "/".join(urllib.parse.quote(part) for part in file.split("/")) return f"https://huggingface.co/{repo}/resolve/{revision or 'main'}/{encoded_file}" def source_url(source: dict[str, Any]) -> str | None: if "HuggingFaceGguf" not in source: return None hf = source["HuggingFaceGguf"] return hf_url(str(hf["repo"]), str(hf["file"]), hf.get("revision")) def source_path(source: dict[str, Any]) -> Path | None: if "LocalPath" not in source: return None return Path(str(source["LocalPath"])).expanduser() def source_kind(source: dict[str, Any]) -> str: if "LocalPath" in source: return "LocalPath" if "HuggingFaceGguf" in source: return "HuggingFaceGguf" return "Unknown" def cache_path_for(url: str) -> Path: parsed = urllib.parse.urlparse(url) basename = Path(parsed.path).name or "model.gguf" digest = hashlib.sha256(url.encode("utf-8")).hexdigest()[:16] return cache_root() / f"{digest}-{basename}" def request_headers() -> dict[str, str]: headers = {"User-Agent": "swactor-mvp-node/0.1"} token = os.environ.get("HF_TOKEN", "").strip() if token: headers["Authorization"] = f"Bearer {token}" return headers def fetch_whole(source: dict[str, Any]) -> Path: local = source_path(source) if local is not None: control(type="GgufLocalPathStatStarted", path=str(local)) if not local.is_file(): fatal("GgufLocalPathMissing", path=str(local)) stat = local.stat() control(type="GgufCacheReady", path=str(local), bytes=stat.st_size, cache_hit=True, source="local") return local url = source_url(source) if not url: fatal("UnsupportedGgufSource", source=source) target = cache_path_for(url) if target.is_file() and target.stat().st_size > 0: control(type="GgufCacheReady", path=str(target), bytes=target.stat().st_size, cache_hit=True, url=url) return target partial = target.with_name(target.name + ".partial") started = time.monotonic() req = urllib.request.Request(url, headers=request_headers()) control(type="GgufDownloadStarted", url=url, path=str(target)) try: with urllib.request.urlopen(req, timeout=60) as response, partial.open("wb") as out: total = int(response.headers.get("Content-Length") or 0) done = 0 last_event = 0.0 while True: chunk = response.read(1024 * 1024) if not chunk: break out.write(chunk) done += len(chunk) now = time.monotonic() if now - last_event >= float(os.environ.get("MVP_DOWNLOAD_PROGRESS_SECS", "5")): control( type="GgufDownloadProgress", bytes_done=done, bytes_total=total, elapsed_ms=int((now - started) * 1000), ) last_event = now partial.replace(target) except Exception as exc: try: partial.unlink(missing_ok=True) except Exception: pass fatal("GgufDownloadFailed", url=url, error=str(exc)) control( type="GgufCacheReady", path=str(target), bytes=target.stat().st_size, cache_hit=False, elapsed_ms=int((time.monotonic() - started) * 1000), url=url, ) return target def require_tinygrad() -> Any: if Tensor is None: fatal("BackendNotInitialized") return Tensor class PipelineStageTinygradModel: def __init__( self, *, block_count: int, dim: int, hidden_dim: int, n_heads: int, n_kv_heads: int, norm_eps: float, vocab_size: int, head_dim: int, rope_theta: float, rope_dim: int, v_head_dim: int, max_context: int, qk_norm: int, num_experts: int, num_experts_per_tok: int, norm_topk_prob: bool, qkv_bias: bool, expert_bias: bool, first_stage: bool, final_stage: bool, nn_mod: Any, config_cls: Any, block_cls: Any, ) -> None: block_config = config_cls( num_blocks=block_count, dim=dim, hidden_dim=hidden_dim, n_heads=n_heads, n_kv_heads=n_kv_heads, norm_eps=norm_eps, vocab_size=vocab_size, head_dim=head_dim, rope_theta=rope_theta, rope_dim=rope_dim, v_head_dim=v_head_dim, max_context=max_context, qk_norm=qk_norm, num_experts=num_experts, num_experts_per_tok=num_experts_per_tok, norm_topk_prob=norm_topk_prob, qkv_bias=qkv_bias, expert_bias=expert_bias, ) self.blk = [block_cls(block_config) for _ in range(block_count)] self.max_context = max_context self.hidden_dim = dim self.first_stage = first_stage self.final_stage = final_stage if first_stage: self.token_embd = nn_mod.Embedding(vocab_size, dim) if final_stage: self.output_norm = nn_mod.RMSNorm(dim, norm_eps) self.output = nn_mod.Linear(dim, vocab_size, bias=False) def token_hidden(self, tokens_tensor: Any) -> Any: return self.token_embd(tokens_tensor).float() def forward_hidden(self, hidden: Any, start_pos: Any) -> Any: for block in self.blk: hidden = block(hidden, start_pos) return hidden.contiguous() def next_token(self, hidden: Any) -> Any: return self.output(self.output_norm(hidden))[:, -1, :].argmax(-1, keepdim=True) def __call__(self, tokens_tensor: Any, start_pos: Any) -> Any: return self.next_token(self.forward_hidden(self.token_hidden(tokens_tensor), start_pos)) def remap_stage_state_dict( state_dict: dict[str, Any], *, layer_start: int, layer_end_exclusive: int, first_stage: bool, final_stage: bool, ) -> dict[str, Any]: if final_stage and "output.weight" not in state_dict and "token_embd.weight" in state_dict: state_dict["output.weight"] = state_dict["token_embd.weight"] remapped: dict[str, Any] = {} prefix = "blk." for key, value in state_dict.items(): if key.startswith(prefix): parts = key.split(".", 2) if len(parts) != 3: continue block_index = int(parts[1]) if layer_start <= block_index < layer_end_exclusive: remapped[f"blk.{block_index - layer_start}.{parts[2]}"] = value elif first_stage and key == "token_embd.weight": remapped[key] = value elif final_stage and (key == "output_norm.weight" or key == "output.weight"): remapped[key] = value return remapped def load_pipeline_stage_model( path: Path, *, max_context: int, layer_start: int, layer_end_exclusive: int, ) -> tuple[PipelineStageTinygradModel, dict[str, Any]]: TensorCls = require_tinygrad() from tinygrad import nn from tinygrad.llm.gguf import gguf_load from tinygrad.llm.model import TransformerBlock, TransformerConfig kv, state_dict = gguf_load(path) state_dict = {key: value.cast("float16") if env_flag("HALF", True) else value for key, value in state_dict.items()} if "output.weight" not in state_dict and "token_embd.weight" in state_dict: state_dict["output.weight"] = state_dict["token_embd.weight"] arch = kv["general.architecture"] max_context = min(max_context, int(kv[f"{arch}.context_length"])) n_heads = int(kv[f"{arch}.attention.head_count"]) n_kv_heads = int(kv[f"{arch}.attention.head_count_kv"]) dim = int(kv[f"{arch}.embedding_length"]) kv_lora_rank = int(kv.get(f"{arch}.attention.kv_lora_rank", 0)) head_dim = int(kv.get(f"{arch}.attention.key_length_mla", kv.get(f"{arch}.attention.key_length", dim // n_heads))) rope_dim = int(kv.get(f"{arch}.rope.dimension_count", head_dim)) for name in list(state_dict): if ("attn_q.weight" in name or "attn_q_b.weight" in name) and (arch == "llama" or kv_lora_rank): weight = state_dict[name].reshape(n_heads, state_dict[name].shape[0] // n_heads, -1) prefix = head_dim - rope_dim state_dict[name] = ( weight[:, :prefix] .cat(weight[:, prefix:].rearrange("n (h two) d -> n (two h) d", two=2), dim=1) .reshape(-1, weight.shape[-1]) ) elif arch == "llama" and "attn_k.weight" in name: weight = state_dict[name].reshape(n_kv_heads, state_dict[name].shape[0] // n_kv_heads, -1) state_dict[name] = weight.rearrange("n (h two) d -> n (two h) d", two=2).reshape(-1, weight.shape[-1]) elif kv_lora_rank and "attn_kv_a_mqa.weight" in name: state_dict[name] = state_dict[name][:kv_lora_rank].cat( state_dict[name][kv_lora_rank:].rearrange("(h two) d -> (two h) d", two=2), dim=0, ) total_layers = int(kv[f"{arch}.block_count"]) - int(kv.get(f"{arch}.nextn_predict_layers", 0)) first_stage = layer_start == 0 final_stage = layer_end_exclusive >= total_layers qk_key = f"blk.{layer_start}.attn_q_norm.weight" qk_norm = int(state_dict[qk_key].shape[0]) if qk_key in state_dict else 0 stage_model = PipelineStageTinygradModel( block_count=layer_end_exclusive - layer_start, dim=dim, hidden_dim=int(kv.get(f"{arch}.expert_feed_forward_length", kv.get(f"{arch}.feed_forward_length", 0))), n_heads=n_heads, n_kv_heads=n_kv_heads, norm_eps=float(kv[f"{arch}.attention.layer_norm_rms_epsilon"]), vocab_size=len(kv["tokenizer.ggml.tokens"]), head_dim=head_dim, rope_theta=float(kv[f"{arch}.rope.freq_base"]), rope_dim=rope_dim, v_head_dim=int(kv.get(f"{arch}.attention.value_length_mla", kv.get(f"{arch}.attention.value_length", head_dim))), max_context=max_context, qk_norm=qk_norm, num_experts=int(kv.get(f"{arch}.expert_count", 0)), num_experts_per_tok=int(kv.get(f"{arch}.expert_used_count", 0)), norm_topk_prob=bool(kv.get(f"{arch}.expert_weights_norm", arch in ("qwen3moe", "qwen35moe"))), qkv_bias="blk.0.attn_q.bias" in state_dict, expert_bias=f"blk.{int(kv.get(f'{arch}.leading_dense_block_count', 0))}.exp_probs_b.bias" in state_dict, first_stage=first_stage, final_stage=final_stage, nn_mod=nn, config_cls=TransformerConfig, block_cls=TransformerBlock, ) stage_state = remap_stage_state_dict( state_dict, layer_start=layer_start, layer_end_exclusive=layer_end_exclusive, first_stage=first_stage, final_stage=final_stage, ) loaded_params = nn.state.load_state_dict(stage_model, stage_state, verbose=False, consume=True, realize=False) for param in loaded_params: param.replace(param.contiguous()) if loaded_params: TensorCls.realize(*loaded_params) return stage_model, kv def load_weights(cmd: dict[str, Any]) -> None: global model, tokenizer started = time.monotonic() model_id = str(cmd["model_id"]) source = cmd["gguf_source"] control( type="LoadWeightsStarted", model_id=model_id, source_kind=source_kind(source), layer_start=int(cmd.get("layer_start", 0)), layer_end_exclusive=int(cmd.get("layer_end_exclusive", 0)), ) control(type="GgufResolveStarted", model_id=model_id, source_kind=source_kind(source)) path = fetch_whole(source) model_bytes = path.stat().st_size control(type="GgufResolveReady", model_id=model_id, path=str(path), bytes=model_bytes) layer_start = int(cmd.get("layer_start", 0)) layer_end_exclusive = int(cmd.get("layer_end_exclusive", 0)) try: control(type="TinygradLlmImportStarted", model_id=model_id) from tinygrad.llm.cli import SimpleTokenizer control(type="TinygradLlmImportReady", model_id=model_id) max_context_raw = os.environ.get("MVP_MAX_CONTEXT", "512") max_context = int(max_context_raw) if max_context_raw else 512 control( type="PipelineStageFromGgufStarted", model_id=model_id, path=str(path), bytes=model_bytes, max_context=max_context, layer_start=layer_start, layer_end_exclusive=layer_end_exclusive, requested_device=os.environ.get("DEV"), ) model, kv = load_pipeline_stage_model( path, max_context=max_context, layer_start=layer_start, layer_end_exclusive=layer_end_exclusive, ) control( type="PipelineStageFromGgufReady", model_id=model_id, path=str(path), bytes=model_bytes, max_context=model.max_context, layer_start=layer_start, layer_end_exclusive=layer_end_exclusive, first_stage=model.first_stage, final_stage=model.final_stage, requested_device=os.environ.get("DEV"), ) tok_src = cmd.get("tokenizer", {"EmbeddedGguf": None}) if "EmbeddedGguf" in tok_src: if kv.get("tokenizer.ggml.pre") == "smollm": kv = dict(kv) kv["tokenizer.ggml.pre"] = "qwen2" control(type="TokenizerBuildStarted", model_id=model_id, source="EmbeddedGguf") tokenizer = SimpleTokenizer.from_gguf_kv(kv) control(type="TokenizerBuildReady", model_id=model_id, source="EmbeddedGguf") else: fatal("UnsupportedTokenizerSource", tokenizer=tok_src) except SystemExit: raise except Exception as exc: tb = traceback.format_exc() print(tb, file=sys.stderr, flush=True) fatal("ModelLoadFailed", error=str(exc), traceback=tb) loaded.clear() loaded.update( model_id=model_id, path=str(path), layer_start=layer_start, layer_end_exclusive=layer_end_exclusive, hidden_dim=int(getattr(model, "hidden_dim", 0)), max_context=int(getattr(model, "max_context", 0)), eos_token_id=int(kv.get("tokenizer.ggml.eos_token_id", 0)), ) control( type="WeightsLoaded", model_id=model_id, path=str(path), layer_start=layer_start, layer_end_exclusive=layer_end_exclusive, elapsed_ms=int((time.monotonic() - started) * 1000), ) def prompt_template_name() -> str: explicit = os.environ.get("MVP_PROMPT_TEMPLATE") if explicit is not None: return explicit.strip().lower() model_id = str(loaded.get("model_id", "")).lower() if "smollm" in model_id: return "smollm-chat" return "llama3-chat" def model_prompt_text(prompt: str) -> tuple[str, str]: template = prompt_template_name() if template in {"", "raw", "none", "off", "false", "0"}: return prompt, "raw" if template in {"llama3", "llama3-chat", "llama-3", "llama-3-chat"}: return ( "<|begin_of_text|>" "<|start_header_id|>user<|end_header_id|>\n\n" f"{prompt}" "<|eot_id|>" "<|start_header_id|>assistant<|end_header_id|>\n\n", "llama3-chat", ) if template in {"smollm", "smollm-chat", "smollm2", "smollm2-chat"}: return ( "<|im_start|>user\n" f"{prompt}" "<|im_end|>\n" "<|im_start|>assistant\n", "smollm-chat", ) return prompt, "raw" def strip_chat_stop_markers(text: str) -> str: cut = len(text) for marker in ( "<|eot_id|>", "<|end_of_text|>", "<|start_header_id|>", "<|im_end|>", "<|endoftext|>", "<|im_start|>", ): index = text.find(marker) if index >= 0: cut = min(cut, index) return text[:cut].rstrip() def decode_greedy_device_resident( prompt_tokens: list[int], max_tokens: int, *, request_id: int | None, model_id: str | None, progress_every: int, ) -> list[int]: if max_tokens <= 0: return [] max_context = int(getattr(model, "max_context", len(prompt_tokens) + max_tokens)) generation_limit = min(max_tokens, max(0, max_context - len(prompt_tokens))) if generation_limit <= 0: control( type="DecodeContextFull", request_id=request_id, model_id=model_id, prompt_tokens=len(prompt_tokens), max_context=max_context, ) return [] if generation_limit < max_tokens: control( type="DecodeLimitedByContext", request_id=request_id, model_id=model_id, prompt_tokens=len(prompt_tokens), requested_tokens=max_tokens, generation_limit=generation_limit, max_context=max_context, ) TensorCls = require_tinygrad() from tinygrad.uop.ops import UOp if hasattr(model, "forward_jit"): model.forward_jit.reset() use_symbolic_pos = os.environ.get("SYM", "1").strip().lower() not in {"0", "false", "no", "off"} pos_upper_bound = max(1, max_context - 1) symbolic_start_pos = UOp.variable("start_pos", 1, pos_upper_bound) next_token = model(TensorCls([prompt_tokens], dtype="int32"), 0).realize() generated_tensors = [] for token_index in range(generation_limit): generated_tensors.append(next_token.clone().realize()) tokens_generated = token_index + 1 if tokens_generated == 1: control( type="FirstTokenReady", request_id=request_id, model_id=model_id, token_index=1, prompt_tokens=len(prompt_tokens), ) elif progress_every > 0 and tokens_generated % progress_every == 0: control( type="TokenProgress", request_id=request_id, model_id=model_id, tokens_generated=tokens_generated, prompt_tokens=len(prompt_tokens), ) if tokens_generated >= generation_limit: break start_pos = len(prompt_tokens) + token_index pos = symbolic_start_pos.bind(start_pos) if use_symbolic_pos else start_pos next_token = model(next_token, pos).realize() generated_tensor = ( generated_tensors[0] if len(generated_tensors) == 1 else generated_tensors[0].cat(*generated_tensors[1:], dim=1) ) generated_array = generated_tensor.numpy().reshape(-1).tolist() return [int(token) for token in generated_array] def infer_prompt(cmd: dict[str, Any]) -> None: if model is None or tokenizer is None: fatal("WeightsNotLoaded") prompt = str(cmd.get("prompt", "")) max_tokens = int(cmd.get("max_tokens", 1)) request_id_raw = cmd.get("request_id") request_id = int(request_id_raw) if request_id_raw is not None else None started = time.monotonic() control( type="PromptStarted", request_id=request_id, model_id=loaded.get("model_id"), prompt_bytes=len(prompt.encode("utf-8")), prompt_chars=len(prompt), max_tokens=max_tokens, ) model_prompt, prompt_template = model_prompt_text(prompt) control( type="PromptEncodeStarted", request_id=request_id, model_id=loaded.get("model_id"), prompt_template=prompt_template, ) prompt_tokens = tokenizer.encode(model_prompt) control( type="PromptEncodeReady", request_id=request_id, model_id=loaded.get("model_id"), prompt_bytes=len(prompt.encode("utf-8")), model_prompt_bytes=len(model_prompt.encode("utf-8")), prompt_template=prompt_template, prompt_tokens=len(prompt_tokens), ) progress_every = int(os.environ.get("MVP_TOKEN_PROGRESS_EVERY", "16") or "16") control( type="DecodeStarted", request_id=request_id, model_id=loaded.get("model_id"), prompt_tokens=len(prompt_tokens), max_tokens=max_tokens, decode_impl="device_resident_greedy", ) cpu_sampler = start_cpu_line_sampler( phase="decode", request_id=request_id, model_id=loaded.get("model_id"), ) try: generated = decode_greedy_device_resident( prompt_tokens, max_tokens, request_id=request_id, model_id=loaded.get("model_id"), progress_every=progress_every, ) finally: stop_cpu_line_sampler(cpu_sampler) control( type="DecodeReady", request_id=request_id, model_id=loaded.get("model_id"), prompt_tokens=len(prompt_tokens), tokens_generated=len(generated), ) control(type="TextDecodeStarted", request_id=request_id, model_id=loaded.get("model_id"), tokens_generated=len(generated)) raw_text = tokenizer.decode(generated) if generated else "" text = strip_chat_stop_markers(raw_text) control( type="TextDecodeReady", request_id=request_id, model_id=loaded.get("model_id"), tokens_generated=len(generated), text_bytes=len(text.encode("utf-8")), ) control( type="PromptCompleted", request_id=request_id, model_id=loaded.get("model_id"), prompt_tokens=prompt_tokens, generated_tokens=generated, text=text, elapsed_ms=int((time.monotonic() - started) * 1000), ) def require_arena() -> mmap.mmap: if arena is None: fatal("ArenaNotMapped") return arena def install_ring(cmd: dict[str, Any]) -> None: ring_id = int(cmd["ring_id"]) layout = cmd["layout"] spec = cmd["object_spec"] rings[ring_id] = { "ring_id": ring_id, "edge_id": int(cmd["edge_id"]), "port": str(cmd.get("port", "")), "direction": str(cmd["direction"]), "data_offset": int(layout["data_offset"]), "data_capacity": int(layout["data_bytes"]), "max_extent": int(spec["max_extent"]), "alignment": int(spec["alignment"]), "next_sequence": 0, } control( type="RingInstalled", ring_id=ring_id, edge_id=rings[ring_id]["edge_id"], direction=rings[ring_id]["direction"], ) def uninstall_ring(cmd: dict[str, Any]) -> None: ring_id = int(cmd["ring_id"]) rings.pop(ring_id, None) control(type="RingUninstalled", ring_id=ring_id) def parse_record(ring: dict[str, Any]) -> tuple[int, int, int, int, bytes]: view = require_arena() base = ring["data_offset"] header = view[base : base + HEADER_LEN] if len(header) < HEADER_LEN: fatal("EofBeforeFullHeader", ring_id=ring["ring_id"]) if header[0:4] != b"MO01": fatal("InvalidObjectMagic", ring_id=ring["ring_id"]) version = struct.unpack_from(" ring["max_extent"]: fatal("ObjectExtentInvalid", ring_id=ring["ring_id"], object_id=object_id, extent=extent) if ring["alignment"] and extent % ring["alignment"] != 0: fatal( "ObjectExtentAlignmentViolation", ring_id=ring["ring_id"], object_id=object_id, extent=extent, alignment=ring["alignment"], ) if sequence != ring["next_sequence"]: fatal("SequenceViolation", ring_id=ring["ring_id"], expected=ring["next_sequence"], actual=sequence) payload = bytes(view[base + HEADER_LEN : base + HEADER_LEN + extent]) ring["next_sequence"] += 1 return object_id, sequence, extent, flags, payload def payload_words(payload: bytes) -> list[int]: if len(payload) % 4 != 0: fatal("PayloadNotU32Aligned", extent=len(payload)) if not payload: return [] return list(struct.unpack(f"<{len(payload) // 4}I", payload)) def object_start_pos(sequence: int, token_count: int) -> int: if sequence == 0: role["prompt_tokens"] = token_count return 0 return int(role.get("prompt_tokens", 1)) + sequence - 1 def materialize_object(payload: bytes, sequence: int) -> dict[str, Any]: if not isinstance(model, PipelineStageTinygradModel): return { "kind": "words", "words": payload_words(payload), "payload": payload, "start_pos": object_start_pos(sequence, max(1, len(payload) // 4)), } TensorCls = require_tinygrad() if bool(getattr(model, "first_stage", False)) and int(role.get("layer_start", 0)) == 0: tokens = payload_words(payload) token_count = len(tokens) return { "kind": "tokens", "tokens": tokens, "tensor": TensorCls([tokens], dtype="int32").realize(), "start_pos": object_start_pos(sequence, token_count), } import numpy as np hidden_dim = int(loaded.get("hidden_dim") or getattr(model, "hidden_dim", 0)) if hidden_dim <= 0: fatal("HiddenDimMissing") bytes_per_token = hidden_dim * 2 if len(payload) % bytes_per_token != 0: fatal("ActivationExtentInvalid", extent=len(payload), hidden_dim=hidden_dim) token_count = len(payload) // bytes_per_token array = np.frombuffer(payload, dtype=np.float16).copy().reshape(1, token_count, hidden_dim) return { "kind": "activation", "tensor": TensorCls(array).realize(), "start_pos": object_start_pos(sequence, token_count), } def ring_readable(cmd: dict[str, Any]) -> None: global next_handle ring_id = int(cmd["ring_id"]) ring = rings[ring_id] if ring["direction"] != "ingress": fatal("WrongRingDirection", ring_id=ring_id, direction=ring["direction"]) object_id, sequence, extent, flags, payload = parse_record(ring) handle = next_handle next_handle += 1 materialized = materialize_object(payload, sequence) materialized.update( object_id=object_id, sequence=sequence, extent=extent, flags=flags, payload=payload, ) device_objects[handle] = materialized control( type="ObjectLoaded", ring_id=ring_id, edge_id=ring["edge_id"], object_id=object_id, sequence=sequence, extent=extent, handle_generation=WORKER_GENERATION, handle_id=handle, ) def encode_record(object_id: int, sequence: int, payload: bytes, flags: int = 0) -> bytes: header = bytearray(HEADER_LEN) header[0:4] = b"MO01" struct.pack_into(" int: if len(payload) > ring["max_extent"]: fatal("OutputExtentInvalid", ring_id=ring["ring_id"], extent=len(payload), max_extent=ring["max_extent"]) if ring["alignment"] and len(payload) % ring["alignment"] != 0: fatal("OutputExtentAlignmentViolation", ring_id=ring["ring_id"], extent=len(payload), alignment=ring["alignment"]) record = encode_record(object_id, sequence, payload, flags) if len(record) > ring["data_capacity"]: fatal("OutputRingCapacityExceeded", ring_id=ring["ring_id"], record_bytes=len(record), capacity=ring["data_capacity"]) view = require_arena() base = ring["data_offset"] view[base : base + len(record)] = record return len(record) def execute_step(cmd: dict[str, Any]) -> None: if not role: fatal("RoleNotConfigured") handle = int(cmd["input_handle_id"]) obj = device_objects.get(handle) if obj is None: fatal("UnknownDeviceObject", handle_id=handle) if int(cmd["input_object_id"]) != obj["object_id"] or int(cmd["input_sequence"]) != obj["sequence"]: fatal("InputBindingMismatch", handle_id=handle, step_id=int(cmd["step_id"])) output_ring_id = int(cmd["output_ring_id"]) ring = rings[output_ring_id] if ring["direction"] != "egress": fatal("WrongRingDirection", ring_id=output_ring_id, direction=ring["direction"]) final_stage = bool(cmd.get("final_stage")) if not isinstance(model, PipelineStageTinygradModel): if final_stage: base = sum(int(word) for word in obj["words"]) + int(role.get("stage_index", 0)) token = 6 if base % 2 else 8 payload = struct.pack("