#!/usr/bin/env python3 """Offline, bounded KV memory capacity estimator for standard full-attention MHA/GQA. This is a memory budget calculation, not an engine deployment or measured GPU test. It deliberately rejects MLA, sliding-window, hybrid models, and multi-GPU sharding. No network, subprocess, CUDA, PyTorch, model loading, or credential access. """ from __future__ import annotations import argparse from dataclasses import dataclass import hashlib import json import math from pathlib import Path GIB = 1024 ** 3 KIB = 1024 MAX_INPUT = 256 * 1024 DTYPE_BYTES = {"fp16": 2, "bf16": 2, "fp8": 1} class UnsupportedArchitecture(ValueError): pass def _integer(cfg: dict, key: str) -> int: val = cfg.get(key) if isinstance(val, bool) or not isinstance(val, int) or not 0 < val <= 10_000_000: raise ValueError("Invalid positive integer field: " + key) return val def dimensions(cfg: dict) -> dict: if not isinstance(cfg, dict): raise ValueError("Model configuration must be a JSON object") if cfg.get("kv_lora_rank"): raise UnsupportedArchitecture("MLA architecture not supported") if cfg.get("sliding_window") not in (None, False, 0): raise UnsupportedArchitecture("Sliding-window layers not supported") if cfg.get("sliding_window_pattern") not in (None, False, 0): raise UnsupportedArchitecture("Layer-pattern sliding attention not supported") if cfg.get("layer_types") not in (None, []): # Even 'full_attention' can hide hybrid details in model-family-specific metadata. raise UnsupportedArchitecture("Per-layer attention configuration requires manual analysis") if cfg.get("attention_type") not in (None, "full_attention"): raise UnsupportedArchitecture("Nonstandard attention configuration") if cfg.get("is_encoder_decoder") is True: raise UnsupportedArchitecture("Encoder-decoder cross-attention requires additional cache accounting") layers = _integer(cfg, "num_hidden_layers") qheads = _integer(cfg, "num_attention_heads") kvheads = cfg.get("num_key_value_heads", qheads) if isinstance(kvheads, bool) or not isinstance(kvheads, int) or not 0 < kvheads <= qheads: raise ValueError("Invalid num_key_value_heads") if qheads % kvheads != 0: raise ValueError("Query heads must be divisible by KV heads") head_dim = cfg.get("head_dim") if head_dim is None: hidden = _integer(cfg, "hidden_size") if hidden % qheads: raise ValueError("hidden_size must be divisible by query heads") head_dim = hidden // qheads if isinstance(head_dim, bool) or not isinstance(head_dim, int) or not 0 < head_dim <= 1_000_000: raise ValueError("Invalid head_dim") return {"layers": layers, "query_heads": qheads, "key_value_heads": kvheads, "head_dim": head_dim, "attention": "MHA" if qheads == kvheads else "GQA"} def _finite(name: str, x: float, *, lower: float=0.0, upper: float=1_000_000.0) -> float: if not isinstance(x,(int,float)) or not math.isfinite(x) or x < lower or x > upper: raise ValueError("Invalid numeric input: "+name) return float(x) def plan(cfg: dict, *, gpu_vram_gib: float, weights_gib: float, runtime_reserve_gib: float, gpu_utilization: float, context_tokens: int, requested_sequences: int, block_size: int, kv_dtype: str) -> dict: dims = dimensions(cfg) capacity = _finite("gpu_vram_gib",gpu_vram_gib,lower=0.001) weights = _finite("weights_gib",weights_gib) reserve = _finite("runtime_reserve_gib",runtime_reserve_gib) util = _finite("gpu_utilization",gpu_utilization,lower=0.001,upper=1.0) if kv_dtype not in DTYPE_BYTES: raise ValueError("KV dtype must be fp16, bf16 or fp8") if isinstance(context_tokens,bool) or not isinstance(context_tokens,int) or not 1<=context_tokens<=10_000_000: raise ValueError("Invalid context token count") if isinstance(requested_sequences,bool) or not isinstance(requested_sequences,int) or not 1<=requested_sequences<=10_000_000: raise ValueError("Invalid concurrent sequence count") if block_size not in (8,16,32,64,128): raise ValueError("Unsupported block size; use 8,16,32,64,128") bytes_per_token = ( 2 * dims["layers"] * dims["key_value_heads"] * dims["head_dim"] * DTYPE_BYTES[kv_dtype] ) padded_tokens = ((context_tokens + block_size - 1) // block_size) * block_size bytes_per_sequence = bytes_per_token * padded_tokens usable_bytes = math.floor(capacity * GIB * util) fixed_bytes = math.ceil(weights * GIB) + math.ceil(reserve * GIB) available_bytes = max(0, usable_bytes - fixed_bytes) max_sequences = available_bytes // bytes_per_sequence max_physical_kv_tokens = available_bytes // (bytes_per_token*block_size)*block_size requested_kv_bytes = bytes_per_sequence * requested_sequences fits = max_sequences >= requested_sequences return { "status":"ESTIMATED_WITHIN_DECLARED_MEMORY_BUDGET" if fits else "EXCEEDS_DECLARED_MEMORY_BUDGET", "assumptions":{ "attention_architecture":dims["attention"], "single_GPU_only":True, "full_attention_layers_only":True, "kv_dtype":kv_dtype, "gpu_vram_gib":capacity,"gpu_utilization":util, "weights_resident_gib_user_supplied":weights, "runtime_reserve_gib_user_supplied":reserve, "context_tokens_per_sequence":context_tokens, "block_size_tokens":block_size, "requested_concurrent_sequences":requested_sequences }, "architecture":dims, "byte_calculation":{ "bytes_per_kv_token":bytes_per_token, "padded_tokens_per_sequence":padded_tokens, "kv_bytes_per_sequence":bytes_per_sequence, "usable_gpu_bytes":usable_bytes, "fixed_reserved_bytes":fixed_bytes, "remaining_cache_budget_bytes":available_bytes, "requested_cache_bytes":requested_kv_bytes, "maximum_concurrent_sequences_under_budget":max_sequences, "maximum_physical_kv_tokens_under_budget":max_physical_kv_tokens }, "fits_declared_budget":fits, "warnings":[ "Not a runtime measurement, deployment acceptance test, or index/SEO metric.", "User supplies weight and runtime reservations; inaccurate values change the estimate.", "Actual serving engines have memory fragmentation, CUDA graphs, KV block scheduling and shared-prefix effects.", "FP8 KV cache requires compatible hardware, backend and numerical validation.", "Does not support MLA, hybrid/sliding-window attention, cross-attention, tensor parallel, or CPU offloading.", "A numerical fit does not guarantee model load, throughput, correctness or absence of OOM." ] } def config_file(path: Path) -> tuple[dict,str]: path=path.expanduser() if not path.is_file() or path.stat().st_size > MAX_INPUT: raise ValueError("Config is absent or larger than 256 KiB") data=path.read_bytes() obj=json.loads(data.decode("utf-8")) return obj,hashlib.sha256(data).hexdigest() def self_test() -> dict: # Synthetic GQA: 64 layers x 8 KV heads x 128 dims x 2 key/value x BF16=262,144B/token. cfg={"num_hidden_layers":64,"num_attention_heads":40, "num_key_value_heads":8,"head_dim":128} opts={"gpu_vram_gib":24,"weights_gib":17.1,"runtime_reserve_gib":1.4, "gpu_utilization":0.9,"context_tokens":8192, "requested_sequences":1,"block_size":16,"kv_dtype":"bf16"} a=plan(cfg,**opts) assert a["byte_calculation"]["bytes_per_kv_token"]==262144 assert a["byte_calculation"]["kv_bytes_per_sequence"]==2*GIB assert a["byte_calculation"]["maximum_concurrent_sequences_under_budget"]==1 assert a["fits_declared_budget"] assert not plan(cfg,**(opts|{"requested_sequences":2}))["fits_declared_budget"] b=plan(cfg,**(opts|{"kv_dtype":"fp8","requested_sequences":2})) assert b["byte_calculation"]["bytes_per_kv_token"]==131072 assert b["fits_declared_budget"] c=plan(cfg,**(opts|{"context_tokens":8193})) assert c["byte_calculation"]["padded_tokens_per_sequence"]==8208 # Confirm invalid configs are rejected before allocating any GPU resources. tests=[ ({"num_hidden_layers":64,"num_attention_heads":40,"num_key_value_heads":8, "head_dim":128,"kv_lora_rank":512},UnsupportedArchitecture), ({"num_hidden_layers":64,"num_attention_heads":40,"num_key_value_heads":8, "head_dim":128,"sliding_window":4096},UnsupportedArchitecture), ({"num_hidden_layers":64,"num_attention_heads":40,"num_key_value_heads":7, "head_dim":128},ValueError), ] for x,error in tests: try: plan(x,**opts) raise AssertionError("Rejected architecture accepted") except error: pass try: plan(cfg,**(opts|{"gpu_utilization":float("nan")})) raise AssertionError("NaN utilization accepted") except ValueError: pass try: plan(cfg,**(opts|{"requested_sequences":0})) raise AssertionError("Zero sequences accepted") except ValueError: pass return {"status":"PASS","synthetic_positive_cases":4, "rejected_invalid_cases":5, "gpu_executions":0,"network_calls":0,"model_downloads":0} def main() -> int: p=argparse.ArgumentParser(description=__doc__) p.add_argument("--config",type=Path,help="Local Hugging Face config.json (no URL, no network)") p.add_argument("--gpu-vram-gib",type=float,help="Physical GPU capacity in GiB") p.add_argument("--weights-gib",type=float,help="Measured/estimated resident model weights in GiB") p.add_argument("--runtime-reserve-gib",type=float,help="Reserved engine/runtime+activation VRAM in GiB") p.add_argument("--gpu-utilization",type=float,default=0.90) p.add_argument("--context-tokens",type=int,default=8192) p.add_argument("--sequences",type=int,default=1) p.add_argument("--block-size",type=int,default=16) p.add_argument("--kv-dtype",choices=tuple(DTYPE_BYTES),default="bf16") p.add_argument("--self-test",action="store_true") args=p.parse_args() if args.self_test: print(json.dumps(self_test(),indent=2)) return 0 for value in ("config","gpu_vram_gib","weights_gib","runtime_reserve_gib"): if getattr(args,value) is None: p.error("required unless --self-test: --"+value.replace("_","-")) try: cfg,digest=config_file(args.config) result=plan(cfg, gpu_vram_gib=args.gpu_vram_gib,weights_gib=args.weights_gib, runtime_reserve_gib=args.runtime_reserve_gib, gpu_utilization=args.gpu_utilization,context_tokens=args.context_tokens, requested_sequences=args.sequences,block_size=args.block_size,kv_dtype=args.kv_dtype) result["input_config_sha256"]=digest print(json.dumps(result,indent=2)) return 0 except (ValueError,UnsupportedArchitecture,json.JSONDecodeError,UnicodeDecodeError) as exc: print(json.dumps({"status":"INVALID_OR_UNSUPPORTED_INPUT", "reason":str(exc)[:180]},indent=2)) return 2 if __name__=="__main__": raise SystemExit(main())