"""Correctness-first packed runtime for WAL-Ternary-8B V77. The runtime never constructs a complete BF16/FP16 Transformer body. It keeps the checkpoint in its physical T3+sparse-k8 / WALB2 / INT3 / INT4 formats and evaluates one output-row chunk at a time. The current implementation decodes logical codes on the CPU and performs the dot products with PyTorch on the requested CPU or CUDA device. It is intentionally a portable reference, not the final fused performance kernel. """ from __future__ import annotations import argparse import hashlib import importlib.util import json import math from pathlib import Path import sys import time from typing import Any, Iterator, Mapping import torch from torch import nn DEFAULT_CHECKPOINT = Path(__file__).resolve().parent def _sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: while block := handle.read(8 << 20): digest.update(block) return digest.hexdigest() def _load_module(name: str, path: Path): spec = importlib.util.spec_from_file_location(name, path) if spec is None or spec.loader is None: raise RuntimeError(f"cannot import {path}") module = importlib.util.module_from_spec(spec) sys.modules[name] = module spec.loader.exec_module(module) return module def bundled_runtime(checkpoint: str | Path) -> tuple[Any, Any]: root = Path(checkpoint).resolve(strict=True) runtime = root / "runtime" source = root / "src" for member in (runtime, source): if not member.is_dir(): raise ValueError(f"missing bundled runtime directory: {member}") for member in (str(runtime), str(source)): if member not in sys.path: sys.path.insert(0, member) endpoint_name = f"wal_v77_endpoint_{abs(hash(str(root)))}" endpoint = sys.modules.get(endpoint_name) if endpoint is None: endpoint = _load_module( endpoint_name, runtime / "qwen_endpoint_lowbit_artifact.py" ) import qwen_static_lowbit_artifact as static # type: ignore endpoint._wal_static = static from wal_tat import binary_lowrank as codec # type: ignore return endpoint, codec def load_manifests(checkpoint: str | Path) -> tuple[Path, dict, Path, dict]: root = Path(checkpoint).resolve(strict=True) manifest = json.loads((root / "manifest.json").read_text(encoding="utf-8")) base_manifest_path = root / str(manifest["base"]["manifest"]) base_root = base_manifest_path.parent base_manifest = json.loads(base_manifest_path.read_text(encoding="utf-8")) return root, manifest, base_root, base_manifest def _iter_t3_sparse_rows( endpoint: Any, path: Path, *, row_chunk: int, ) -> Iterator[tuple[int, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]: """Yield base codes, sparse codes, alpha and beta without dense weights.""" header = endpoint._read_matrix_header(path) if header["spec"].name != "t3-sparse-k8": raise ValueError(f"unsupported body format: {header['spec'].name}") rows = int(header["rows"]) columns = int(header["columns"]) group_size = int(header["group_size"]) groups_per_row = columns // group_size static = endpoint._wal_static code_bytes_per_row = groups_per_row * static.SPARSE_CODE_RECORD_BYTES metadata_bytes_per_row = groups_per_row * 2 * 2 metadata_offset = endpoint.MATRIX_HEADER.size + int(header["code_bytes"]) with path.open("rb") as handle: for start in range(0, rows, row_chunk): count = min(row_chunk, rows - start) handle.seek(endpoint.MATRIX_HEADER.size + start * code_bytes_per_row) code_raw = handle.read(count * code_bytes_per_row) handle.seek(metadata_offset + start * metadata_bytes_per_row) metadata_raw = handle.read(count * metadata_bytes_per_row) if len(code_raw) != count * code_bytes_per_row: raise ValueError("truncated T3+sparse-k8 codes") if len(metadata_raw) != count * metadata_bytes_per_row: raise ValueError("truncated T3+sparse-k8 metadata") packed = torch.frombuffer(bytearray(code_raw), dtype=torch.uint8) base, sparse = static._unpack_sparse_k8_records(packed) shape = (count, groups_per_row, group_size) base = base.reshape(shape) sparse = sparse.reshape(shape) metadata = torch.frombuffer( bytearray(metadata_raw), dtype=torch.float16 ).reshape(count, groups_per_row, 2) if not torch.isfinite(metadata).all() or torch.any(metadata <= 0): raise ValueError("invalid T3+sparse-k8 metadata") yield start, base, sparse, metadata[..., 0], metadata[..., 1] def _group_dot( value: torch.Tensor, codes: torch.Tensor, scales: torch.Tensor, ) -> torch.Tensor: """Compute grouped low-bit rows without constructing a dense weight.""" flat = value.reshape(-1, value.shape[-1]).float() rows, groups, group_size = codes.shape if flat.shape[-1] != groups * group_size: raise ValueError("activation width does not match packed matrix") grouped = flat.reshape(flat.shape[0], groups, group_size) logical = codes.to(device=value.device, dtype=torch.float32) scale = scales.to(device=value.device, dtype=torch.float32) dots = torch.einsum("bgi,rgi->brg", grouped, logical) return (dots * scale.unsqueeze(0)).sum(-1).reshape(*value.shape[:-1], rows) class PackedWALLinear(nn.Module): """T3+sparse-k8 base plus optional WALB2 correction.""" def __init__( self, base_path: Path, *, endpoint: Any, codec: Any, overlay_path: Path | None, in_features: int, out_features: int, row_chunk: int = 64, output_dtype: torch.dtype = torch.bfloat16, ) -> None: super().__init__() self.base_path = Path(base_path) self.overlay_path = None if overlay_path is None else Path(overlay_path) self.in_features = int(in_features) self.out_features = int(out_features) self.row_chunk = int(row_chunk) self.output_dtype = output_dtype self.endpoint = endpoint self.codec = codec self._overlay_cache: dict[str, tuple[Any, ...]] = {} def _overlay(self, device: torch.device) -> tuple[Any, ...]: if self.overlay_path is None: return () key = str(device) cached = self._overlay_cache.get(key) if cached is not None: return cached bundle = self.codec.read_binary_lowrank_bundle(self.overlay_path) if tuple(bundle.shape) != (self.out_features, self.in_features): raise ValueError("WALB2 shape does not match T3 base") paths = [] for item in bundle.paths: paths.append( ( item.u_codes.to(device=device, dtype=torch.int8), item.v_codes.to(device=device, dtype=torch.int8), item.row_scales_fp16.to(device=device), item.latent_scales_fp16.to(device=device), item.column_scales_fp16.to(device=device), ) ) cached = tuple(paths) self._overlay_cache[key] = cached return cached @torch.no_grad() def forward(self, value: torch.Tensor) -> torch.Tensor: chunks = [] for _, base, sparse, alpha, beta in _iter_t3_sparse_rows( self.endpoint, self.base_path, row_chunk=self.row_chunk ): part = _group_dot(value, base, alpha) part.add_(_group_dot(value, sparse, beta)) chunks.append(part) result = torch.cat(chunks, dim=-1) flat = value.reshape(-1, value.shape[-1]).float() correction = torch.zeros( (flat.shape[0], self.out_features), device=value.device, dtype=torch.float32, ) for u_codes, v_codes, row, latent, column in self._overlay(value.device): hidden = (flat * column.float().unsqueeze(0)) @ v_codes.float().transpose(0, 1) hidden.mul_(latent.float().unsqueeze(0)) left = u_codes.float() * row.float().unsqueeze(1) correction.add_(hidden @ left.transpose(0, 1)) result.add_(correction.reshape(*value.shape[:-1], self.out_features)) return result.to(dtype=self.output_dtype) class PackedEndpointLinear(nn.Module): """Streaming INT3/INT4 output projection without a dense weight.""" def __init__( self, path: Path, *, endpoint: Any, row_chunk: int, output_dtype: torch.dtype, ) -> None: super().__init__() self.path = Path(path) self.endpoint = endpoint header = endpoint._read_endpoint_header(self.path) payload = header["payload"] self.in_features = int(payload.columns) self.out_features = int(payload.rows) self.row_chunk = int(row_chunk) self.output_dtype = output_dtype @torch.no_grad() def forward(self, value: torch.Tensor) -> torch.Tensor: outputs = [] for _, codes, scales in self.endpoint.iter_quantized_endpoint_rows( self.path, row_chunk=self.row_chunk ): outputs.append(_group_dot(value, codes, scales)) return torch.cat(outputs, dim=-1).to(dtype=self.output_dtype) class PackedEmbedding(nn.Module): """Random-access INT3 embedding lookup from the packed endpoint file.""" def __init__(self, path: Path, *, endpoint: Any, output_dtype: torch.dtype) -> None: super().__init__() self.path = Path(path) self.endpoint = endpoint header = endpoint._read_endpoint_header(self.path) self.spec = header["spec"] self.payload = header["payload"] self.num_embeddings = int(self.payload.rows) self.embedding_dim = int(self.payload.columns) self.output_dtype = output_dtype def _rows(self, rows: list[int]) -> torch.Tensor: code_bytes = self.embedding_dim * int(self.spec.code_bits) // 8 groups = self.embedding_dim // int(self.payload.group_size) metadata_bytes = groups * int(self.spec.scale_bits) // 8 metadata_offset = ( self.endpoint.ENDPOINT_HEADER.size + int(self.payload.code_bits) // 8 ) decoded = [] with self.path.open("rb") as handle: for row in rows: handle.seek(self.endpoint.ENDPOINT_HEADER.size + row * code_bytes) raw_codes = handle.read(code_bytes) handle.seek(metadata_offset + row * metadata_bytes) raw_scales = handle.read(metadata_bytes) if len(raw_codes) != code_bytes or len(raw_scales) != metadata_bytes: raise ValueError("truncated packed embedding row") codes = self.endpoint._unpack_signed_codes( raw_codes, count=self.embedding_dim, spec=self.spec ).reshape(groups, int(self.payload.group_size)) scales = torch.frombuffer( bytearray(raw_scales), dtype=torch.float16 ).reshape(groups) decoded.append( (codes.float() * scales.float().unsqueeze(-1)).reshape(-1) ) return torch.stack(decoded) @torch.no_grad() def forward(self, input_ids: torch.Tensor) -> torch.Tensor: if input_ids.numel() == 0: return torch.empty( (*input_ids.shape, self.embedding_dim), device=input_ids.device, dtype=self.output_dtype, ) minimum = int(input_ids.min().item()) maximum = int(input_ids.max().item()) if minimum < 0 or maximum >= self.num_embeddings: raise IndexError("embedding index out of range") unique, inverse = torch.unique(input_ids.detach().cpu(), sorted=True, return_inverse=True) rows = self._rows([int(item) for item in unique.tolist()]) result = rows.index_select(0, inverse.reshape(-1)).reshape( *input_ids.shape, self.embedding_dim ) return result.to(device=input_ids.device, dtype=self.output_dtype) def _replace_module(model: nn.Module, name: str, replacement: nn.Module) -> None: parent_name, child_name = name.rsplit(".", 1) parent = model.get_submodule(parent_name) setattr(parent, child_name, replacement) def _load_raw_tensor(endpoint: Any, path: Path) -> torch.Tensor: rows = [value for _, value in endpoint.iter_raw_tensor_rows(path, row_chunk=256)] return torch.cat(rows, dim=0) @torch.no_grad() def load_packed_model( checkpoint: str | Path, *, device: str | torch.device, dtype: torch.dtype, row_chunk: int = 64, verify_hashes: bool = True, ) -> tuple[nn.Module, dict[str, Any]]: """Load Qwen3 with packed operators and no dense Transformer weights.""" started = time.monotonic() root, manifest, base_root, base_manifest = load_manifests(checkpoint) endpoint, codec = bundled_runtime(root) if verify_hashes: for row in manifest["checkpoint_files"]: member = root / str(row["file"]) if member.stat().st_size != int(row["bytes"]) or _sha256(member) != row["sha256"]: raise ValueError(f"checkpoint integrity failure: {row['file']}") from accelerate import init_empty_weights from transformers import AutoConfig, AutoModelForCausalLM requested = torch.device(device) config = AutoConfig.from_pretrained(root, local_files_only=True, trust_remote_code=False) with init_empty_weights(): model = AutoModelForCausalLM.from_config(config) base_by_name = {str(row["name"]): row for row in base_manifest["matrices"]} overlay_by_name = {str(row["name"]): row for row in manifest["overlays"]} original_modules = dict(model.named_modules()) for weight_name, row in base_by_name.items(): module_name = weight_name.removesuffix(".weight") original = original_modules[module_name] replacement = PackedWALLinear( base_root / str(row["file"]), endpoint=endpoint, codec=codec, overlay_path=( None if weight_name not in overlay_by_name else root / str(overlay_by_name[weight_name]["file"]) ), in_features=int(original.in_features), out_features=int(original.out_features), row_chunk=row_chunk, output_dtype=dtype, ) _replace_module(model, module_name, replacement) endpoints = {str(row["name"]): row for row in base_manifest["endpoints"]} embedding_row = endpoints["model.embed_tokens.weight"] head_row = endpoints["lm_head.weight"] model.model.embed_tokens = PackedEmbedding( base_root / str(embedding_row["file"]), endpoint=endpoint, output_dtype=dtype ) model.lm_head = PackedEndpointLinear( base_root / str(head_row["file"]), endpoint=endpoint, row_chunk=row_chunk, output_dtype=dtype, ) buffer_snapshots = { name: value.detach().contiguous().cpu().clone() for name, value in model.named_buffers() if not value.is_meta } model.to_empty(device=requested) model.to(dtype=dtype) buffers = dict(model.named_buffers()) for name, snapshot in buffer_snapshots.items(): buffers[name].copy_(snapshot.to(device=requested, dtype=buffers[name].dtype)) if hasattr(model.model.rotary_emb, "inv_freq"): model.model.rotary_emb.original_inv_freq = model.model.rotary_emb.inv_freq parameters = dict(model.named_parameters()) for row in base_manifest["unquantized_tensors"]: name = str(row["name"]) if name not in parameters: raise ValueError(f"raw parameter absent after packed replacement: {name}") value = _load_raw_tensor(endpoint, base_root / str(row["file"])) parameters[name].copy_(value.to(device=requested, dtype=dtype)) model.eval() dense_body_parameters = sum( parameter.numel() for name, parameter in model.named_parameters() if ".layers." in name and name.endswith(".weight") and parameter.ndim == 2 ) report = { "status": "ok", "checkpoint": str(root), "device": str(requested), "dtype": str(dtype), "body_packed_matrices": len(base_by_name), "walb2_matrices": len(overlay_by_name), "dense_body_matrix_parameters": int(dense_body_parameters), "full_bf16_body_materialized": False, "temporary_decode_scope": f"at most {row_chunk} output rows", "load_seconds": time.monotonic() - started, } return model, report def operator_check( checkpoint: str | Path, *, device: str, row_chunk: int, matrix_name: str | None, seed: int, ) -> dict[str, Any]: root, manifest, base_root, base_manifest = load_manifests(checkpoint) endpoint, codec = bundled_runtime(root) base_by_name = {str(row["name"]): row for row in base_manifest["matrices"]} overlay_by_name = {str(row["name"]): row for row in manifest["overlays"]} if matrix_name is None: matrix_name = next(name for name in base_by_name if name in overlay_by_name) row = base_by_name[matrix_name] shape = tuple(int(item) for item in row["shape"]) module = PackedWALLinear( base_root / str(row["file"]), endpoint=endpoint, codec=codec, overlay_path=( None if matrix_name not in overlay_by_name else root / str(overlay_by_name[matrix_name]["file"]) ), in_features=shape[1], out_features=shape[0], row_chunk=row_chunk, output_dtype=torch.float32, ) generator = torch.Generator(device="cpu").manual_seed(seed) value = torch.randn((2, shape[1]), generator=generator, dtype=torch.float32).to(device) started = time.monotonic() actual = module(value).float() if str(device).startswith("cuda"): torch.cuda.synchronize(torch.device(device)) packed_seconds = time.monotonic() - started dense_parts = [] for _, rows in endpoint.iter_dequantized_matrix_rows( base_root / str(row["file"]), row_chunk=row_chunk, dtype=torch.float32 ): dense_parts.append(rows) dense = torch.cat(dense_parts) if matrix_name in overlay_by_name: dense.add_( codec.decode_binary_lowrank_bundle( codec.read_binary_lowrank_bundle( root / str(overlay_by_name[matrix_name]["file"]) ), dtype=torch.float32, ) ) expected = value.float() @ dense.to(device).transpose(0, 1) error = (actual - expected).abs() denominator = expected.abs().clamp_min(1e-6) return { "status": "ok", "matrix": matrix_name, "shape": list(shape), "device": device, "row_chunk": row_chunk, "packed_seconds": packed_seconds, "max_abs_error": float(error.max().item()), "mean_abs_error": float(error.mean().item()), "mean_relative_error": float((error / denominator).mean().item()), "dense_body_materialized_by_runtime": False, "dense_matrix_materialized_only_by_test_oracle": True, } def smoke( checkpoint: str | Path, *, device: str, row_chunk: int, prompt: str, ) -> dict[str, Any]: from transformers import AutoTokenizer dtype = torch.bfloat16 if str(device).startswith("cuda") else torch.float32 model, report = load_packed_model( checkpoint, device=device, dtype=dtype, row_chunk=row_chunk, verify_hashes=True ) tokenizer = AutoTokenizer.from_pretrained( checkpoint, local_files_only=True, trust_remote_code=False ) encoded = tokenizer(prompt, return_tensors="pt") encoded = {key: value.to(device) for key, value in encoded.items()} if str(device).startswith("cuda"): torch.cuda.reset_peak_memory_stats(torch.device(device)) started = time.monotonic() logits = model(**encoded, use_cache=False).logits[:, -1, :].float() if str(device).startswith("cuda"): torch.cuda.synchronize(torch.device(device)) report["forward"] = { "seconds": time.monotonic() - started, "input_tokens": int(encoded["input_ids"].numel()), "logits_finite": bool(torch.isfinite(logits).all()), "argmax_token_id": int(logits.argmax(dim=-1).item()), "argmax_token": tokenizer.decode(logits.argmax(dim=-1).tolist()), "cuda_peak_allocated_bytes": ( int(torch.cuda.max_memory_allocated(torch.device(device))) if str(device).startswith("cuda") else None ), } return report def main() -> None: parser = argparse.ArgumentParser(description=__doc__) sub = parser.add_subparsers(dest="command", required=True) check = sub.add_parser("operator-check") check.add_argument("checkpoint", nargs="?", default=str(DEFAULT_CHECKPOINT)) check.add_argument("--device", default="cpu") check.add_argument("--row-chunk", type=int, default=64) check.add_argument("--matrix") check.add_argument("--seed", type=int, default=20260806) run = sub.add_parser("smoke") run.add_argument("checkpoint", nargs="?", default=str(DEFAULT_CHECKPOINT)) run.add_argument("--device", default="cpu") run.add_argument("--row-chunk", type=int, default=64) run.add_argument("--prompt", default="Hello") args = parser.parse_args() if args.command == "operator-check": result = operator_check( args.checkpoint, device=args.device, row_chunk=args.row_chunk, matrix_name=args.matrix, seed=args.seed, ) else: result = smoke( args.checkpoint, device=args.device, row_chunk=args.row_chunk, prompt=args.prompt, ) print(json.dumps(result, ensure_ascii=False, indent=2, allow_nan=False)) if __name__ == "__main__": main()