File size: 9,827 Bytes
533964d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 | """Shared-prefix scoring with a bounded memory fork, used by both engines.
Rows that start with the same tokens (one state, one question per row) run the shared prefix once and then only the
question suffixes, against a copy of the prefix cache. Until 1.1.0 the copy was made with
`cache.reorder_cache(zeros(n))`, which expands the prefix to all n rows at once: a 31k-token state with 32 questions
costs n times the prefix cache, which is 133 GB on a 31B model and is wasteful on the 2B. Here the prefix cache is
forked in chunks of m rows, m chosen so that one fork fits a byte budget, and the rows are scored chunk by chunk.
The budget is `DECIDER_SHARED_FORK_GB` (8 GB), capped at half of the memory currently free on the device (device-free
plus the caching allocator's reserved-but-unused blocks); `m = clamp(budget // prefix_bytes, 1, n)`. Chunking changes
which rows share a forward and how far each chunk's suffixes are padded, so answers can move by the usual bf16
reduction-order amount; they do not depend on it mathematically (right padding, causal layers).
Cache layout. transformers 5.17 keeps one layer object per layer in `cache.layers`. An attention layer carries
`keys` and `values` tensors `[batch, heads, seq, head_dim]`; a sparse-attention layer adds `indexer_keys`
`[batch, seq, dim]`; a linear-attention layer carries `conv_states` and `recurrent_states` dicts of tensors whose first
dimension is the batch. Qwen3.5 gives a `DynamicCache` whose 24 layers are a mix of the two objects. `cache_state_tensors`
enumerates whatever of `STATE_NAMES` is present rather than assuming a layout, and `fork_cache` builds a new cache
object from them without touching the original (the linear-attention layers update their states with `copy_()`, so a
fork that shared storage with the prefix would corrupt it for the next chunk).
A layout we do not know is not chunked. If a layer holds any other tensor attribute, or a dict of tensors, outside
`STATE_NAMES`, we cannot tell whether it carries a batch dimension, so `score_shared` takes one fork of all n rows --
the 1.1.0 behaviour, correct but unbounded -- and counts it in `engine.stats["shared_unchunked_layout"]`.
The prefix and suffix forwards are eager, at request-specific shapes, and are correct only with the cuDNN SDPA backend
off (`decider.engine.set_attention_backend_policy`, applied in both engines' `__init__`).
"""
import copy
import os
import torch
import torch.nn.functional as F
from decider.engine import fill_ids, read_slots
from decider.temperature import item_slice, slot_temperatures
DEFAULT_FORK_GB = 8.0
STATE_NAMES = ("keys", "values", "indexer_keys", "conv_states", "recurrent_states")
def common_prefix_len(ids):
"""Length of the longest common prefix of the rows, capped one token below the shortest row."""
lcp = 0
short = min(len(x) for x in ids) - 1
while lcp < short and all(x[lcp] == ids[0][lcp] for x in ids):
lcp += 1
return lcp
def _get(container, key):
return container[key] if isinstance(container, dict) else getattr(container, key)
def _set(container, key, value):
if isinstance(container, dict):
container[key] = value
else:
setattr(container, key, value)
def cache_state_tensors(cache):
"""Every tensor in the cache whose first dimension is the batch, as (container, key) pairs.
Covers `STATE_NAMES`: the attention layers' `keys`/`values`, a sparse-attention layer's `indexer_keys` and the
linear-attention layers' `conv_states`/`recurrent_states`, whether those are dicts of tensors (transformers 5.17) or
single tensors, and ignores anything not present."""
out = []
for layer in getattr(cache, "layers", None) or []:
for attr in STATE_NAMES:
v = getattr(layer, attr, None)
if isinstance(v, torch.Tensor):
if v.numel():
out.append((layer, attr))
elif isinstance(v, dict):
for k, t in v.items():
if isinstance(t, torch.Tensor) and t.numel():
out.append((v, k))
return out
def cache_row_bytes(cache):
"""Bytes one row of this cache holds, summed over the layers and over `STATE_NAMES`."""
return sum(_get(c, k).nbytes // max(_get(c, k).shape[0], 1) for c, k in cache_state_tensors(cache))
def unknown_state_names(cache):
"""Attribute names the cache's layers hold that carry tensors and are not in `STATE_NAMES`.
A tensor outside the enumerated set may or may not have a batch dimension, and `fork_cache` would leave it at one
row. When this is not empty the caller must not chunk."""
out = set()
for layer in getattr(cache, "layers", None) or []:
for name, v in vars(layer).items():
if name in STATE_NAMES:
continue
if isinstance(v, torch.Tensor) or (isinstance(v, dict) and any(isinstance(t, torch.Tensor) for t in v.values())):
out.add(name)
return out
def fork_budget_bytes(device=None, gb=None):
"""DECIDER_SHARED_FORK_GB, capped at half of the memory free on the device right now.
The cap applies only when the CUDA memory queries succeed: on CPU, on MPS and when the driver does not know the
device string, the configured budget is kept as it is."""
gb = float(os.environ.get("DECIDER_SHARED_FORK_GB", DEFAULT_FORK_GB)) if gb is None else float(gb)
budget = int(gb * (1 << 30))
try:
free, _ = torch.cuda.mem_get_info(device)
free += torch.cuda.memory_reserved(device) - torch.cuda.memory_allocated(device)
budget = min(budget, free // 2)
except Exception: # no CUDA device, or a device string the driver does not know
pass
return max(int(budget), 1)
def chunk_rows(prefix_bytes, n, budget_bytes=None, device=None):
"""Rows per fork: as many copies of the prefix cache as the budget holds, at least 1 and at most n. One row is the
minimum even when a single copy is over the budget: the budget bounds the fork, it cannot make it free."""
if budget_bytes is None:
budget_bytes = fork_budget_bytes(device)
if prefix_bytes <= 0:
return max(1, int(n))
return max(1, min(int(n), int(budget_bytes // prefix_bytes)))
def fork_cache(cache, m, row=0):
"""A new cache holding `m` copies of `cache`'s row `row`. The original is not read from again and not modified."""
fork = copy.copy(cache)
layers = getattr(cache, "layers", None)
if layers is not None:
new = []
for layer in layers:
nl = copy.copy(layer)
for k, v in list(vars(nl).items()): # the per-state dicts are mutated by the forward: give the fork its own
if isinstance(v, dict):
setattr(nl, k, dict(v))
new.append(nl)
fork.layers = new
idx = {}
for container, key in cache_state_tensors(fork):
t = _get(container, key)
i = idx.get(t.device)
if i is None:
i = idx[t.device] = torch.full((m,), row, dtype=torch.long, device=t.device)
_set(container, key, t.index_select(0, i)) # a fresh contiguous tensor: the fork never shares storage
return fork
def _count(engine, key):
stats = getattr(engine, "stats", None)
if isinstance(stats, dict):
stats[key] = stats.get(key, 0) + 1
@torch.no_grad()
def score_shared(engine, items, temperature=1.0, min_prefix=192, budget_bytes=None, rows_per_fork=None):
"""Score `items` through the shared prefix. -> one probability tensor per item, in item order, or None when the
request does not qualify (fewer than two rows, or a common prefix below `min_prefix`) and the caller should use
`score_items`.
`temperature`: a number, or one entry per item (decider.temperature.slot_temperatures).
`rows_per_fork` forces the chunk size; it exists for the tests that compare chunked against unchunked answers."""
ids = [it["ids"] for it in items]
n = len(ids)
if n < 2:
return None
lcp = common_prefix_len(ids)
if lcp < min_prefix:
return None
slot_temperatures(temperature, items) # a length mismatch fails before any forward
core, W, dev, pad = engine.core, engine.W, engine.dev, engine.tok.pad_token_id
pre = torch.tensor(ids[0][:lcp], device=dev)[None]
cache = core(input_ids=pre, use_cache=True).past_key_values
unknown = unknown_state_names(cache)
if unknown: # a state we cannot fork row by row: one fork of everything, as in 1.1.0
_count(engine, "shared_unchunked_layout")
m = n
elif rows_per_fork:
m = int(rows_per_fork)
else:
m = chunk_rows(cache_row_bytes(cache), n, budget_bytes, dev)
m = max(1, min(m, n))
out = []
for i in range(0, n, m):
part = items[i:i + m]
b = len(part)
fork = fork_cache(cache, b)
Ts = max(len(it["ids"]) for it in part) - lcp
suf = fill_ids([it["ids"][lcp:] for it in part], b, Ts, pad)
h = core(input_ids=suf.to(dev), past_key_values=fork, use_cache=True).last_hidden_state
rows = [j for j, it in enumerate(part) for _ in it["slots"]]
sl = [s - lcp for it in part for s in it["slots"]]
idx = torch.tensor([rows, sl], device=dev)
out += read_slots(F.linear(h[idx[0], idx[1]], W).float()[:, None, :], list(range(len(rows))), [0] * len(rows),
[k for it in part for k in it["nopts"]],
slot_temperatures(item_slice(temperature, i, i + b), part), [len(it["slots"]) for it in part])
del fork, h, suf # drop this chunk's fork before the next one is built
return out
|