"""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) @torch.no_grad() 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() @torch.no_grad() 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))