| """Decision Index engine for a d3 package loaded through d3q (text and the public vision suite). |
| |
| PYTHONPATH=kit_engine python -m decision_index run --engine d3q_engine:D3QEngine \ |
| --option model=~/models/d3-NVFP4 --out runs/d3-NVFP4 --compact |
| |
| The request path is the package's own d3_engine.D3Engine (prepare, run, respond; batch_size questions per forward |
| pass, request order), so results compare one to one with the official BF16 run. Only loading differs: |
| d3q.load.load_d3 keeps INT4/NVFP4 weights packed and runs them on the 4-bit kernels, and loads a BF16 package with |
| the stock runtime. The package's own d3q/ is used when it ships one, otherwise this project's. |
| |
| The vision runner passes ``image_root`` (row images are relative to it) and a per-row ``prompt``: the prompt is the |
| letter-readout cue for generative engines and is ignored here, since d3 reads every option's probability directly. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import os |
| import sys |
| from pathlib import Path |
|
|
| from decision_index.engines import Engine, Unsupported |
|
|
| PROJECT = Path(__file__).resolve().parent.parent |
| |
| os.environ.setdefault("TRITON_CACHE_AUTOTUNING", "1") |
|
|
|
|
| class D3QEngine(Engine): |
| name = "d3" |
| latency = ( |
| "Device-synchronized in-process request wall time including prompt rendering and tokenization (and, " |
| "for image requests, their decoding and preprocessing); one request per call, its questions in " |
| "request order in batches of batch_size; excludes model loading." |
| ) |
|
|
| def __init__(self, model=None, device="cuda:0", batch_size=8, image_root=None, **options): |
| if options: |
| raise TypeError(f"unknown engine options {sorted(options)}") |
| if not model: |
| raise ValueError("pass --option model=<package dir>") |
| package = Path(model).expanduser().resolve() |
| |
| |
| fraction = os.environ.get("D3Q_MEM_FRACTION") |
| if fraction: |
| import torch |
|
|
| torch.cuda.set_per_process_memory_fraction(float(fraction), torch.device(device)) |
| sys.path.insert(0, str(package if (package / "d3q").is_dir() else PROJECT)) |
| from d3q.load import load_d3 |
|
|
| super().__init__(model=str(package), device=device, batch_size=batch_size) |
| self.decision = load_d3(package, device=device, batch_size=int(batch_size)) |
| self.provenance = self.decision.provenance() |
| self.image_root = Path(image_root) if image_root else None |
| import torch |
|
|
| self.after_load_bytes = torch.cuda.memory_allocated(device) |
| torch.cuda.reset_peak_memory_stats(device) |
| self.device = device |
|
|
| def warmup(self): |
| super().warmup() |
| self.warmup_seconds = self.decision.warmup() |
|
|
| def runtime(self): |
| return self.decision.runtime_info() |
|
|
| def synchronize(self): |
| self.decision.synchronize() |
|
|
| def close(self): |
| |
| path = os.environ.get("D3Q_MEMORY_LOG") |
| if path: |
| import json |
|
|
| import torch |
|
|
| Path(path).write_text(json.dumps({ |
| "after_load_gib": round(self.after_load_bytes / 2**30, 2), |
| "peak_gib": round(torch.cuda.max_memory_allocated(self.device) / 2**30, 2), |
| "gpu": torch.cuda.get_device_name(self.device), |
| })) |
|
|
| def __call__(self, state, questions, images=None, videos=None, prompt=None): |
| if images and self.image_root: |
| images = [ |
| str(self.image_root / i) if isinstance(i, str) and not os.path.isabs(i) and "://" not in i else i |
| for i in images |
| ] |
| prepared = self.decision.prepare(state, questions, images, videos) |
| over = [e["message"] for e in prepared.errors.values() if e["error"] == "max_length_exceeded"] |
| if over: |
| raise Unsupported(over[0]) |
| if prepared.errors: |
| raise ValueError( |
| "invalid questions: " + "; ".join(f"{k}: {e['message']}" for k, e in prepared.errors.items()) |
| ) |
| probabilities, tokens = self.decision.run(prepared) |
| response = self.decision.respond(prepared, probabilities, tokens) |
| failed = {k: a["message"] for k, a in response["answers"].items() if "error" in a} |
| if failed: |
| raise ValueError(f"invalid model output: {failed}") |
| return response, None |
|
|
