mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-07 14:21:06 -10:00
Drafting: strata launch --drafting auto|mtp|dflash2|off, one suggestion table, and the DFlash2 drafter fetched pinned and packed on the person's PC
- strata_control/drafting.py is the one place: the choice (saved in the model's run config as "drafting"), the table card class -> suggested drafter with the measured evidence (Auto picks DFlash2 only where it was faster on every tested prompt; today no class qualifies, so Auto is the MTP everywhere), the pinned drafter (incoai/Qwen3.8-Flash-Next-DFlash2 rev a22a1744, sizes and SHA-256 of the repository), what each choice puts on the engine arguments, and the plain "Drafting: MTP (DFlash2 didn't fit)" words read from the engine's own log. - serve/server.py applies the choice at every start and reports "drafting" / "drafting_note" in /metrics. - strata launch --drafting writes it to the launch config and presets, shows it in --plan, adds the drafter to the download plan when it would be used, fetches it before the start (a failure starts the MTP and says so), and prints the engine's fallback line. - strata drafting plan|fetch (JSON/JSONL for the app). tools/dflash2_pack.py now reads safetensors with numpy alone (no torch on the person's PC), byte for byte the same file as before (checked against the earlier version), and ships in the Terminal bundle. - DFlash2 is for Qwen3.8-Flash-Next (qwen, gyro) only. The engine in this release runs it only in its CLI run, not with --serve, so while drafting.ENGINE_SERVES_DFLASH2 is False the choice starts the MTP, says "DFlash2 isn't available in this mode yet" and downloads nothing.
This commit is contained in:
Niko1221 committed
1 parent
f1b6a2144c
commit
437a60676f
12 files changed
+1197
-17
No files matched your search
@@ -107,6 +107,7 @@ PAYLOAD = (
|
||||
TOOLS_SUBSET = (
|
||||
"strata_tokenizer.py",
|
||||
"gguf_reader.py",
|
||||
"dflash2_pack.py", # the DFlash2 drafter: strata_control/drafting_fetch.py packs the download with it (numpy only)
|
||||
)
|
||||
|
||||
#: What `strata launch` and the desktop app's installer need on top of the read-only terminal: the model catalog and GPU probe
|
||||
|
||||
@@ -72,6 +72,7 @@ tls.install() # #3722: image URLs and tool servers over HTTPS are checked
|
||||
from strata_control import hipblaslt_tuning # noqa: E402
|
||||
from strata_control import vision_reader # noqa: E402
|
||||
from strata_control import draft_layer # noqa: E402
|
||||
from strata_control import drafting # noqa: E402
|
||||
from strata_control import user_args # noqa: E402
|
||||
from strata_control import gpu_roles # noqa: E402
|
||||
from serve import ready as ready_file # noqa: E402
|
||||
@@ -496,6 +497,18 @@ def start_failure_hint(log: str | None, offset: int) -> str:
|
||||
return ". mtp: " + DRAFT_HEAD_FAIL + ". " + (" ".join(said) if said else "Hint: " + DRAFT_HEAD_HINT)
|
||||
|
||||
|
||||
def start_log_text(log: str | None, offset: int, limit: int = 1 << 20) -> str:
|
||||
"""Everything this engine start has written to the log so far (at most `limit` bytes), "" when there is no log."""
|
||||
if not log:
|
||||
return ""
|
||||
try:
|
||||
with open(log, "rb") as f:
|
||||
f.seek(max(0, offset))
|
||||
return f.read(limit).decode("utf-8", "replace")
|
||||
except (OSError, ValueError):
|
||||
return ""
|
||||
|
||||
|
||||
def start_log_tail(log: str | None, offset: int, n: int = 20) -> str:
|
||||
"""#496: the last n lines this start wrote to the engine log, for the error when the engine ended before READY -
|
||||
whatever the failure, the engine's own reason is in them (people posted the traceback without the log). "" when
|
||||
@@ -1109,6 +1122,10 @@ class StrataEngine:
|
||||
if self.info.get("engine"):
|
||||
self.info["version"] = str(self.info["engine"])
|
||||
self.terminal_stages = read_engine_stage_log(log, log_start) if log else []
|
||||
try: # which drafter runs, and in plain words why not the one that was asked for
|
||||
self.info.update(drafting.engine_info(args, start_log_text(log, log_start), drafting.take_note()))
|
||||
except Exception: # noqa: BLE001 - a note must never stop a start
|
||||
pass
|
||||
self.ended = False # READY: alive from here (restart() set it True, #344)
|
||||
# the engine's stdout on a thread, so a request can wait with a timeout (heartbeats, cancel checks)
|
||||
self.lines: queue.Queue = queue.Queue()
|
||||
@@ -3027,6 +3044,22 @@ def layer_split_of(cfg: dict) -> bool:
|
||||
return len(gpu_list(cfg)) > 1 and "--peer-device" not in cfg["args"]
|
||||
|
||||
|
||||
_CARD_CLASS: list = []
|
||||
|
||||
|
||||
def drafting_card_class() -> str:
|
||||
"""The class of this computer's graphics card for the drafter table (strata_control/drafting.py), looked up once."""
|
||||
if not _CARD_CLASS:
|
||||
archs: list[str] = []
|
||||
try:
|
||||
from strata_control import engine_component
|
||||
archs = list(engine_component.cuda_archs_present()) + list(engine_component.hip_archs_present())
|
||||
except Exception: # noqa: BLE001 - an unknown card is the "other" class, never an error
|
||||
pass
|
||||
_CARD_CLASS.append(drafting.card_class(archs))
|
||||
return _CARD_CLASS[0]
|
||||
|
||||
|
||||
def engine_args(cfg: dict) -> list[str]:
|
||||
"""The engine's arguments: the config's, and with several GPUs the layer split across them ("layer_split" in the
|
||||
config: "auto" by default, or the first layer of each later GPU's share, e.g. "18" or "16,32"; see
|
||||
@@ -3036,6 +3069,17 @@ def engine_args(cfg: dict) -> list[str]:
|
||||
# so a setup run or the app rewriting "args" cannot undo them (#3216 #3215 #3270)
|
||||
if cfg.get("user_args"):
|
||||
args = user_args.merge(args, cfg["user_args"])
|
||||
# "drafting": auto | mtp | dflash2 | off (strata_control/drafting.py; the Models screen and `strata launch --drafting` write it): which drafter
|
||||
# guesses ahead. A choice that cannot be applied starts the MTP and says why in the engine's /metrics ("drafting_note").
|
||||
if cfg.get("drafting") is not None or drafting.suggest(drafting_card_class()) != "mtp":
|
||||
shaped = {**cfg, "args": args}
|
||||
decision = drafting.apply(shaped, klass=drafting_card_class(), data_dir=drafting.data_dir_of(args))
|
||||
args = shaped["args"]
|
||||
drafting.remember_note(decision.note)
|
||||
if decision.note:
|
||||
print(f"[strata] drafting: {decision.note}", flush=True)
|
||||
elif decision.flags:
|
||||
print(f"[strata] drafting: DFlash2 ({' '.join(decision.flags[:2])})", flush=True)
|
||||
# "mtp": false (the Models settings, `strata launch --mtp off`, or by hand): no draft layer, whatever the args say. A setup run
|
||||
# that writes --mtp back into the args cannot undo it, because this is applied at every start.
|
||||
if draft_layer.is_off(cfg):
|
||||
|
||||
@@ -169,6 +169,8 @@ def parser() -> argparse.ArgumentParser:
|
||||
models_move.add_argument("--plan", action="store_true", help="say what would happen and change nothing")
|
||||
models_move.add_argument("--cancel-file", metavar="PATH", help="stop (safely) when this file appears")
|
||||
models_move.add_argument("--output", choices=("human", "json", "jsonl"), default="human")
|
||||
from strata_cli import drafting_cmd
|
||||
drafting_cmd.add_parser(sp)
|
||||
engine = sp.add_parser("engine", help="engine tuning and component actions")
|
||||
engine_sp = engine.add_subparsers(dest="engine_command", required=True)
|
||||
calibrate = engine_sp.add_parser("calibrate", help="measure, review and apply engine tuning")
|
||||
@@ -2139,6 +2141,9 @@ def _run(a: argparse.Namespace) -> int:
|
||||
return _models(a)
|
||||
if a.command == "engine":
|
||||
return _calibration(a)
|
||||
if a.command == "drafting":
|
||||
from strata_cli import drafting_cmd as _drafting_cmd
|
||||
return _drafting_cmd.run(a)
|
||||
if a.command == "bench":
|
||||
return _bench(a)
|
||||
if a.command == "timeskip":
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
"""``strata drafting plan|fetch``: the DFlash2 drafter file (strata_control/drafting_fetch.py). The app calls it with ``--output jsonl``."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from strata_control import acquisition, drafting, drafting_fetch
|
||||
|
||||
|
||||
def add_parser(sp) -> None:
|
||||
drafting_p = sp.add_parser("drafting", help="the DFlash2 drafter file: what it would download, and fetch it")
|
||||
d_sp = drafting_p.add_subparsers(dest="drafting_command", required=True)
|
||||
for name, text in (("plan", "say what fetching would download; nothing is changed"),
|
||||
("fetch", "download the drafter (pinned, checked) and pack it for the engine")):
|
||||
one = d_sp.add_parser(name, help=text)
|
||||
one.add_argument("--data-dir", metavar="DIR", help="the model data folder (default: the one Strata uses)")
|
||||
one.add_argument("--output", choices=("human", "json", "jsonl"), default="human")
|
||||
|
||||
|
||||
def _data_dir(a) -> Path:
|
||||
if a.data_dir:
|
||||
return Path(a.data_dir)
|
||||
from strata_cli.launch.catalog import Catalog
|
||||
return Catalog().data_dir()
|
||||
|
||||
|
||||
def run(a) -> int:
|
||||
machine = a.output in ("json", "jsonl")
|
||||
folder = _data_dir(a)
|
||||
if a.drafting_command == "plan":
|
||||
result = drafting_fetch.plan(folder)
|
||||
if machine:
|
||||
print(json.dumps({"schema_version": 1, "ok": True, **result}, ensure_ascii=False))
|
||||
else:
|
||||
print(result["text"])
|
||||
return 0
|
||||
|
||||
def show(step: str, done: int, total: int) -> None:
|
||||
if a.output == "jsonl":
|
||||
print(json.dumps({"event": "progress", "step": step, "done": done, "total": total}, ensure_ascii=False), flush=True)
|
||||
elif not machine and total:
|
||||
print(f"\r{step} {done / 1e9:.2f} of {total / 1e9:.2f} GB", end="", flush=True)
|
||||
|
||||
try:
|
||||
out = drafting_fetch.ensure(folder, progress=show)
|
||||
except (drafting_fetch.DraftingError, acquisition.AcquisitionRefused) as exc:
|
||||
if machine:
|
||||
print(json.dumps({"schema_version": 1, "ok": False, "error": {"code": "drafting_failed", "message": str(exc)}}, ensure_ascii=False))
|
||||
else:
|
||||
print(f"\nstrata: {exc}", file=sys.stderr)
|
||||
return 5
|
||||
if machine:
|
||||
print(json.dumps({"schema_version": 1, "ok": True, "event": "result", "pack": str(out)}, ensure_ascii=False))
|
||||
else:
|
||||
print(f"\nThe DFlash2 drafter is ready: {out}")
|
||||
return 0
|
||||
@@ -99,6 +99,10 @@ def add_parsers(sp: argparse._SubParsersAction) -> list[argparse.ArgumentParser]
|
||||
g.add_argument("--mtp", type=_choice("a draft layer setting", S.MTP_MODES), metavar="on|off",
|
||||
help="the draft layer that makes answers faster. off starts without it: about 1 GB less graphics memory and less RAM, "
|
||||
"slower answers; use it when the engine stops at start on a small card or a PC with little RAM. Default: on")
|
||||
g.add_argument("--drafting", type=_choice("a drafting setting", S.DRAFTING_MODES), metavar="auto|mtp|dflash2|off",
|
||||
help="which drafter guesses words ahead to make answers faster: auto (recommended: Strata picks from what it measured on cards like yours; "
|
||||
"that is the MTP today), mtp, dflash2 (Qwen3.8-Flash-Next only; downloads a 1.2 GB drafter once) or off. If DFlash2 cannot run "
|
||||
"(not enough graphics memory, or a model it is not for) the MTP is used and the start says why")
|
||||
g.add_argument("--gpu", metavar="LIST", help="graphics card numbers to use, e.g. 0 or 0,1 (default: the engine's choice)")
|
||||
g.add_argument("--gpu-roles", metavar="LIST", help="a role for each graphics card: main, helper or off, e.g. 1=main,0=helper (cards you do not name are off). "
|
||||
"Main does the last steps of every answer; a helper holds part of the model")
|
||||
@@ -216,6 +220,8 @@ def settings_from_args(a: argparse.Namespace, catalog=None) -> tuple[S.LaunchSet
|
||||
s.parallel = a.parallel
|
||||
if a.mtp:
|
||||
s.mtp = a.mtp
|
||||
if a.drafting:
|
||||
s.drafting = a.drafting
|
||||
if a.gpu and getattr(a, "gpu_roles", None):
|
||||
raise ValueError("Use --gpu (which cards) or --gpu-roles (a role for each card), not both.")
|
||||
if a.gpu:
|
||||
@@ -619,7 +625,7 @@ def run_launch(a: argparse.Namespace) -> int:
|
||||
s, notices = settings_from_args(a, catalog)
|
||||
if not machine:
|
||||
_interrupted_move_notice()
|
||||
have_choices = bool(a.preset or a.model or a.model_path or a.context or a.vision or a.parallel or a.access or a.tune)
|
||||
have_choices = bool(a.preset or a.model or a.model_path or a.context or a.vision or a.parallel or a.access or a.tune or a.drafting)
|
||||
interactive = sys.stdin.isatty() and sys.stdout.isatty() and not a.no_tui and not machine and not a.plan
|
||||
plain_questions = False
|
||||
if not have_choices and not interactive and not a.plan:
|
||||
@@ -672,6 +678,8 @@ def run_launch(a: argparse.Namespace) -> int:
|
||||
_say(f" model {m.get('title') or f'{s.family} {s.model}'}{size}")
|
||||
_say(f" context {S.context_label(s.context)} pictures {PICTURE_WORDS.get(s.vision or 'off', s.vision)} parallel {s.parallel or 'fit to your card'}")
|
||||
_say(f" download {(plan.get('download') or {}).get('text', '')}")
|
||||
if plan.get("drafting"):
|
||||
_say(f" drafting {plan['drafting'].replace('Drafting: ', '', 1)}")
|
||||
_say(f" tuning {s.tune} access {S.ACCESS_WORDS[s.access]}, port {s.port}"
|
||||
+ (" API key on" if s.api_key else ""))
|
||||
for n in notes:
|
||||
|
||||
@@ -356,6 +356,8 @@ def derive_config(base: dict, s: S.LaunchSettings, *, bool_flags: frozenset[str]
|
||||
cfg["api_key"] = "none" # served without a key on purpose (the server warns once at start)
|
||||
else:
|
||||
cfg.pop("api_key", None)
|
||||
if s.drafting: # strata_control/drafting.py: serve/server.py applies it at every start
|
||||
cfg["drafting"] = s.drafting
|
||||
cfg.setdefault("exe", "component:strata-engine")
|
||||
cfg["open_browser"] = False # strata launch never opens a browser by itself (only `--open` does, after the start)
|
||||
if isinstance(cfg.get("args"), list):
|
||||
@@ -1042,6 +1044,40 @@ class Runner:
|
||||
self.emit("model", f"These files are {fam} {model}.", None)
|
||||
|
||||
# ---- 3. config and profile
|
||||
def drafting_inputs(self, cfg: dict) -> dict:
|
||||
"""What strata_control/drafting.py needs to decide for this launch: the choice, the card class, the model and the data folder."""
|
||||
from strata_control import drafting, engine_component
|
||||
try:
|
||||
archs = list(engine_component.cuda_archs_present()) + list(engine_component.hip_archs_present())
|
||||
except Exception: # noqa: BLE001 - an unknown card is the "other" class
|
||||
archs = []
|
||||
mode = self.s.drafting if self.s.drafting is not None else cfg.get("drafting")
|
||||
return {"mode": mode, "klass": drafting.card_class(archs), "model_name": cfg.get("model_name"), "family": self.s.family,
|
||||
"data_dir": drafting.data_dir_of(cfg.get("args") or []) or (Path(self.s.data_dir) if self.s.data_dir else None)}
|
||||
|
||||
def ensure_drafter(self, base_path: Path) -> None:
|
||||
"""Fetch the DFlash2 drafter when this launch will use it and it is not here yet. Never stops the launch: if it cannot be had, the model
|
||||
starts with the MTP and the person is told (recommend, never force)."""
|
||||
from strata_control import drafting, drafting_fetch
|
||||
try:
|
||||
want = self.drafting_inputs(read_config(base_path))
|
||||
if want["data_dir"] is None or not drafting.needs_pack(**want):
|
||||
return
|
||||
self.emit("model", "Downloading the DFlash2 drafter (1.2 GB, once)…", 0)
|
||||
|
||||
def show(step: str, done: int, total: int) -> None:
|
||||
self.check_cancel()
|
||||
self.emit("model", f"{step} {done / 1e9:.2f} of {total / 1e9:.2f} GB" if total > 1 and done < total else step,
|
||||
100.0 * done / total if total > 0 else None)
|
||||
drafting_fetch.ensure(want["data_dir"], progress=show, cancel=self.cancel)
|
||||
except Cancelled:
|
||||
raise
|
||||
except Exception as e: # noqa: BLE001 - the MTP still runs
|
||||
text = f"{e} The model starts with the MTP instead." if isinstance(e, drafting_fetch.DraftingError) else \
|
||||
"The DFlash2 drafter could not be downloaded. The model starts with the MTP instead."
|
||||
self.notices.append(text)
|
||||
self.emit("model", text, None, level="warn")
|
||||
|
||||
def build_config(self, base_path: Path) -> Path:
|
||||
from strata_control import tune_knobs
|
||||
base = read_config(base_path)
|
||||
@@ -1205,6 +1241,21 @@ class Runner:
|
||||
text, nxt = explain_engine_exit(r, self.config_path)
|
||||
raise LaunchError(text, nxt, 6)
|
||||
self.emit("start", "The model is loaded and answering.", 100)
|
||||
self.report_drafting()
|
||||
|
||||
def report_drafting(self) -> None:
|
||||
"""Say which drafter runs when the engine fell back from the one that was asked for ("Drafting: MTP (DFlash2 didn't fit)"). Best effort."""
|
||||
try:
|
||||
import urllib.request
|
||||
headers = {"Authorization": f"Bearer {self.s.api_key}"} if self.s.api_key else {}
|
||||
with urllib.request.urlopen(urllib.request.Request(f"http://127.0.0.1:{int(self.s.port)}/metrics", headers=headers), timeout=4) as resp:
|
||||
info = (json.loads(resp.read().decode("utf-8", "replace")).get("engine") or {})
|
||||
except Exception: # noqa: BLE001 - the line is a convenience
|
||||
return
|
||||
note = str(info.get("drafting_note") or "")
|
||||
if note:
|
||||
self.notices.append(note)
|
||||
self.emit("start", note, None)
|
||||
|
||||
# ---- 6. access
|
||||
def open_access(self) -> str | None:
|
||||
@@ -1262,6 +1313,15 @@ class Runner:
|
||||
parts.append({"kind": "mmproj", "label": "the picture reader", "gb": MMPROJ_GB})
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
from strata_control import drafting
|
||||
_b, _cfg = self.installed_state()
|
||||
want = self.drafting_inputs(_cfg or {})
|
||||
want["data_dir"] = want["data_dir"] or self.existing_data_dir() or self.catalog.data_dir()
|
||||
if drafting.needs_pack(**want): # chosen DFlash2 (or Auto picking it) and the drafter is not here: it is part of the download
|
||||
parts.append({"kind": "dflash2", "label": "the DFlash2 drafter", "gb": drafting.DOWNLOAD_GB})
|
||||
except Exception: # noqa: BLE001 - the plan still shows the rest
|
||||
pass
|
||||
gap = self.ready_pack_gap() if st.own_files else None
|
||||
if gap:
|
||||
# a ready-made pack (NVFP4) is taken from setup's own models folder, not from where the person's files are (round 8, N2: the plan
|
||||
@@ -1292,6 +1352,8 @@ class Runner:
|
||||
text += " To skip the picture reader (the model then answers text only), add --vision off; or put its mmproj file next to your model files."
|
||||
if any(p["kind"] == "mtp" for p in parts):
|
||||
text += " To skip the draft layer (a little slower, less memory), add --mtp off."
|
||||
if any(p["kind"] == "dflash2" for p in parts):
|
||||
text += " To skip the DFlash2 drafter, add --drafting mtp."
|
||||
if not engine:
|
||||
text += " The engine for your graphics card is downloaded first (from the update source; its size is shown when it starts)."
|
||||
return {"gb": gb, "parts": parts, "engine_missing": not engine, "own_files": st.own_files, "text": text}
|
||||
@@ -1356,7 +1418,18 @@ class Runner:
|
||||
model = {"title": title, "family": s.family, "model": s.model, "model_gb": row.download_gb if row is not None else None,
|
||||
"own_files": bool(s.model_path), "folder": s.model_path}
|
||||
return {"profile": self.profile, "settings": s.to_dict(with_key=False), "warnings": warns,
|
||||
"flags": S.flags_text(s), "steps": list(STAGES), "model": model, "download": download}
|
||||
"flags": S.flags_text(s), "steps": list(STAGES), "model": model, "download": download, "drafting": self.drafting_plan_line()}
|
||||
|
||||
def drafting_plan_line(self) -> str:
|
||||
"""The drafting line of --plan: what drafts, or why not what was asked for. Nothing is changed or downloaded."""
|
||||
from strata_control import drafting
|
||||
try:
|
||||
base, cfg = self.installed_state()
|
||||
want = self.drafting_inputs(cfg or {})
|
||||
return drafting.summary_line(drafting.decide(want["mode"], klass=want["klass"], model_name=want["model_name"], family=want["family"],
|
||||
data_dir=want["data_dir"]))
|
||||
except Exception: # noqa: BLE001
|
||||
return ""
|
||||
|
||||
def run(self, *, start: bool = True) -> Result:
|
||||
try:
|
||||
@@ -1365,6 +1438,8 @@ class Runner:
|
||||
self.check_cancel()
|
||||
base = self.ensure_model()
|
||||
self.check_cancel()
|
||||
self.ensure_drafter(base)
|
||||
self.check_cancel()
|
||||
self.emit("config", "Writing this launch's settings…")
|
||||
cfg = self.build_config(base)
|
||||
self.ensure_profile(cfg)
|
||||
|
||||
@@ -26,6 +26,8 @@ ACCESS_WORDS = {"local": "This computer only", "lan": "My local network", "cloud
|
||||
TUNE_MODES = ("quick", "thorough", "skip")
|
||||
VISION_MODES = ("off", "on", "gpu", "cpu")
|
||||
MTP_MODES = ("on", "off")
|
||||
#: strata_control/drafting.py MODES (tests/terminal/test_drafting.py keeps the two the same)
|
||||
DRAFTING_MODES = ("auto", "mtp", "dflash2", "off")
|
||||
KV_MODES = ("int8", "q4_0", "k8v4", "fp16")
|
||||
DEFAULT_PORT = 8080
|
||||
|
||||
@@ -87,6 +89,7 @@ class LaunchSettings:
|
||||
kv: str | None = None # None = int8 above 8K, as setup
|
||||
vision: str | None = None # None = the default (on when the model can read pictures and the engine has the reader), off | on | gpu | cpu
|
||||
parallel: int | None = None # None = fit to the card
|
||||
drafting: str | None = None # None = Auto (recommended): the table in strata_control/drafting.py; mtp | dflash2 | off choose one. Saved in the launch config
|
||||
mtp: str | None = None # None = as installed (draft layer on); "off" starts without it (less graphics memory and RAM, slower); "on" overrides an installed "off"
|
||||
tune: str = "skip" # quick | thorough | skip
|
||||
tune_allow_small: bool = False # the "answer-changing" tick, off by default, as in the app
|
||||
@@ -287,6 +290,8 @@ def flags_text(s: LaunchSettings, *, preset: str | None = None) -> str:
|
||||
parts.append(f"--parallel {s.parallel}")
|
||||
if s.mtp:
|
||||
parts.append(f"--mtp {s.mtp}")
|
||||
if s.drafting:
|
||||
parts.append(f"--drafting {s.drafting}")
|
||||
parts.append(f"--tune {s.tune}")
|
||||
if s.tune_allow_small:
|
||||
parts.append("--tune-allow-small")
|
||||
|
||||
@@ -0,0 +1,285 @@
|
||||
"""Which drafter a model uses to guess ahead: MTP, DFlash2 or none. One place, read by the Terminal and the app.
|
||||
|
||||
The choice is the person's, and it is saved in the model's run config as ``"drafting"``: ``auto`` (the default), ``mtp``, ``dflash2`` or
|
||||
``off``. ``strata launch --drafting ...`` writes it, the app's Models screen writes it, and ``serve/server.py`` applies it at every start
|
||||
(like ``"mtp": false``, strata_control/draft_layer.py), so a setup run that rewrites the engine arguments cannot undo it.
|
||||
|
||||
mtp the engine's default drafter, the model's own draft layer (no extra flags)
|
||||
dflash2 ``--draft dflash2 --draft-model <pack>``: the DFlash2 block drafter (incoai/Qwen3.8-Flash-Next-DFlash2) for Qwen3.8-Flash-Next
|
||||
off no draft layer (the same as ``"mtp": false``: the engine still looks ahead in the text with ``--spec``)
|
||||
auto the table below says which of mtp / dflash2 this card class gets
|
||||
|
||||
Recommend, never force: a choice that cannot be applied (DFlash2 for another model, no drafter file yet) does not stop the start. The
|
||||
model starts with the MTP and ``Decision.note`` says why in plain words.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
MODES = ("auto", "mtp", "dflash2", "off")
|
||||
DEFAULT_MODE = "auto"
|
||||
LABELS = {"auto": "Auto (recommended)", "mtp": "MTP", "dflash2": "DFlash2", "off": "Off"}
|
||||
DRAFTERS = ("mtp", "dflash2", "off")
|
||||
|
||||
# --------------------------------------------------------------------------------------------------- the suggestion table
|
||||
#
|
||||
# "Suggest DFlash2 where the win is 100%": Auto picks DFlash2 only on a card class where it measured faster than the MTP on EVERY tested
|
||||
# prompt (``faster_everywhere``). A class goes here when it was measured; a row with ``faster_everywhere`` False is MTP whatever the average
|
||||
# says. Fill this in as measurements come in (and change nothing else: the Terminal and the app both read it, tests/terminal/test_drafting.py
|
||||
# checks that a DFlash2 row has the evidence).
|
||||
#
|
||||
# Evidence so far (phaseA-tests/notes-dflash2-p1.md, -p02.md, -ports.md; DFlash2 against MTP, tokens per second, same prompts and flags):
|
||||
# strix_halo Radeon 8060S (gfx1151): faster on chat, code and json, SLOWER on tool_call (-6%).
|
||||
# blackwell RTX 5070 (sm_120, 12 GB): faster on code, SLOWER on chat (-1%).
|
||||
# rdna4 Radeon AI PRO R9700 (gfx1201): slower on every prompt.
|
||||
# ampere RTX 3060 (sm_86, 12 GB): slower on every prompt.
|
||||
# other not measured.
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Suggestion:
|
||||
drafter: str # "mtp" | "dflash2"
|
||||
faster_everywhere: bool # DFlash2 measured faster than the MTP on every tested prompt of this class
|
||||
evidence: str # one line: what was measured
|
||||
|
||||
|
||||
SUGGESTED: dict[str, Suggestion] = {
|
||||
"strix_halo": Suggestion("mtp", False, "gfx1151: faster on chat, code and json; slower on tool_call (-6%)"),
|
||||
"blackwell": Suggestion("mtp", False, "sm_120 (RTX 5070, 12 GB): faster on code; slower on chat (-1%)"),
|
||||
"rdna4": Suggestion("mtp", False, "gfx1201 (R9700): slower on every prompt"),
|
||||
"ampere": Suggestion("mtp", False, "sm_86 (RTX 3060, 12 GB): slower on every prompt"),
|
||||
"other": Suggestion("mtp", False, "not measured on this kind of card"),
|
||||
}
|
||||
|
||||
#: GPU architecture (what nvidia-smi's compute capability and the kernel driver report) -> card class
|
||||
ARCH_CLASS = (
|
||||
(re.compile(r"^gfx1151$"), "strix_halo"),
|
||||
(re.compile(r"^gfx120[01]$"), "rdna4"),
|
||||
(re.compile(r"^sm_12\d$"), "blackwell"),
|
||||
(re.compile(r"^sm_86$"), "ampere"),
|
||||
)
|
||||
|
||||
|
||||
def card_class(archs: Any) -> str:
|
||||
"""The class of the first card whose architecture the table knows (``sm_120``, ``gfx1151``...); ``other`` for an unknown or empty list."""
|
||||
for arch in ([archs] if isinstance(archs, str) else list(archs or [])):
|
||||
for pattern, name in ARCH_CLASS:
|
||||
if pattern.match(str(arch).strip().lower()):
|
||||
return name
|
||||
return "other"
|
||||
|
||||
|
||||
def suggest(klass: str) -> str:
|
||||
"""What Auto picks on this card class: DFlash2 only where it was faster on every tested prompt, MTP otherwise."""
|
||||
row = SUGGESTED.get(klass) or SUGGESTED["other"]
|
||||
return "dflash2" if row.drafter == "dflash2" and row.faster_everywhere else "mtp"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------------------- the drafter file
|
||||
#
|
||||
# incoai/Qwen3.8-Flash-Next-DFlash2 (Apache-2.0), pinned to one revision. The safetensors file is 1.2 GB; tools/dflash2_pack.py turns it into the
|
||||
# engine's runtime file (Q8 pack, about 0.76 GB). The sizes and hashes are the repository's own for that revision: a download is refused
|
||||
# if either differs.
|
||||
|
||||
REPO = "incoai/Qwen3.8-Flash-Next-DFlash2"
|
||||
REVISION = "a22a1744e85c17364dcc54a6a420af7135d4c7f6"
|
||||
LICENSE = "Apache-2.0"
|
||||
BASE_URL = f"https://huggingface.co/{REPO}/resolve/{REVISION}/"
|
||||
FILES = (
|
||||
("config.json", 1236, "2409575986f9ea24d455343eb469c0e6940def5edab1d1be503bb8eb4cc3b1b9"),
|
||||
("model.safetensors", 1206140008, "d6723f2a359eaf33ec0c95f22ae8edda896774020c552296a85e1913b4d3186f"),
|
||||
)
|
||||
DOWNLOAD_GB = 1.2 # what is fetched (the safetensors file)
|
||||
PACK_GB = 0.76 # what stays on disk (the source file is removed after packing)
|
||||
PACK_NAME = "qwen3.8-flash-next-dflash2.q8.pack"
|
||||
VRAM_GB = 1.2 # what the drafter needs in graphics memory (the engine's own estimate)
|
||||
FOLDER = "dflash2" # under the model data folder, next to ``mtp``
|
||||
|
||||
#: model families (strata launch --family) the drafter was trained for and measured on: Qwen3.8-Flash-Next and its Gyro build.
|
||||
SUPPORTED_FAMILIES = ("qwen", "gyro")
|
||||
_SUPPORTED_NAME = re.compile(r"^qwen3\.8-flash-next(-gyro-[a-z0-9._-]+|-(q|iq|tq)\d[a-z0-9_.]*)?$", re.IGNORECASE)
|
||||
|
||||
|
||||
def supported_model(model_name: str | None = None, family: str | None = None) -> bool:
|
||||
"""DFlash2 is for Qwen3.8-Flash-Next only: by family when it is known (the Terminal), else by the run config's model name."""
|
||||
if family is not None:
|
||||
return family.lower() in SUPPORTED_FAMILIES
|
||||
return bool(model_name and _SUPPORTED_NAME.match(str(model_name)))
|
||||
|
||||
|
||||
def pack_path(data_dir: str | Path) -> Path:
|
||||
return Path(data_dir) / FOLDER / PACK_NAME
|
||||
|
||||
|
||||
def pack_ready(data_dir: str | Path | None) -> bool:
|
||||
try:
|
||||
return bool(data_dir) and pack_path(data_dir).is_file() and pack_path(data_dir).stat().st_size > (1 << 20)
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------------------- a decision
|
||||
|
||||
|
||||
@dataclass
|
||||
class Decision:
|
||||
asked: str # what the person chose (auto | mtp | dflash2 | off)
|
||||
drafter: str # what the engine is started with: mtp | dflash2 | off
|
||||
note: str = "" # plain words when the choice could not be applied; "" otherwise
|
||||
flags: tuple[str, ...] = () # engine flags added
|
||||
|
||||
|
||||
def normalize(value: Any) -> str:
|
||||
text = str(value if value is not None else DEFAULT_MODE).strip().lower()
|
||||
return text if text in MODES else DEFAULT_MODE
|
||||
|
||||
|
||||
#: Does the engine in this release run DFlash2 when it serves (the app and `strata launch`)? In f-023-int src/program/generate.cpp, `--draft dflash2`
|
||||
#: is "for one card and the CLI run (not --serve, --batch or a layer split)": a served model prints "--draft dflash2 is off: ... the MTP drafts
|
||||
#: instead". So today the choice cannot take effect when serving, and nothing is downloaded or packed for it (1.2 GB for no effect would be a bad
|
||||
#: surprise); the start says so ("Drafting: MTP (DFlash2 isn't available in this mode yet)"). Flip this when the engine serves it; everything
|
||||
#: else (flags, download, pack, fallback text) is in place and tested with it on.
|
||||
ENGINE_SERVES_DFLASH2 = False
|
||||
|
||||
NOT_SUPPORTED = "DFlash2 is for Qwen3.8-Flash-Next only; this model uses the MTP."
|
||||
NOT_SERVED = "DFlash2 isn't available in this mode yet; this start uses the MTP."
|
||||
NOT_DOWNLOADED = "The DFlash2 drafter is not downloaded yet; this start uses the MTP."
|
||||
|
||||
|
||||
def decide(mode: Any, *, klass: str, model_name: str | None = None, family: str | None = None, data_dir: str | Path | None = None) -> Decision:
|
||||
"""The drafter a start uses. Never raises: a choice that cannot be applied becomes MTP with a note."""
|
||||
asked = normalize(mode)
|
||||
if asked == "off":
|
||||
return Decision(asked, "off")
|
||||
want = suggest(klass) if asked == "auto" else asked
|
||||
if want != "dflash2":
|
||||
return Decision(asked, "mtp")
|
||||
if not supported_model(model_name, family):
|
||||
return Decision(asked, "mtp", NOT_SUPPORTED if asked == "dflash2" else "")
|
||||
if not ENGINE_SERVES_DFLASH2:
|
||||
return Decision(asked, "mtp", NOT_SERVED if asked == "dflash2" else "")
|
||||
if not pack_ready(data_dir):
|
||||
return Decision(asked, "mtp", NOT_DOWNLOADED)
|
||||
return Decision(asked, "dflash2", flags=("--draft", "dflash2", "--draft-model", str(pack_path(data_dir))))
|
||||
|
||||
|
||||
def needs_pack(mode: Any, *, klass: str, model_name: str | None = None, family: str | None = None, data_dir: str | Path | None = None) -> bool:
|
||||
"""Would this choice use the DFlash2 drafter if its file were here, and is it not here yet? (Then it is fetched before the start.)"""
|
||||
asked = normalize(mode)
|
||||
want = suggest(klass) if asked == "auto" else asked
|
||||
return want == "dflash2" and ENGINE_SERVES_DFLASH2 and supported_model(model_name, family) and not pack_ready(data_dir)
|
||||
|
||||
|
||||
DRAFT_FLAGS = {"--draft": True, "--draft-model": True}
|
||||
|
||||
|
||||
def strip_draft_flags(args: list[Any]) -> list[str]:
|
||||
"""The arguments without ``--draft`` / ``--draft-model`` (and their values)."""
|
||||
out: list[str] = []
|
||||
items = [str(a) for a in args]
|
||||
i = 0
|
||||
while i < len(items):
|
||||
if items[i] in DRAFT_FLAGS:
|
||||
i += 2 if i + 1 < len(items) and not items[i + 1].startswith("--") else 1
|
||||
continue
|
||||
out.append(items[i])
|
||||
i += 1
|
||||
return out
|
||||
|
||||
|
||||
def data_dir_of(args: list[Any]) -> Path | None:
|
||||
"""The model data folder an argument list belongs to: `--pack <data>/packs/<tag>`, else `--mtp <data>/mtp/rt`."""
|
||||
items = [str(a) for a in args]
|
||||
for flag, up in (("--pack", 2), ("--mtp", 2)):
|
||||
if flag in items and items.index(flag) + 1 < len(items):
|
||||
path = Path(items[items.index(flag) + 1])
|
||||
if flag == "--mtp" and path.name.isdigit():
|
||||
continue
|
||||
return path.parents[up - 1] if len(path.parents) >= up else None
|
||||
return None
|
||||
|
||||
|
||||
def apply(cfg: dict[str, Any], *, klass: str, data_dir: str | Path | None = None, family: str | None = None) -> Decision:
|
||||
"""Put the config's ``"drafting"`` choice on its engine arguments (in place). ``off`` is the draft layer's own switch (``"mtp": false``)."""
|
||||
from strata_control import draft_layer
|
||||
decision = decide(cfg.get("drafting"), klass=klass, model_name=cfg.get("model_name"), family=family, data_dir=data_dir)
|
||||
if cfg.get("drafting") is None and decision.drafter == "mtp":
|
||||
return decision # nobody chose and Auto is the MTP: the arguments stay exactly as they are
|
||||
args = strip_draft_flags(cfg.get("args") or [])
|
||||
if decision.drafter == "off":
|
||||
cfg["mtp"] = False
|
||||
cfg["args"] = args
|
||||
draft_layer.apply(cfg)
|
||||
else:
|
||||
if cfg.get("drafting") is not None and normalize(cfg.get("drafting")) != "off" and cfg.get("mtp") is False:
|
||||
cfg["mtp"] = True # an explicit drafter choice overrides an older "mtp": false
|
||||
cfg["args"] = args + list(decision.flags)
|
||||
return decision
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------------------- what the engine said
|
||||
|
||||
_FALLBACK = re.compile(r"--draft dflash2(?: is off)?:\s*(?P<why>[^\n]*?)(?:\.\s|;\s|$)")
|
||||
_ON = re.compile(r"DFlash2 drafter on\b")
|
||||
|
||||
DIDNT_FIT = "DFlash2 didn't fit"
|
||||
NOT_HERE = "DFlash2 isn't available in this mode"
|
||||
|
||||
|
||||
def fallback_note(log_text: str) -> str:
|
||||
"""From one engine start's log: ``Drafting: MTP (DFlash2 didn't fit)`` when the engine fell back from DFlash2, else ``""``. The engine
|
||||
names the reason in its own words; two are told apart: no room in graphics memory, and a mode this build does not support."""
|
||||
text = str(log_text or "")
|
||||
if _ON.search(text):
|
||||
return ""
|
||||
if "--draft dflash2" not in text:
|
||||
return ""
|
||||
low = text.lower()
|
||||
if "it is for one card and the cli run" in low or "needs a native pack" in low or "it needs --mtp" in low:
|
||||
reason = NOT_HERE
|
||||
else:
|
||||
reason = DIDNT_FIT
|
||||
return f"Drafting: MTP ({reason})"
|
||||
|
||||
|
||||
def summary_line(decision: Decision, engine_note: str = "") -> str:
|
||||
"""The line the Terminal and the app show: what drafts."""
|
||||
if engine_note:
|
||||
return engine_note
|
||||
if decision.note:
|
||||
return f"Drafting: MTP ({decision.note.rstrip('.')})"
|
||||
return f"Drafting: {LABELS.get(decision.drafter, decision.drafter) if decision.drafter != 'mtp' else 'MTP'}"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------------------- what a start did
|
||||
|
||||
_pending_note = ""
|
||||
|
||||
|
||||
def remember_note(note: str) -> None:
|
||||
"""The server calls this when it builds the engine arguments: why the choice was not applied, for the engine that starts next."""
|
||||
global _pending_note
|
||||
_pending_note = note or ""
|
||||
|
||||
|
||||
def take_note() -> str:
|
||||
global _pending_note
|
||||
note, _pending_note = _pending_note, ""
|
||||
return note
|
||||
|
||||
|
||||
def engine_info(args: list[Any], log_text: str = "", note: str = "") -> dict[str, str]:
|
||||
"""The facts of one engine start for /metrics and the apps: ``drafting`` (mtp | dflash2 | off) and ``drafting_note`` (plain words, or "").
|
||||
The engine's own log wins over the choice: if it says it fell back, that is what ran."""
|
||||
items = [str(a) for a in args]
|
||||
asked_dflash2 = any(items[i] == "--draft" and i + 1 < len(items) and items[i + 1] == "dflash2" for i in range(len(items)))
|
||||
fell_back = fallback_note(log_text) if asked_dflash2 else ""
|
||||
if asked_dflash2 and not fell_back:
|
||||
return {"drafting": "dflash2", "drafting_note": ""}
|
||||
drafter = "mtp" if "--mtp" in items else "off"
|
||||
text = fell_back or (f"Drafting: MTP ({note.rstrip('.')})" if note and drafter == "mtp" else "")
|
||||
return {"drafting": drafter, "drafting_note": text}
|
||||
@@ -0,0 +1,102 @@
|
||||
"""Getting the DFlash2 drafter onto this computer: a pinned download, checked against the repository's own SHA-256, packed for the engine.
|
||||
|
||||
strata drafting plan [--data-dir DIR] what it would download (nothing is changed)
|
||||
strata drafting fetch [--data-dir DIR] download, check, pack; progress in plain words
|
||||
|
||||
Only an explicit command reaches the network (strata_control/acquisition.py): a start never downloads. The drafter is
|
||||
incoai/Qwen3.8-Flash-Next-DFlash2 at one revision (strata_control/drafting.py has the revision, the file sizes and hashes). The 1.2 GB
|
||||
safetensors file is packed (tools/dflash2_pack.py, numpy only) into the engine's 0.76 GB runtime file and then removed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
|
||||
from strata_control import acquisition, drafting, transfer
|
||||
|
||||
Progress = Callable[[str, int, int], None] # (step in plain words, bytes done, bytes total)
|
||||
|
||||
|
||||
class DraftingError(RuntimeError):
|
||||
"""The drafter could not be fetched or packed. The message is plain words and says what to do."""
|
||||
|
||||
|
||||
def plan(data_dir: str | Path) -> dict:
|
||||
"""What fetching would do, in numbers and words. ``gb`` is 0 when the drafter is already here."""
|
||||
ready = drafting.pack_ready(data_dir)
|
||||
gb = 0.0 if ready else drafting.DOWNLOAD_GB
|
||||
if ready:
|
||||
text = "The DFlash2 drafter is already on this computer."
|
||||
else:
|
||||
text = (f"The DFlash2 drafter is {drafting.DOWNLOAD_GB:.1f} GB to download (it is packed to {drafting.PACK_GB:.2f} GB for the engine and the "
|
||||
f"download is removed), from {drafting.REPO}, {drafting.LICENSE}. It needs about {drafting.VRAM_GB:.1f} GB of graphics memory; "
|
||||
"if it does not fit, the model starts with the MTP instead.")
|
||||
return {"ready": ready, "gb": gb, "pack": str(drafting.pack_path(data_dir)), "repo": drafting.REPO, "revision": drafting.REVISION, "text": text}
|
||||
|
||||
|
||||
def _packer():
|
||||
"""tools/dflash2_pack.py as a module: beside strata_control in a source checkout and in the packaged Control (TOOLS_SUBSET)."""
|
||||
path = Path(__file__).resolve().parents[1] / "tools" / "dflash2_pack.py"
|
||||
if not path.is_file():
|
||||
raise DraftingError("The packing tool is missing from this installation. Update Strata, then try again.")
|
||||
spec = importlib.util.spec_from_file_location("strata_dflash2_pack", path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules.setdefault("strata_dflash2_pack", module)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def ensure(data_dir: str | Path, *, progress: Progress | None = None, cancel: threading.Event | None = None,
|
||||
base_url: str | None = None, files=None) -> Path:
|
||||
"""Make the packed drafter exist under ``data_dir`` and return its path. A pack that is there is kept. Resumable: an interrupted
|
||||
download is continued, a file whose hash differs is deleted and the person is told. ``base_url`` / ``files`` are for tests."""
|
||||
report = progress or (lambda *_: None)
|
||||
out = drafting.pack_path(data_dir)
|
||||
if drafting.pack_ready(data_dir):
|
||||
report("The DFlash2 drafter is ready.", 1, 1)
|
||||
return out
|
||||
acquisition.require("the DFlash2 drafter", "strata drafting fetch")
|
||||
folder = out.parent
|
||||
src = folder / "src"
|
||||
src.mkdir(parents=True, exist_ok=True)
|
||||
spec = tuple(files if files is not None else drafting.FILES)
|
||||
total = sum(size for _n, size, _h in spec)
|
||||
done_before = 0
|
||||
for name, size, sha in spec:
|
||||
target = src / name
|
||||
if not (target.is_file() and target.stat().st_size == size and transfer.sha256_file(target) == sha):
|
||||
base = done_before
|
||||
partial = src / (name + ".part")
|
||||
try:
|
||||
transfer.download((base_url or drafting.BASE_URL) + name, partial, size=size, sha256=sha, cancel=cancel,
|
||||
progress=lambda d, _s, base=base: report("Downloading the DFlash2 drafter…", base + d, total))
|
||||
except transfer.ChecksumMismatch:
|
||||
raise DraftingError("The downloaded DFlash2 file did not pass the safety check, so it was deleted. Try again; if it keeps "
|
||||
"happening, use MTP for now.") from None
|
||||
except transfer.TransferError as exc:
|
||||
raise DraftingError(f"The DFlash2 download did not finish: {exc}. Try again; it continues where it stopped.") from None
|
||||
os.replace(partial, target)
|
||||
done_before += size
|
||||
report("Packing the DFlash2 drafter for the engine…", total, total)
|
||||
tmp = out.with_suffix(".tmp")
|
||||
try:
|
||||
_packer().pack(str(src / "model.safetensors"), str(src / "config.json"), str(tmp), say=lambda *_: None)
|
||||
os.replace(tmp, out)
|
||||
except DraftingError:
|
||||
raise
|
||||
except Exception as exc: # noqa: BLE001 - a plain sentence for any packing failure; the cause stays in the chain
|
||||
tmp.unlink(missing_ok=True)
|
||||
raise DraftingError("The DFlash2 drafter could not be prepared on this computer. The model keeps using the MTP.") from exc
|
||||
for name, _size, _sha in spec: # the source files are 1.2 GB of nothing the engine reads
|
||||
(src / name).unlink(missing_ok=True)
|
||||
try:
|
||||
src.rmdir()
|
||||
except OSError:
|
||||
pass
|
||||
report("The DFlash2 drafter is ready.", total, total)
|
||||
return out
|
||||
@@ -0,0 +1,412 @@
|
||||
"""Which drafter a model uses (MTP, DFlash2 or none): the suggestion table, the Terminal flag, the engine arguments, the download and the
|
||||
fallback words. strata_control/drafting.py is the one place; the Terminal and the app read it."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import http.server
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
from strata_cli.launch import settings as S
|
||||
from strata_cli.launch.runner import derive_config
|
||||
from strata_control import acquisition, drafting, drafting_fetch
|
||||
from tests.terminal.launch_support import run_cli
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
GYRO = "qwen3.8-flash-next-gyro-gyro-s"
|
||||
|
||||
|
||||
def serving_dflash2():
|
||||
"""The engine of this release runs DFlash2 only in its CLI run (drafting.ENGINE_SERVES_DFLASH2): the tests turn the switch on."""
|
||||
return mock.patch.object(drafting, "ENGINE_SERVES_DFLASH2", True)
|
||||
|
||||
|
||||
class SuggestionTable(unittest.TestCase):
|
||||
def test_every_row_is_a_real_drafter_with_evidence(self):
|
||||
for klass, row in drafting.SUGGESTED.items():
|
||||
self.assertIn(row.drafter, ("mtp", "dflash2"), klass)
|
||||
self.assertTrue(row.evidence.strip(), f"{klass}: say what was measured")
|
||||
|
||||
def test_auto_is_dflash2_only_where_it_won_on_every_tested_prompt(self):
|
||||
# The owner's rule (2026-10-22): "suggest DFlash2 where the win is 100%". Today no card class qualifies.
|
||||
for klass in drafting.SUGGESTED:
|
||||
self.assertEqual(drafting.suggest(klass), "mtp", klass)
|
||||
for klass in ("strix_halo", "blackwell", "rdna4", "ampere", "other", "never-heard-of-it"):
|
||||
self.assertEqual(drafting.suggest(klass), "mtp")
|
||||
# a row that says DFlash2 but not "faster everywhere" is still the MTP; only both together switch Auto
|
||||
with mock.patch.dict(drafting.SUGGESTED, {"blackwell": drafting.Suggestion("dflash2", False, "wins on code only")}):
|
||||
self.assertEqual(drafting.suggest("blackwell"), "mtp")
|
||||
with mock.patch.dict(drafting.SUGGESTED, {"blackwell": drafting.Suggestion("dflash2", True, "faster on chat, code, json and tool_call")}):
|
||||
self.assertEqual(drafting.suggest("blackwell"), "dflash2")
|
||||
self.assertEqual(drafting.suggest("strix_halo"), "mtp", "the other classes are unchanged")
|
||||
|
||||
def test_the_measured_losses_are_written_down(self):
|
||||
text = {k: v.evidence for k, v in drafting.SUGGESTED.items()}
|
||||
self.assertIn("-6%", text["strix_halo"])
|
||||
self.assertIn("tool_call", text["strix_halo"])
|
||||
self.assertIn("-1%", text["blackwell"])
|
||||
self.assertIn("chat", text["blackwell"])
|
||||
self.assertIn("every prompt", text["rdna4"])
|
||||
self.assertIn("every prompt", text["ampere"])
|
||||
|
||||
def test_cards_are_classified_by_their_architecture(self):
|
||||
self.assertEqual(drafting.card_class(["sm_120"]), "blackwell")
|
||||
self.assertEqual(drafting.card_class(["gfx1151"]), "strix_halo")
|
||||
self.assertEqual(drafting.card_class(["gfx1201"]), "rdna4")
|
||||
self.assertEqual(drafting.card_class(["sm_86"]), "ampere")
|
||||
self.assertEqual(drafting.card_class([]), "other")
|
||||
self.assertEqual(drafting.card_class(["sm_61"]), "other")
|
||||
self.assertEqual(drafting.card_class(["sm_61", "sm_120"]), "blackwell", "the first card the table knows")
|
||||
self.assertEqual(drafting.card_class("SM_120"), "blackwell")
|
||||
|
||||
def test_both_sides_know_the_same_modes(self):
|
||||
self.assertEqual(drafting.MODES, S.DRAFTING_MODES)
|
||||
self.assertEqual(drafting.normalize(None), "auto")
|
||||
self.assertEqual(drafting.normalize("DFlash2"), "dflash2")
|
||||
self.assertEqual(drafting.normalize("fastest"), "auto", "an unknown value is Auto, never an error")
|
||||
|
||||
|
||||
class Decisions(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.data = Path(self.tmp.name)
|
||||
self.addCleanup(self.tmp.cleanup)
|
||||
|
||||
def ready(self):
|
||||
pack = drafting.pack_path(self.data)
|
||||
pack.parent.mkdir(parents=True)
|
||||
pack.write_bytes(b"x" * (2 << 20))
|
||||
|
||||
def test_auto_and_mtp_run_the_mtp_with_no_extra_flags(self):
|
||||
with serving_dflash2():
|
||||
for mode in ("auto", "mtp", None):
|
||||
d = drafting.decide(mode, klass="blackwell", model_name=GYRO, data_dir=self.data)
|
||||
self.assertEqual((d.drafter, d.flags, d.note), ("mtp", (), ""), mode)
|
||||
|
||||
def test_off_is_off(self):
|
||||
self.assertEqual(drafting.decide("off", klass="blackwell", model_name=GYRO).drafter, "off")
|
||||
|
||||
def test_dflash2_needs_the_drafter_file_and_gets_the_engine_flags_when_it_is_there(self):
|
||||
with serving_dflash2():
|
||||
gone = drafting.decide("dflash2", klass="blackwell", model_name=GYRO, data_dir=self.data)
|
||||
self.assertEqual((gone.drafter, gone.note), ("mtp", drafting.NOT_DOWNLOADED))
|
||||
self.assertTrue(drafting.needs_pack("dflash2", klass="blackwell", model_name=GYRO, data_dir=self.data))
|
||||
self.ready()
|
||||
d = drafting.decide("dflash2", klass="blackwell", model_name=GYRO, data_dir=self.data)
|
||||
self.assertEqual(d.drafter, "dflash2")
|
||||
self.assertEqual(d.flags, ("--draft", "dflash2", "--draft-model", str(drafting.pack_path(self.data))))
|
||||
self.assertFalse(drafting.needs_pack("dflash2", klass="blackwell", model_name=GYRO, data_dir=self.data))
|
||||
|
||||
def test_dflash2_is_for_qwen38_flash_next_only(self):
|
||||
with serving_dflash2():
|
||||
self.ready()
|
||||
for name in ("qwen3.6-35b-a3b-q4", "qwen3.8-flash-next-coder-iq3_xxs", "qwen3.8-flash-next-swift-iq2_xs", None):
|
||||
d = drafting.decide("dflash2", klass="blackwell", model_name=name, data_dir=self.data)
|
||||
self.assertEqual((d.drafter, d.note), ("mtp", drafting.NOT_SUPPORTED), name)
|
||||
self.assertFalse(drafting.needs_pack("dflash2", klass="blackwell", model_name=name, data_dir=self.data))
|
||||
for name in (GYRO, "qwen3.8-flash-next-q2_0", "qwen3.8-flash-next-iq3_xxs"):
|
||||
self.assertEqual(drafting.decide("dflash2", klass="blackwell", model_name=name, data_dir=self.data).drafter, "dflash2", name)
|
||||
self.assertTrue(drafting.supported_model(family="qwen") and drafting.supported_model(family="gyro"))
|
||||
self.assertFalse(drafting.supported_model(family="coder") or drafting.supported_model(family="swift") or drafting.supported_model(family="nvfp4"))
|
||||
|
||||
def test_while_the_engine_cannot_serve_it_the_choice_says_so_and_nothing_is_downloaded(self):
|
||||
self.assertFalse(drafting.ENGINE_SERVES_DFLASH2, "src/program/generate.cpp: --draft dflash2 is off with --serve in this release")
|
||||
d = drafting.decide("dflash2", klass="blackwell", model_name=GYRO, data_dir=self.data)
|
||||
self.assertEqual((d.drafter, d.note), ("mtp", drafting.NOT_SERVED))
|
||||
self.assertFalse(drafting.needs_pack("dflash2", klass="blackwell", model_name=GYRO, data_dir=self.data))
|
||||
self.assertEqual(drafting.decide("auto", klass="blackwell", model_name=GYRO, data_dir=self.data).note, "")
|
||||
|
||||
def test_auto_follows_the_table(self):
|
||||
with serving_dflash2(), mock.patch.dict(drafting.SUGGESTED, {"blackwell": drafting.Suggestion("dflash2", True, "wins everywhere")}):
|
||||
self.ready()
|
||||
self.assertEqual(drafting.decide("auto", klass="blackwell", model_name=GYRO, data_dir=self.data).drafter, "dflash2")
|
||||
self.assertEqual(drafting.decide("auto", klass="strix_halo", model_name=GYRO, data_dir=self.data).drafter, "mtp")
|
||||
self.assertEqual(drafting.decide("auto", klass="blackwell", model_name="qwen3.6-35b-a3b-q4", data_dir=self.data).note, "",
|
||||
"Auto never nags about a model DFlash2 is not for")
|
||||
|
||||
|
||||
class EngineArguments(unittest.TestCase):
|
||||
ARGS = ["--pack", "/d/packs/gyro", "--native", "/d/m.gguf", "--mtp", "/d/mtp/rt", "--spec", "4", "--spec-min-p", "0.5"]
|
||||
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.data = Path(self.tmp.name)
|
||||
self.addCleanup(self.tmp.cleanup)
|
||||
|
||||
def cfg(self, **extra):
|
||||
return {"model_name": GYRO, "args": list(self.ARGS), **extra}
|
||||
|
||||
def test_a_config_nobody_touched_keeps_its_arguments_exactly(self):
|
||||
c = self.cfg()
|
||||
drafting.apply(c, klass="blackwell", data_dir=self.data)
|
||||
self.assertEqual(c["args"], self.ARGS)
|
||||
self.assertNotIn("mtp", c)
|
||||
|
||||
def test_dflash2_adds_the_flags_and_replaces_older_ones(self):
|
||||
pack = drafting.pack_path(self.data)
|
||||
pack.parent.mkdir(parents=True)
|
||||
pack.write_bytes(b"x" * (2 << 20))
|
||||
with serving_dflash2():
|
||||
c = self.cfg(drafting="dflash2", args=[*self.ARGS, "--draft", "mtp", "--draft-model", "/old.pack"])
|
||||
d = drafting.apply(c, klass="blackwell", data_dir=self.data)
|
||||
self.assertEqual(d.drafter, "dflash2")
|
||||
self.assertEqual(c["args"], [*self.ARGS, "--draft", "dflash2", "--draft-model", str(pack)])
|
||||
self.assertEqual(c["args"].count("--draft"), 1)
|
||||
self.assertIn("--mtp", c["args"], "the MTP stays as the fallback the engine names")
|
||||
|
||||
def test_mtp_removes_a_drafter_that_was_written_by_hand(self):
|
||||
c = self.cfg(drafting="mtp", args=[*self.ARGS, "--draft", "dflash2", "--draft-model", "/p"])
|
||||
drafting.apply(c, klass="blackwell", data_dir=self.data)
|
||||
self.assertEqual(c["args"], self.ARGS)
|
||||
|
||||
def test_off_is_the_draft_layers_own_switch(self):
|
||||
c = self.cfg(drafting="off")
|
||||
d = drafting.apply(c, klass="blackwell", data_dir=self.data)
|
||||
self.assertEqual(d.drafter, "off")
|
||||
self.assertIs(c["mtp"], False)
|
||||
for flag in ("--mtp", "--spec-min-p", "--batch-mtp", "--draft"):
|
||||
self.assertNotIn(flag, c["args"])
|
||||
self.assertEqual(c["args"][c["args"].index("--spec") + 1], "4", "prompt lookup keeps its windows")
|
||||
|
||||
def test_a_choice_that_cannot_be_applied_starts_the_mtp_and_says_why(self):
|
||||
c = self.cfg(drafting="dflash2", model_name="qwen3.8-flash-next-coder-iq3_xxs")
|
||||
with serving_dflash2():
|
||||
d = drafting.apply(c, klass="blackwell", data_dir=self.data)
|
||||
self.assertEqual((d.drafter, d.note), ("mtp", drafting.NOT_SUPPORTED))
|
||||
self.assertEqual(c["args"], self.ARGS)
|
||||
|
||||
def test_the_server_applies_it_at_every_start(self):
|
||||
from serve.server import engine_args
|
||||
pack = drafting.pack_path(Path("/d"))
|
||||
with mock.patch("serve.server.drafting_card_class", return_value="blackwell"), serving_dflash2(), \
|
||||
mock.patch.object(drafting, "pack_ready", return_value=True), mock.patch.object(drafting, "data_dir_of", return_value=Path("/d")):
|
||||
args = engine_args(self.cfg(drafting="dflash2"))
|
||||
self.assertEqual(args[-4:], ["--draft", "dflash2", "--draft-model", str(pack)])
|
||||
with mock.patch("serve.server.drafting_card_class", return_value="blackwell"):
|
||||
self.assertEqual(engine_args(self.cfg()), self.ARGS, "no choice: nothing changes")
|
||||
off = engine_args(self.cfg(drafting="off"))
|
||||
self.assertNotIn("--mtp", off)
|
||||
self.assertEqual(drafting.take_note(), "")
|
||||
engine_args(self.cfg(drafting="dflash2"))
|
||||
self.assertEqual(drafting.take_note(), drafting.NOT_SERVED, "the engine that starts next is told why")
|
||||
|
||||
def test_data_folder_of_an_argument_list(self):
|
||||
self.assertEqual(drafting.data_dir_of(["--pack", "/d/packs/x"]), Path("/d"))
|
||||
self.assertEqual(drafting.data_dir_of(["--mtp", "/d/mtp/rt"]), Path("/d"))
|
||||
self.assertIsNone(drafting.data_dir_of(["--mtp", "4"]))
|
||||
self.assertIsNone(drafting.data_dir_of([]))
|
||||
|
||||
|
||||
class WhatTheEngineSaid(unittest.TestCase):
|
||||
ARGS = ["--mtp", "/m", "--draft", "dflash2", "--draft-model", "/p"]
|
||||
|
||||
def test_no_room_in_graphics_memory(self):
|
||||
log = ("strata generate: --draft dflash2: could not allocate the drafter's buffers. DFlash2 needs about 1.2 GiB of VRAM (the pack 0.76 GiB, "
|
||||
"its buffers and K/V ring 0.25 GiB, a draft head 0.2 GiB); start with --vram-reserve-mib 1900 to keep it. The MTP drafts instead.\n")
|
||||
info = drafting.engine_info(self.ARGS, log)
|
||||
self.assertEqual(info, {"drafting": "mtp", "drafting_note": "Drafting: MTP (DFlash2 didn't fit)"})
|
||||
|
||||
def test_a_mode_this_build_does_not_run(self):
|
||||
log = "strata generate: --draft dflash2 is off: it is for one card and the CLI run (not --serve, --batch or a layer split) in this build; the MTP drafts instead\n"
|
||||
self.assertEqual(drafting.engine_info(self.ARGS, log)["drafting_note"], "Drafting: MTP (DFlash2 isn't available in this mode)")
|
||||
|
||||
def test_it_ran(self):
|
||||
log = "strata generate: DFlash2 drafter on (/p, graphs captured in 3 ms), the MTP layer stays as the fallback\n"
|
||||
self.assertEqual(drafting.engine_info(self.ARGS, log), {"drafting": "dflash2", "drafting_note": ""})
|
||||
self.assertEqual(drafting.fallback_note(log), "")
|
||||
|
||||
def test_choices_that_never_reached_the_engine_keep_their_own_words(self):
|
||||
self.assertEqual(drafting.engine_info(["--mtp", "/m"], "", drafting.NOT_SUPPORTED)["drafting_note"],
|
||||
"Drafting: MTP (DFlash2 is for Qwen3.8-Flash-Next only; this model uses the MTP)")
|
||||
self.assertEqual(drafting.engine_info(["--mtp", "/m"]), {"drafting": "mtp", "drafting_note": ""})
|
||||
self.assertEqual(drafting.engine_info(["--spec", "4"]), {"drafting": "off", "drafting_note": ""})
|
||||
|
||||
def test_summary_line(self):
|
||||
self.assertEqual(drafting.summary_line(drafting.Decision("auto", "mtp")), "Drafting: MTP")
|
||||
self.assertEqual(drafting.summary_line(drafting.Decision("dflash2", "dflash2")), "Drafting: DFlash2")
|
||||
self.assertEqual(drafting.summary_line(drafting.Decision("off", "off")), "Drafting: Off")
|
||||
self.assertEqual(drafting.summary_line(drafting.Decision("dflash2", "mtp", drafting.NOT_DOWNLOADED)),
|
||||
"Drafting: MTP (The DFlash2 drafter is not downloaded yet; this start uses the MTP)")
|
||||
|
||||
|
||||
class TerminalFlag(unittest.TestCase):
|
||||
BASE = {"exe": "component:strata-engine", "args": ["--pack", "/d/packs/x", "--mtp", "/d/mtp/rt"], "port": 8080, "model_name": GYRO}
|
||||
|
||||
def test_the_flag_is_listed_planned_copyable_and_saved_in_the_launch_config(self):
|
||||
code, out, _ = run_cli(["launch", "--help"])
|
||||
self.assertIn("--drafting", out)
|
||||
self.assertIn("auto|mtp|dflash2|off", out)
|
||||
cfg = derive_config(self.BASE, S.LaunchSettings(drafting="dflash2"), bool_flags=frozenset())
|
||||
self.assertEqual(cfg["drafting"], "dflash2")
|
||||
self.assertNotIn("drafting", derive_config(self.BASE, S.LaunchSettings(), bool_flags=frozenset()), "Auto writes nothing: nobody chose")
|
||||
self.assertIn("--drafting dflash2", S.flags_text(S.LaunchSettings(drafting="dflash2")))
|
||||
self.assertNotIn("--drafting", S.flags_text(S.LaunchSettings()))
|
||||
code, out, err = run_cli(["launch", "--plan", "--drafting", "fastest"])
|
||||
self.assertNotEqual(code, 0)
|
||||
|
||||
def test_a_saved_preset_keeps_the_choice(self):
|
||||
s = S.LaunchSettings.from_dict(S.LaunchSettings(drafting="mtp").to_dict())
|
||||
self.assertEqual(s.drafting, "mtp")
|
||||
self.assertIsNone(S.LaunchSettings.from_dict({}).drafting)
|
||||
|
||||
|
||||
class _Handler(http.server.BaseHTTPRequestHandler):
|
||||
files: dict = {}
|
||||
hits: list = []
|
||||
|
||||
def do_GET(self):
|
||||
body = self.files.get(self.path.rsplit("/", 1)[-1])
|
||||
self.hits.append(self.path)
|
||||
if body is None:
|
||||
self.send_error(404)
|
||||
return
|
||||
start = 0
|
||||
rng = self.headers.get("Range")
|
||||
if rng and rng.startswith("bytes="):
|
||||
start = int(rng[6:].split("-")[0])
|
||||
self.send_response(206 if start else 200)
|
||||
if start:
|
||||
self.send_header("Content-Range", f"bytes {start}-{len(body) - 1}/{len(body)}")
|
||||
self.send_header("Content-Length", str(len(body) - start))
|
||||
self.end_headers()
|
||||
self.wfile.write(body[start:])
|
||||
|
||||
def log_message(self, *_):
|
||||
pass
|
||||
|
||||
|
||||
class FetchAndPack(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
sys.path.insert(0, str(ROOT / "tools"))
|
||||
spec = importlib.util.spec_from_file_location("t_pack", ROOT / "tools" / "test_dflash2_pack.py")
|
||||
cls.synth = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(cls.synth)
|
||||
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.tmp.cleanup)
|
||||
self.dir = Path(self.tmp.name)
|
||||
self.synth.write_safetensors(self.dir / "m.safetensors", self.synth.tensors())
|
||||
(self.dir / "c.json").write_text(json.dumps(self.synth.CFG))
|
||||
import hashlib
|
||||
self.blobs = {"config.json": (self.dir / "c.json").read_bytes(), "model.safetensors": (self.dir / "m.safetensors").read_bytes()}
|
||||
self.files = tuple((n, len(b), hashlib.sha256(b).hexdigest()) for n, b in self.blobs.items())
|
||||
_Handler.files, _Handler.hits = dict(self.blobs), []
|
||||
self.server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), _Handler)
|
||||
threading.Thread(target=self.server.serve_forever, daemon=True).start()
|
||||
self.addCleanup(self.server.shutdown)
|
||||
self.base = f"http://127.0.0.1:{self.server.server_address[1]}/"
|
||||
self.data = self.dir / "data"
|
||||
self.steps: list = []
|
||||
# the packed test file is a few hundred kB, not the real 0.76 GB: "ready" asks for more than 1 MiB by default
|
||||
patcher = mock.patch.object(drafting, "pack_ready", lambda d: bool(d) and drafting.pack_path(d).is_file())
|
||||
patcher.start()
|
||||
self.addCleanup(patcher.stop)
|
||||
|
||||
def fetch(self, **kw):
|
||||
with mock.patch.dict("os.environ", {"STRATA_UPDATE_ALLOW_HTTP": "1"}):
|
||||
return drafting_fetch.ensure(self.data, progress=lambda *a: self.steps.append(a), base_url=self.base, files=self.files, **kw)
|
||||
|
||||
def test_the_pinned_files_are_the_repositorys_own(self):
|
||||
self.assertEqual(drafting.REVISION, "a22a1744e85c17364dcc54a6a420af7135d4c7f6")
|
||||
self.assertTrue(drafting.BASE_URL.endswith(f"/{drafting.REVISION}/"), "a revision, never a branch name")
|
||||
names = {n: (size, sha) for n, size, sha in drafting.FILES}
|
||||
self.assertEqual(names["model.safetensors"], (1206140008, "d6723f2a359eaf33ec0c95f22ae8edda896774020c552296a85e1913b4d3186f"))
|
||||
self.assertEqual(len(names["config.json"][1]), 64)
|
||||
self.assertEqual(drafting.LICENSE, "Apache-2.0")
|
||||
|
||||
def test_it_downloads_checks_packs_and_removes_the_source(self):
|
||||
out = self.fetch()
|
||||
self.assertEqual(out, drafting.pack_path(self.data))
|
||||
self.assertEqual(out.read_bytes()[:4], b"DF2P")
|
||||
self.assertFalse((out.parent / "src").exists(), "the 1.2 GB source file does not stay on disk")
|
||||
words = [s[0] for s in self.steps]
|
||||
self.assertIn("Downloading the DFlash2 drafter…", words)
|
||||
self.assertIn("Packing the DFlash2 drafter for the engine…", words)
|
||||
self.assertEqual(words[-1], "The DFlash2 drafter is ready.")
|
||||
self.assertEqual(max(s[1] for s in self.steps if s[0].startswith("Downloading")), sum(f[1] for f in self.files))
|
||||
|
||||
def test_a_drafter_that_is_here_is_not_fetched_again(self):
|
||||
self.fetch()
|
||||
_Handler.hits.clear()
|
||||
self.steps.clear()
|
||||
self.fetch()
|
||||
self.assertEqual(_Handler.hits, [])
|
||||
self.assertEqual(self.steps[0][0], "The DFlash2 drafter is ready.")
|
||||
|
||||
def test_a_file_that_does_not_match_its_hash_is_deleted_and_the_words_are_plain(self):
|
||||
bad = tuple((n, size, "0" * 64) for n, size, _h in self.files)
|
||||
with mock.patch.dict("os.environ", {"STRATA_UPDATE_ALLOW_HTTP": "1"}):
|
||||
with self.assertRaises(drafting_fetch.DraftingError) as ctx:
|
||||
drafting_fetch.ensure(self.data, base_url=self.base, files=bad)
|
||||
self.assertIn("did not pass the safety check", str(ctx.exception))
|
||||
self.assertFalse(drafting.pack_path(self.data).exists())
|
||||
self.assertEqual(list((drafting.pack_path(self.data).parent / "src").glob("*.part")), [])
|
||||
|
||||
def test_an_interrupted_download_is_continued(self):
|
||||
name, size, _sha = self.files[1]
|
||||
src = drafting.pack_path(self.data).parent / "src"
|
||||
src.mkdir(parents=True)
|
||||
(src / (name + ".part")).write_bytes(self.blobs[name][: size // 2])
|
||||
self.fetch()
|
||||
ranges = [h for h in _Handler.hits if h.endswith(name)]
|
||||
self.assertEqual(len(ranges), 1)
|
||||
self.assertTrue(drafting.pack_path(self.data).is_file())
|
||||
|
||||
def test_a_start_never_downloads(self):
|
||||
with acquisition.forbidden():
|
||||
with self.assertRaises(acquisition.AcquisitionRefused):
|
||||
drafting_fetch.ensure(self.data, base_url=self.base, files=self.files)
|
||||
self.assertFalse(drafting.pack_path(self.data).exists())
|
||||
|
||||
def test_the_plan_names_the_size_and_where_it_comes_from(self):
|
||||
with mock.patch.object(drafting, "pack_ready", return_value=False):
|
||||
plan = drafting_fetch.plan(self.data)
|
||||
self.assertEqual(plan["gb"], 1.2)
|
||||
self.assertIn("1.2 GB to download", plan["text"])
|
||||
self.assertIn("0.76 GB", plan["text"])
|
||||
self.assertIn("incoai/Qwen3.8-Flash-Next-DFlash2", plan["text"])
|
||||
self.assertIn("Apache-2.0", plan["text"])
|
||||
self.assertIn("starts with the MTP instead", plan["text"], "and what happens if it does not fit")
|
||||
with mock.patch.object(drafting, "pack_ready", return_value=True):
|
||||
self.assertEqual(drafting_fetch.plan(self.data)["gb"], 0.0)
|
||||
|
||||
def test_the_command_line_says_the_same(self):
|
||||
code, out, err = run_cli(["drafting", "plan", "--data-dir", str(self.data), "--output", "json"])
|
||||
self.assertEqual(code, 0, err)
|
||||
payload = json.loads(out)
|
||||
self.assertEqual((payload["ok"], payload["revision"]), (True, drafting.REVISION))
|
||||
self.assertEqual(payload["gb"], 1.2)
|
||||
|
||||
|
||||
class TheDownloadPlanListsIt(unittest.TestCase):
|
||||
"""`strata launch --plan --drafting dflash2` adds the drafter to the download when it would be used."""
|
||||
|
||||
def test_the_plan_line_and_the_part(self):
|
||||
from tests.terminal.launch_support import HW_12, Home, catalog
|
||||
from strata_cli.launch.runner import Runner
|
||||
with Home() as home:
|
||||
runner = Runner(S.LaunchSettings(family="qwen", model="Q2_0", context=32768, vision="off", drafting="dflash2"), catalog=catalog(HW_12))
|
||||
line = runner.drafting_plan_line()
|
||||
self.assertEqual(line, "Drafting: MTP (DFlash2 isn't available in this mode yet; this start uses the MTP)")
|
||||
self.assertNotIn("dflash2", [p["kind"] for p in runner.download_plan()["parts"]], "nothing to fetch while the engine cannot serve it")
|
||||
with serving_dflash2():
|
||||
kinds = {p["kind"]: p for p in runner.download_plan()["parts"]}
|
||||
self.assertEqual(kinds["dflash2"]["gb"], 1.2)
|
||||
self.assertIn("the DFlash2 drafter (1.2 GB)", runner.download_plan()["text"])
|
||||
self.assertIn("--drafting mtp", runner.download_plan()["text"])
|
||||
auto = Runner(S.LaunchSettings(family="qwen", model="Q2_0", context=32768, vision="off"), catalog=catalog(HW_12))
|
||||
self.assertNotIn("dflash2", [p["kind"] for p in auto.download_plan()["parts"]], "Auto is the MTP: no extra download")
|
||||
other = Runner(S.LaunchSettings(family="coder", model="IQ3_XXS", context=32768, vision="off", drafting="dflash2"), catalog=catalog(HW_12))
|
||||
self.assertNotIn("dflash2", [p["kind"] for p in other.download_plan()["parts"]], "not for this model: hidden, not offered")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+52
-15
@@ -17,21 +17,52 @@ import sys
|
||||
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
from safetensors import safe_open
|
||||
except ImportError: # pragma: no cover
|
||||
sys.exit("needs the safetensors package")
|
||||
|
||||
class SafeTensors:
|
||||
"""A .safetensors file read with numpy alone (no torch, no safetensors package: this runs on the person's PC, where neither is
|
||||
installed). The format is a little-endian u64 header length, a JSON header {name: {dtype, shape, data_offsets}}, then the raw data."""
|
||||
|
||||
def __init__(self, path):
|
||||
self.path = path
|
||||
with open(path, "rb") as fh:
|
||||
(n,) = struct.unpack("<Q", fh.read(8))
|
||||
if n > (64 << 20):
|
||||
raise ValueError("not a safetensors file (header too large)")
|
||||
self.header = json.loads(fh.read(n))
|
||||
self.base = 8 + n
|
||||
self.header.pop("__metadata__", None)
|
||||
|
||||
def raw(self, name):
|
||||
"""-> (dtype string, shape, uint8 memmap of the tensor's bytes)"""
|
||||
info = self.header[name]
|
||||
lo, hi = info["data_offsets"]
|
||||
return info["dtype"], info["shape"], np.memmap(self.path, dtype=np.uint8, mode="r", offset=self.base + lo, shape=(hi - lo,))
|
||||
|
||||
def get_float32(self, name):
|
||||
"""-> float32 ndarray (bf16 and f16 tensors are widened exactly)"""
|
||||
dtype, shape, data = self.raw(name)
|
||||
if dtype == "BF16":
|
||||
return (np.frombuffer(data, dtype=np.uint16).astype(np.uint32) << 16).view(np.float32).reshape(shape)
|
||||
if dtype == "F16":
|
||||
return np.frombuffer(data, dtype=np.float16).astype(np.float32).reshape(shape)
|
||||
if dtype == "F32":
|
||||
return np.frombuffer(data, dtype=np.float32).reshape(shape).copy()
|
||||
raise ValueError("%s: unsupported dtype %s" % (name, dtype))
|
||||
|
||||
def get_bf16_bits(self, name):
|
||||
dtype, shape, data = self.raw(name)
|
||||
if dtype != "BF16":
|
||||
raise ValueError("%s: expected BF16, found %s" % (name, dtype))
|
||||
return np.frombuffer(data, dtype=np.uint16).reshape(shape).copy()
|
||||
|
||||
|
||||
def load_tensor(f, name):
|
||||
"""-> float32 ndarray (bf16 tensors are widened exactly)"""
|
||||
return f.get_tensor(name).float().numpy()
|
||||
return f.get_float32(name)
|
||||
|
||||
|
||||
def raw_bf16(f, name):
|
||||
import torch
|
||||
t = f.get_tensor(name)
|
||||
return np.ascontiguousarray(t.view(torch.int16).numpy().view(np.uint16))
|
||||
return np.ascontiguousarray(f.get_bf16_bits(name))
|
||||
|
||||
|
||||
def quant_q8_0(w):
|
||||
@@ -49,13 +80,12 @@ def quant_q8_0(w):
|
||||
return out.tobytes()
|
||||
|
||||
|
||||
def main():
|
||||
if len(sys.argv) < 4:
|
||||
sys.exit(__doc__)
|
||||
st_path, cfg_path, out_path = sys.argv[1:4]
|
||||
cfg = json.load(open(cfg_path))
|
||||
def pack(st_path, cfg_path, out_path, say=print):
|
||||
"""Write the engine's runtime file for the drafter in `st_path` (+ its config.json). Returns the number of bytes written."""
|
||||
with open(cfg_path, encoding="utf-8") as fh:
|
||||
cfg = json.load(fh)
|
||||
dc = cfg["dflash_config"]
|
||||
f = safe_open(st_path, "pt")
|
||||
f = SafeTensors(st_path)
|
||||
tensors = []
|
||||
blobs = []
|
||||
|
||||
@@ -120,7 +150,14 @@ def main():
|
||||
for t, b in zip(tensors, blobs):
|
||||
o.seek(base + t["offset"])
|
||||
o.write(b)
|
||||
print("wrote %s: %d tensors, %.1f MB" % (out_path, len(tensors), (base + off) / 1e6))
|
||||
say("wrote %s: %d tensors, %.1f MB" % (out_path, len(tensors), (base + off) / 1e6))
|
||||
return base + off
|
||||
|
||||
|
||||
def main():
|
||||
if len(sys.argv) < 4:
|
||||
sys.exit(__doc__)
|
||||
pack(*sys.argv[1:4])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
#!/usr/bin/env python3
|
||||
"""tools/dflash2_pack.py reads a .safetensors file with numpy alone (the person's PC has neither torch nor the safetensors package) and writes
|
||||
the same runtime file the torch-based version wrote.
|
||||
|
||||
python tools/test_dflash2_pack.py
|
||||
"""
|
||||
import importlib.util
|
||||
import json
|
||||
import struct
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
sys.path.insert(0, str(HERE))
|
||||
import dflash2_pack as new # noqa: E402
|
||||
|
||||
L, H, INTER, HEADS, KV, HD, RANK, VOCAB = 2, 64, 96, 2, 1, 32, 32, 128
|
||||
CFG = {"num_hidden_layers": L, "hidden_size": H, "num_attention_heads": HEADS, "num_key_value_heads": KV, "head_dim": HD,
|
||||
"intermediate_size": INTER, "sliding_window": 128, "rope_parameters": {"rope_theta": 1000000.0}, "rms_norm_eps": 1e-6, "vocab_size": VOCAB,
|
||||
"dflash_config": {"block_size": 8, "conv_group_size": 16, "conv_kernel_size": 2, "mask_token_id": 7, "selector_rank": RANK, "selector_top_k": 4,
|
||||
"target_layer_ids": [1, 2, 3]}}
|
||||
|
||||
|
||||
def bf16(a):
|
||||
"""float32 -> BF16 bytes by truncation (what torch's .to(bfloat16) does for these values up to rounding; the test only needs a valid BF16 file)."""
|
||||
return (np.ascontiguousarray(a, np.float32).view(np.uint32) >> 16).astype(np.uint16)
|
||||
|
||||
|
||||
def tensors():
|
||||
rng = np.random.default_rng(7)
|
||||
r = lambda *shape: rng.standard_normal(shape).astype(np.float32) * 0.1
|
||||
t = {"fc.weight": r(H, H * 3), "hidden_norm.weight": r(H), "norm.weight": r(H)}
|
||||
for l in range(L):
|
||||
p = "layers.%d." % l
|
||||
t[p + "self_attn.q_proj.weight"] = r(HEADS * HD, H)
|
||||
t[p + "self_attn.k_proj.weight"] = r(KV * HD, H)
|
||||
t[p + "self_attn.v_proj.weight"] = r(KV * HD, H)
|
||||
t[p + "self_attn.o_proj.weight"] = r(H, HEADS * HD)
|
||||
t[p + "mlp.gate_proj.weight"] = r(INTER, H)
|
||||
t[p + "mlp.up_proj.weight"] = r(INTER, H)
|
||||
t[p + "mlp.down_proj.weight"] = r(H, INTER)
|
||||
t[p + "attention_conv.kernel_projection.weight"] = r(32, H)
|
||||
t[p + "mlp_conv.kernel_projection.weight"] = r(32, H)
|
||||
for n in ("input_layernorm", "post_attention_layernorm", "self_attn.q_norm", "self_attn.k_norm"):
|
||||
t[p + n + ".weight"] = r(HD if "_norm" in n else H)
|
||||
t[p + "attention_conv.base_kernel"] = r(2, 16)
|
||||
t[p + "mlp_conv.base_kernel"] = r(2, 16)
|
||||
t["candidate_selector.hidden_projection.weight"] = r(RANK, H)
|
||||
t["candidate_selector.predecessor_codebook"] = r(VOCAB, RANK)
|
||||
t["candidate_selector.successor_codebook"] = r(VOCAB, RANK)
|
||||
return t
|
||||
|
||||
|
||||
def write_safetensors(path, arrays):
|
||||
"""A BF16 safetensors file written by hand (the format is documented: u64 header length, JSON header, raw data)."""
|
||||
header, blobs, at = {}, [], 0
|
||||
for name, a in arrays.items():
|
||||
data = bf16(a).tobytes()
|
||||
header[name] = {"dtype": "BF16", "shape": list(a.shape), "data_offsets": [at, at + len(data)]}
|
||||
blobs.append(data)
|
||||
at += len(data)
|
||||
js = json.dumps(header).encode()
|
||||
js += b" " * ((8 - len(js) % 8) % 8)
|
||||
with open(path, "wb") as fh:
|
||||
fh.write(struct.pack("<Q", len(js)) + js + b"".join(blobs))
|
||||
|
||||
|
||||
class PackWithoutTorch(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.tmp = tempfile.TemporaryDirectory()
|
||||
cls.dir = Path(cls.tmp.name)
|
||||
cls.arrays = tensors()
|
||||
write_safetensors(cls.dir / "model.safetensors", cls.arrays)
|
||||
(cls.dir / "config.json").write_text(json.dumps(CFG))
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
cls.tmp.cleanup()
|
||||
|
||||
def pack(self, out):
|
||||
return new.pack(str(self.dir / "model.safetensors"), str(self.dir / "config.json"), str(out), say=lambda *_: None)
|
||||
|
||||
def test_reader_widens_bf16_exactly(self):
|
||||
f = new.SafeTensors(str(self.dir / "model.safetensors"))
|
||||
got = f.get_float32("fc.weight")
|
||||
want = (bf16(self.arrays["fc.weight"]).astype(np.uint32) << 16).view(np.float32)
|
||||
self.assertTrue(np.array_equal(got, want))
|
||||
self.assertEqual(f.get_bf16_bits("candidate_selector.predecessor_codebook").dtype, np.uint16)
|
||||
self.assertEqual(f.raw("norm.weight")[0], "BF16")
|
||||
|
||||
def test_the_file_has_the_header_the_engine_reads(self):
|
||||
out = self.dir / "out.pack"
|
||||
size = self.pack(out)
|
||||
data = out.read_bytes()
|
||||
self.assertEqual(data[:4], b"DF2P")
|
||||
version, index_len = struct.unpack("<II", data[4:12])
|
||||
self.assertEqual(version, 1)
|
||||
self.assertEqual((12 + index_len) % 4096, 0, "the data starts on a 4096 multiple")
|
||||
index = data[12:12 + index_len].decode().splitlines()
|
||||
self.assertTrue(any(l.startswith("cfg hidden %d" % H) for l in index))
|
||||
self.assertIn("end", [l.strip() for l in index])
|
||||
names = [l.split()[1] for l in index if l.startswith("t ")]
|
||||
self.assertIn("fc.weight", names)
|
||||
self.assertIn("candidate_selector.successor_codebook", names)
|
||||
self.assertEqual(len(data), size)
|
||||
|
||||
def test_a_matrix_is_ggml_q8_0_of_the_widened_values(self):
|
||||
out = self.dir / "q8.pack"
|
||||
self.pack(out)
|
||||
data = out.read_bytes()
|
||||
index_len = struct.unpack("<I", data[8:12])[0]
|
||||
base = 12 + index_len
|
||||
row = next(l.split() for l in data[12:base].decode().splitlines() if l.startswith("t fc.weight "))
|
||||
_t, _name, dtype, rows, cols, off, nbytes = row
|
||||
self.assertEqual(dtype, "q8_0")
|
||||
got = data[base + int(off): base + int(off) + int(nbytes)]
|
||||
widened = (bf16(self.arrays["fc.weight"]).astype(np.uint32) << 16).view(np.float32)
|
||||
self.assertEqual(got, new.quant_q8_0(widened))
|
||||
|
||||
def test_it_matches_the_torch_version_when_torch_is_here(self):
|
||||
try:
|
||||
import torch # noqa: F401
|
||||
import safetensors # noqa: F401
|
||||
except ImportError:
|
||||
self.skipTest("torch or safetensors is not installed here (the point of the change)")
|
||||
old_src = subprocess.run(["git", "show", "f1b6a2144:tools/dflash2_pack.py"], cwd=HERE, capture_output=True, text=True)
|
||||
if old_src.returncode != 0:
|
||||
self.skipTest("the earlier version is not in this clone")
|
||||
old_path = self.dir / "old_pack.py"
|
||||
old_path.write_text(old_src.stdout)
|
||||
spec = importlib.util.spec_from_file_location("old_pack", old_path)
|
||||
old = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(old)
|
||||
a, b = self.dir / "a.pack", self.dir / "b.pack"
|
||||
sys.argv = ["x", str(self.dir / "model.safetensors"), str(self.dir / "config.json"), str(a)]
|
||||
old.main()
|
||||
self.pack(b)
|
||||
self.assertEqual(a.read_bytes(), b.read_bytes(), "byte for byte the same file")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user