Spaces:
Running on Zero
Running on Zero
| """Functional hybrid prefixes with native multimodal embeddings and 3-axis RoPE. | |
| Prefix and the layer-wise training schedule are adapted from | |
| jaredpalmer/kev@f1535963cea021439370c23127bc970b6788e730/kev/shared_prefix.py | |
| (Apache-2.0; see third_party/kev/). Branches never mutate a saved prefix. | |
| """ | |
| import os | |
| from collections import OrderedDict | |
| import torch | |
| from torch.nn import functional as F | |
| from torch.utils.checkpoint import checkpoint | |
| class Prefix: | |
| record_past = True | |
| def __init__(self): | |
| self.keys = self.values = self.conv = None | |
| self.recurrent_states = {} | |
| def layers(self): | |
| return self | |
| def __getitem__(self, layer_idx): | |
| return self | |
| def has_previous_state(self, layer_idx=None, state_idx=0): | |
| return self.conv is not None | |
| def update_conv_state(self, conv_input, layer_idx, conv_kernel_size, **kwargs): | |
| if self.conv is not None: | |
| return torch.cat([self.conv, conv_input], -1) | |
| self.conv = F.pad(conv_input, (max(0, conv_kernel_size - conv_input.shape[-1]), 0))[ | |
| ..., -conv_kernel_size: | |
| ] | |
| return conv_input | |
| def update_recurrent_state(self, recurrent_state, layer_idx, **kwargs): | |
| self.recurrent_states.setdefault(0, recurrent_state) | |
| return recurrent_state | |
| def update(self, keys, values, layer_idx, *args, **kwargs): | |
| if self.keys is None: | |
| self.keys, self.values = keys, values | |
| return keys, values | |
| return torch.cat([self.keys, keys], -2), torch.cat([self.values, values], -2) | |
| def branch(self, owner): | |
| other = Prefix() | |
| pick = lambda t: None if t is None else t.index_select(0, owner) | |
| other.keys, other.values, other.conv = pick(self.keys), pick(self.values), pick(self.conv) | |
| other.recurrent_states = {k: pick(v) for k, v in self.recurrent_states.items()} | |
| return other | |
| def nbytes(self): | |
| tensors = [self.keys, self.values, self.conv, *self.recurrent_states.values()] | |
| return sum(t.numel() * t.element_size() for t in tensors if t is not None) | |
| def compact(self): | |
| # Cached slices must own their storage so byte accounting does not hide a | |
| # complete prefill tensor behind a small convolution window or KV view. | |
| self.keys = None if self.keys is None else self.keys.clone().contiguous() | |
| self.values = None if self.values is None else self.values.clone().contiguous() | |
| self.conv = None if self.conv is None else self.conv.clone().contiguous() | |
| self.recurrent_states = { | |
| k: v.clone().contiguous() for k, v in self.recurrent_states.items() | |
| } | |
| return self | |
| def attention_mask(allow, hidden, implementation): | |
| allow = allow[:, None] | |
| if implementation == "sdpa": | |
| return allow | |
| return hidden.new_zeros(allow.shape).masked_fill(~allow, torch.finfo(hidden.dtype).min) | |
| def split_rows(hidden, positions, packed): | |
| prefixes = packed["prefix_lengths"].tolist() | |
| lengths = packed["lengths"].tolist() | |
| unique, first, owners = {}, [], [] | |
| for i, key in enumerate(packed["state_keys"]): | |
| if key not in unique: | |
| unique[key] = len(first) | |
| first.append(i) | |
| reference = first[unique[key]] | |
| n = prefixes[i] | |
| if n != prefixes[reference] or not torch.equal( | |
| packed["inputs"]["input_ids"][i, :n], packed["inputs"]["input_ids"][reference, :n] | |
| ): | |
| raise ValueError("Same state identity compiled to different prefix tokens") | |
| if not torch.equal(positions[:, i, :n], positions[:, reference, :n]): | |
| raise ValueError("Same state identity compiled to different multimodal positions") | |
| owners.append(unique[key]) | |
| ns, nb = max(prefixes), max(n - p for n, p in zip(lengths, prefixes)) | |
| s = hidden.new_zeros((len(first), ns, hidden.shape[-1])) | |
| b = hidden.new_zeros((len(lengths), nb, hidden.shape[-1])) | |
| sp = positions.new_zeros((3, len(first), ns)) | |
| bp = positions.new_zeros((3, len(lengths), nb)) | |
| sm = torch.zeros(s.shape[:2], dtype=torch.bool, device=hidden.device) | |
| bm = torch.zeros(b.shape[:2], dtype=torch.bool, device=hidden.device) | |
| for j, i in enumerate(first): | |
| n = prefixes[i] | |
| s[j, -n:], sp[:, j, -n:], sm[j, -n:] = hidden[i, :n], positions[:, i, :n], True | |
| for i, (p, n) in enumerate(zip(prefixes, lengths)): | |
| b[i, : n - p], bp[:, i, : n - p], bm[i, : n - p] = ( | |
| hidden[i, p:n], | |
| positions[:, i, p:n], | |
| True, | |
| ) | |
| return s, b, sp, bp, sm, bm, torch.tensor(owners, device=hidden.device) | |
| def masks(lm, s, b, sm, bm, owner): | |
| ns, nb, device = s.shape[1], b.shape[1], s.device | |
| causal = lambda n: torch.ones(n, n, dtype=torch.bool, device=device).tril() | |
| eye = lambda n: torch.eye(n, dtype=torch.bool, device=device) | |
| impl = lm.config._attn_implementation | |
| s_allow = (causal(ns)[None] & sm[:, None, :]) | eye(ns)[None] | |
| b_allow = torch.cat( | |
| [ | |
| sm[owner][:, None, :].expand(-1, nb, -1), | |
| (causal(nb)[None] & bm[:, None, :]) | eye(nb)[None], | |
| ], | |
| -1, | |
| ) | |
| ms = { | |
| "full_attention": None | |
| if impl == "sdpa" and bool(sm.all()) | |
| else attention_mask(s_allow, s, impl), | |
| "linear_attention": sm.to(s.dtype), | |
| } | |
| mb = {"full_attention": attention_mask(b_allow, b, impl), "linear_attention": bm.to(b.dtype)} | |
| return ms, mb | |
| def shared_hidden(lm, hidden, positions, packed): | |
| s, b, sp, bp, sm, bm, owner = split_rows(hidden, positions, packed) | |
| ms, mb = masks(lm, s, b, sm, bm, owner) | |
| rs, rb = lm.rotary_emb(s, sp), lm.rotary_emb(b, bp) | |
| for layer in lm.layers: | |
| # Capture the current layer, never a late-bound loop variable in recomputation. | |
| def step(s, b, layer=layer): | |
| prefix = Prefix() | |
| s = layer.forward( | |
| s, | |
| position_embeddings=rs, | |
| attention_mask=ms[layer.block_type], | |
| past_key_values=prefix, | |
| ) | |
| b = layer.forward( | |
| b, | |
| position_embeddings=rb, | |
| attention_mask=mb[layer.block_type], | |
| past_key_values=prefix.branch(owner), | |
| ) | |
| return s * sm[..., None], b | |
| if lm.gradient_checkpointing and lm.training and torch.is_grad_enabled(): | |
| s, b = checkpoint(step, s, b, use_reentrant=False) | |
| else: | |
| s, b = step(s, b) | |
| return lm.norm(b) | |
| class ByteLRU: | |
| def __init__(self, max_bytes): | |
| self.max_bytes, self.bytes = max_bytes, 0 | |
| self.entries = OrderedDict() | |
| self.hits = self.misses = 0 | |
| def get(self, key): | |
| if key not in self.entries: | |
| self.misses += 1 | |
| return None | |
| self.hits += 1 | |
| self.entries.move_to_end(key) | |
| return self.entries[key][0] | |
| def put(self, key, value, size): | |
| if size > self.max_bytes: | |
| return | |
| if key in self.entries: | |
| self.bytes -= self.entries.pop(key)[1] | |
| while self.entries and self.bytes + size > self.max_bytes: | |
| self.bytes -= self.entries.popitem(last=False)[1][1] | |
| self.entries[key] = (value, size) | |
| self.bytes += size | |
| def clear(self): | |
| self.entries.clear() | |
| self.bytes = 0 | |
| class KevRuntime: | |
| """An immutable-model session with byte-bounded image and hybrid prefix caches. | |
| Every round recomputes only the private question history from the original state. | |
| All branches of one state are batched; attention/conv/recurrent prefixes are read-only. | |
| """ | |
| def __init__(self, model, config, prefix_bytes=8 * 1024**3, image_bytes=2 * 1024**3, | |
| enable_native_qk=False): | |
| from .roi_replay import REPLAY_VARIANTS | |
| if config.get('point_variant') in REPLAY_VARIANTS: | |
| raise ValueError('ROI replay uses uncached full rows, not the persistent hybrid-prefix runtime') | |
| if config.get("point_variant") == "visual_tokens": | |
| raise ValueError("Visual-token pointing requires prefix keys; cached branch-only readout is unsupported") | |
| if not enable_native_qk and config.get("readout_kind") in ( | |
| "native_last_qk_mean", | |
| "native_merged_qk256", | |
| "native_last_qk_concat", | |
| ): | |
| raise ValueError("Native QK currently uses the uncached causal-row forward") | |
| if enable_native_qk and (config.get("readout_kind") != "native_last_qk_mean" | |
| or config.get("readout_layer_indices") is not None | |
| or config.get("gradient_checkpointing", False)): | |
| raise ValueError("Experimental native cache requires single-layer mean and no checkpointing") | |
| import hashlib | |
| import json | |
| self.model = model | |
| self.signature = hashlib.sha256(json.dumps(config, sort_keys=True).encode()).hexdigest() | |
| self.versions = tuple((p._version, p.dtype, p.device) for p in model.parameters()) | |
| self.prefix_cache, self.image_cache = ByteLRU(prefix_bytes), ByteLRU(image_bytes) | |
| def clear(self): | |
| self.prefix_cache.clear() | |
| self.image_cache.clear() | |
| def __call__(self, packed): | |
| if packed.get("visual_pointing", False): | |
| raise ValueError("Visual-token pointing requires the uncached causal-row forward") | |
| if self.model.training or self.versions != tuple( | |
| (p._version, p.dtype, p.device) for p in self.model.parameters() | |
| ): | |
| raise ValueError("Create a fresh runtime after changing model weights or training mode") | |
| model, inputs = self.model, packed["inputs"] | |
| device_type = inputs["input_ids"].device.type | |
| precision = ( | |
| torch.is_autocast_enabled(device_type), | |
| str(torch.get_autocast_dtype(device_type)), | |
| torch.get_float32_matmul_precision(), | |
| os.environ.get("TRITON_F32_DEFAULT", "tf32"), | |
| ) | |
| scope = (self.signature, precision) | |
| features = [self.image_cache.get((scope, k)) for k in packed["image_keys"]] | |
| missing = [i for i, f in enumerate(features) if f is None] | |
| if missing: | |
| chunks = inputs["pixel_values"].split(inputs["unique_grid_thw"].prod(-1).tolist()) | |
| fs = model.native.get_image_features( | |
| torch.cat([chunks[i] for i in missing]), | |
| inputs["unique_grid_thw"][missing], | |
| return_dict=True, | |
| ).pooler_output | |
| for i, f in zip(missing, fs): | |
| f = f.clone().contiguous() | |
| features[i] = f | |
| self.image_cache.put( | |
| (scope, packed["image_keys"][i]), f, f.numel() * f.element_size() | |
| ) | |
| hidden, positions = model.embeddings(inputs, features=features) | |
| with model.language_context(hidden.device.type): | |
| return self.language_forward(packed, hidden, positions, scope) | |
| def language_forward(self, packed, hidden, positions, scope): | |
| model, inputs = self.model, packed["inputs"] | |
| s, b, sp, bp, sm, bm, owner = split_rows(hidden, positions, packed) | |
| lm = model.native.language_model | |
| result = torch.zeros_like(b) | |
| for state in range(len(s)): | |
| rows = (owner == state).nonzero().flatten() | |
| first = int(rows[0]) | |
| ns = int(packed["prefix_lengths"][first]) | |
| key = ( | |
| scope, | |
| packed["state_keys"][first], | |
| str(hidden.dtype), | |
| tuple(inputs["input_ids"][first, :ns].tolist()), | |
| tuple(positions[:, first, :ns].flatten().tolist()), | |
| ) | |
| entry = self.prefix_cache.get(key) | |
| if entry is None: | |
| hs = s[state : state + 1, -ns:] | |
| rs = lm.rotary_emb(hs, sp[:, state : state + 1, -ns:]) | |
| entry = [] | |
| for layer in lm.layers: | |
| prefix = Prefix() | |
| hs = layer.forward(hs, position_embeddings=rs, past_key_values=prefix) | |
| entry.append(prefix.compact()) | |
| self.prefix_cache.put(key, entry, sum(p.nbytes() for p in entry)) | |
| # Different states need not inherit the longest unrelated question's | |
| # padding. Their final readouts only index these real branch lengths. | |
| branch_length = int(bm[rows].sum(-1).max()) | |
| hb, mask = b[rows, :branch_length], bm[rows, :branch_length] | |
| rb = lm.rotary_emb(hb, bp[:, rows, :branch_length]) | |
| owners = torch.zeros(len(rows), dtype=torch.long, device=hb.device) | |
| _, mb = masks( | |
| lm, s[state : state + 1, -ns:], hb, sm[state : state + 1, -ns:], mask, owners | |
| ) | |
| for layer, prefix in zip(lm.layers, entry): | |
| hb = layer.forward( | |
| hb, | |
| position_embeddings=rb, | |
| attention_mask=mb[layer.block_type], | |
| past_key_values=prefix.branch(owners), | |
| ) | |
| result[rows, :branch_length] = lm.norm(hb) | |
| if model.pointer is None: | |
| # The cached image prefix is unchanged. Candidates and decisions are | |
| # private text suffix tokens; retain their ORIGINAL M-RoPE coordinates | |
| # while translating tensor indices into the returned suffix buffer. | |
| selected = dict(packed) | |
| for key in ("positions", "option_positions"): | |
| indices = packed[key].clone() | |
| indices[:, 1] -= packed["prefix_lengths"][indices[:, 0]] | |
| if (indices[:, 1] < 0).any(): | |
| raise ValueError("Native suffix cache cannot read keys inside the prefix") | |
| selected[key] = indices | |
| return lm.native_qk(result, selected, lm.rotary_emb(result, bp)) | |
| return model.readout(result, packed, branch_only=True) | |