qev-s-0.8b: qev-0.8b + router + games and e2e-tester specialists (2-bit), bundle qevS2d
2c045dc verified | """Serving engines for an A1 model on one GPU (OPTIMIZATION.md, fast_app.infer_bench). | |
| - `load_a1(base, trunk)`: the base Qwen3.5 checkpoint in bf16 with an A1 trunk (DiLoCo global: {"{layer}.{param}"} and | |
| "24.weight" for the final norm) copied in, exactly as Society + jev_app._load_globals load it for evaluation. | |
| - `GraphRunner`: batch 1 on CUDA graphs. One graph per length bucket (the request right-padded to it: causal layers, | |
| so padding never reaches an earlier token), the forward torch.compile'd before capture (the fla kernels stay | |
| graph breaks), up to `max_slots` answer slots read inside the graph. Longer requests run the packed forward eagerly. | |
| - `Packed`: batches of requests as one unpadded varlen forward (fla `cu_seqlens`, varlen FlashAttention-2). | |
| - `quantize(model, mode)`: the fused projections as int8 weight-only (torchao), FP8 W8A8 (e4m3, per-channel weight and | |
| per-token activation scales, torch._scaled_mm) or int8 W8A8 (torchao dynamic int8); embeddings / letter head stay bf16. | |
| """ | |
| import json | |
| import time | |
| import torch | |
| from torch import nn | |
| from . import systemone as S1 | |
| from .model import Batch, FastQwen35 | |
| # The pre-tokenizer the training data was encoded with (the Qwen2 split pattern). transformers' Qwen tokenizer | |
| # classes impose their own pattern whatever tokenizer.json says (5.17: this one; 5.20+: the Qwen3.5 pattern with | |
| # combining marks), so it is set explicitly: otherwise words in scripts with combining marks (Indic and others) | |
| # split differently from training. | |
| TRAINING_PRE_TOKENIZER = json.loads('{"type": "Sequence", "pretokenizers": [{"type": "Split", "pattern": {"Regex": "(?i:\'s|\'t|\'re|\'ve|\'m|\'ll|\'d)|[^\\\\r\\\\n\\\\p{L}\\\\p{N}]?\\\\p{L}+|\\\\p{N}| ?[^\\\\s\\\\p{L}\\\\p{N}]+[\\\\r\\\\n]*|\\\\s*[\\\\r\\\\n]+|\\\\s+(?!\\\\S)|\\\\s+"}, "behavior": "Isolated", "invert": false}, {"type": "ByteLevel", "add_prefix_space": false, "trim_offsets": true, "use_regex": false}]}') | |
| def use_training_pretokenizer(tok): | |
| """Set the training-time pre-tokenizer on a transformers fast tokenizer (in place) -> tok.""" | |
| from tokenizers import Tokenizer | |
| t = json.loads(tok.backend_tokenizer.to_str()) | |
| t["pre_tokenizer"] = TRAINING_PRE_TOKENIZER | |
| tok.backend_tokenizer.pre_tokenizer = Tokenizer.from_str(json.dumps(t)).pre_tokenizer | |
| return tok | |
| BUCKETS = (128, 192, 256, 320, 384, 448, 512, 640, 768, 896, 1024, 1280, 1536, 2048, 2560, 3072, 4096) | |
| def load_a1(base: str, trunk: str | None, device="cuda", dtype=torch.bfloat16): | |
| """-> (lm, tok, letters): transformers' causal LM of `base` with the trunk at `trunk` (a global.pt) loaded.""" | |
| from decider.prompt import letter_ids | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| tok = use_training_pretokenizer(AutoTokenizer.from_pretrained(base)) | |
| lm = AutoModelForCausalLM.from_pretrained(base, dtype=dtype).to(device).eval() | |
| if trunk: | |
| glob = torch.load(trunk, map_location="cpu", weights_only=True) | |
| mods = list(lm.model.layers) + [lm.model.norm] | |
| n = 0 | |
| with torch.no_grad(): | |
| for i, mod in enumerate(mods): | |
| for name, p in mod.named_parameters(): | |
| key = f"{i}.{name}" | |
| if key in glob: | |
| p.copy_(glob[key].to(p.device, p.dtype)) | |
| n += 1 | |
| assert n == len(glob), (n, len(glob)) | |
| for p in lm.parameters(): | |
| p.requires_grad_(False) | |
| return lm, tok, letter_ids(tok) | |
| # ---- quantised linears ----------------------------------------------------------------------------------------------- | |
| E4M3_MAX = 448.0 | |
| def _rowwise_fp8(x): | |
| s = x.abs().amax(dim=-1, keepdim=True).float().clamp(min=1e-12) / E4M3_MAX | |
| return (x.float() / s).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn), s | |
| class FP8Linear(nn.Module): | |
| """y = x W^T with W in e4m3 (per output channel scale) and x quantised per token at run time (W8A8 FP8).""" | |
| def __init__(self, lin: nn.Linear): | |
| super().__init__() | |
| wq, sw = _rowwise_fp8(lin.weight.detach()) | |
| self.register_buffer("wq", wq.contiguous()) | |
| self.register_buffer("sw_t", sw.t().contiguous()) | |
| self.in_features, self.out_features, self.out_dtype = lin.in_features, lin.out_features, lin.weight.dtype | |
| def forward(self, x): | |
| shp = x.shape[:-1] | |
| xq, sx = _rowwise_fp8(x.reshape(-1, self.in_features)) | |
| y = torch._scaled_mm(xq, self.wq.t(), scale_a=sx, scale_b=self.sw_t, out_dtype=self.out_dtype) | |
| return y.reshape(*shp, self.out_features) | |
| def quantize(model: FastQwen35, mode: str) -> dict: | |
| """mode: "bf16" (nothing), "int8wo", "int8dq", "fp8", "int4wo". -> {"mode", "linears", "weight_gb"}.""" | |
| assert model.fuse, "quantisation acts on the fused inference copies" | |
| lins = [(mods, name) for mods in model.fused for name in list(mods.keys())] | |
| if mode == "fp8": | |
| for mods, name in lins: | |
| mods[name] = FP8Linear(mods[name]) | |
| elif mode in ("int8wo", "int8dq", "int4wo"): | |
| from torchao.quantization import quantize_ | |
| if mode == "int8wo": | |
| from torchao.quantization import Int8WeightOnlyConfig as C | |
| cfg = C() | |
| elif mode == "int8dq": | |
| from torchao.quantization import Int8DynamicActivationInt8WeightConfig as C | |
| cfg = C() | |
| else: | |
| from torchao.quantization import Int4WeightOnlyConfig as C | |
| cfg = C(group_size=128) | |
| quantize_(model.fused, cfg) | |
| elif mode != "bf16": | |
| raise ValueError(mode) | |
| size = 0 | |
| for mods in model.fused: | |
| for m in mods.values(): | |
| for t in list(m.parameters()) + list(m.buffers()): | |
| try: | |
| size += t.numel() * t.element_size() | |
| except Exception: # noqa: BLE001 -- tensor subclasses without a plain element size | |
| size += t.nbytes if hasattr(t, "nbytes") else 0 | |
| return {"mode": mode, "linears": len(lins), "weight_gb": round(size / 1e9, 3)} | |
| # ---- batch 1 on CUDA graphs ------------------------------------------------------------------------------------------ | |
| class GraphRunner: | |
| def __init__(self, model: FastQwen35, pad_id: int, buckets=BUCKETS, max_slots: int = 32, compile: bool = True): | |
| self.model, self.pad_id, self.buckets, self.max_slots = model, pad_id, tuple(buckets), max_slots | |
| self.compile = compile | |
| self.fn = torch.compile(model.forward, dynamic=False) if compile else model.forward | |
| self.graphs: dict = {} | |
| self.pool = torch.cuda.graph_pool_handle() | |
| self.stats = {"graph": 0, "eager": 0, "captures": 0, "capture_s": 0.0} | |
| def _bucket(self, n): | |
| for b in self.buckets: | |
| if n <= b: | |
| return b | |
| return None | |
| def _capture(self, T): | |
| t0 = time.perf_counter() | |
| b = Batch([[self.pad_id] * T], [[0] * self.max_slots], "cuda", pad_to=T, pad_id=self.pad_id) | |
| st = torch.cuda.Stream() | |
| st.wait_stream(torch.cuda.current_stream()) | |
| with torch.cuda.stream(st), torch.no_grad(): | |
| for _ in range(3): | |
| self.fn(b) | |
| torch.cuda.current_stream().wait_stream(st) | |
| g = torch.cuda.CUDAGraph() | |
| with torch.cuda.graph(g, pool=self.pool), torch.no_grad(): | |
| out = self.fn(b) | |
| self.stats["captures"] += 1 | |
| self.stats["capture_s"] += time.perf_counter() - t0 | |
| return b, out, g | |
| def warm(self, lengths): | |
| for T in sorted({self._bucket(n) for n in lengths if self._bucket(n)}): | |
| if T not in self.graphs: | |
| self.graphs[T] = self._capture(T) | |
| def logits(self, ids: list, slots: list): | |
| """One request's prompt row -> letter logits at its slots (CPU float tensor [len(slots), 255]).""" | |
| T = self._bucket(len(ids)) | |
| if T is None or len(slots) > self.max_slots: | |
| self.stats["eager"] += 1 | |
| return self.model(Batch([ids], [slots], "cuda")).cpu() | |
| hit = self.graphs.get(T) | |
| if hit is None: | |
| hit = self.graphs[T] = self._capture(T) | |
| b, out, g = hit | |
| b.ids.fill_(self.pad_id) | |
| b.ids[: len(ids)].copy_(torch.as_tensor(ids, dtype=torch.long), non_blocking=True) | |
| s = torch.zeros(self.max_slots, dtype=torch.long) | |
| s[: len(slots)] = torch.as_tensor(slots) | |
| b.slots.copy_(s, non_blocking=True) | |
| g.replay() | |
| self.stats["graph"] += 1 | |
| return out[: len(slots)].cpu() | |
| class Packed: | |
| """Requests in batches as one unpadded varlen forward (the throughput path).""" | |
| def __init__(self, model: FastQwen35, compile_pieces: bool = False): | |
| self.model = model | |
| if compile_pieces: | |
| compile_elementwise() | |
| def logits(self, rows: list) -> list: | |
| """rows: [(ids, slots)] -> per row its letter logits (CPU).""" | |
| b = Batch([r[0] for r in rows], [r[1] for r in rows], "cuda") | |
| out = self.model(b).cpu() | |
| res, k = [], 0 | |
| for _, sl in rows: | |
| res.append(out[k: k + len(sl)]) | |
| k += len(sl) | |
| return res | |
| _COMPILED = {"done": False} | |
| def compile_elementwise(): | |
| """torch.compile (dynamic shapes) the norm / rope / gate pieces of qev_fast.model in place: one fused kernel | |
| each instead of chains of fp32 and bf16 ops, also for varlen batches whose length changes every call.""" | |
| from . import model as M | |
| if _COMPILED["done"]: | |
| return | |
| for name in ("rmsnorm", "gated_rmsnorm", "apply_rope", "gdn_gates"): | |
| setattr(M, name, torch.compile(getattr(M, name), dynamic=True)) | |
| _COMPILED["done"] = True | |
| # ---- System One on top ----------------------------------------------------------------------------------------------- | |
| class A1Server: | |
| """state + questions -> wire-format answers (decider's /v1/systemone shape) with an A1 model on a runner.""" | |
| def __init__(self, tok, runner, temps: dict): | |
| self.tok, self.runner, self.temps = tok, runner, temps | |
| def serve(self, state, questions) -> dict: | |
| enc = S1.encode(self.tok, state, questions) | |
| lg = self.runner.logits(enc["ids"], enc["slots"]) | |
| return S1.answers(enc, S1.probabilities(enc, lg.tolist(), self.temps)) | |