Spaces:
Running on Zero
Running on Zero
File size: 4,711 Bytes
36cdb93 | 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 | """Preallocated K/V cache for block-causal streaming.
Replaces streaming_blocks.KVFrameCache, which stored a deque of per-frame tensors
and rebuilt the context with `torch.cat` on every `get()`. The profile in
PROGRESS.md §6 counted 302 `aten::cat` calls per single frame-forward.
Here the cache is one preallocated buffer per layer and the context is a *view*,
so steady-state streaming does no allocation and no copying beyond writing the
new frame's K/V.
Commit semantics matter for multi-step denoising. The current chunk's K/V changes
at every denoising step, so it must not be committed until the chunk is final:
for step in steps: # chunk still noisy
cache.write(layer, k, v) # scratch region at [len, len+S)
ctx_k, ctx_v = cache.context(layer, S) # past + current, one view
cache.commit(S) # chunk finalised, becomes past
Because the cache only ever holds past + current frames, full attention over that
view IS block-causal -- bidirectional within the current chunk, unrestricted over
the past, and no future keys exist to leak. No mask is needed; see
tests/test_attention_fallback.py case 5.
"""
import torch
class StreamingKVCache:
"""Per-layer ring-free preallocated K/V store with an uncommitted scratch tail."""
def __init__(self, num_layers, max_tokens, num_heads, head_dim,
device, dtype=torch.bfloat16, scratch_tokens=None):
self.num_layers = num_layers
self.max_tokens = max_tokens
self.scratch = scratch_tokens if scratch_tokens is not None else max_tokens
cap = max_tokens + self.scratch
self.k = torch.zeros(num_layers, cap, num_heads, head_dim,
device=device, dtype=dtype)
self.v = torch.zeros(num_layers, cap, num_heads, head_dim,
device=device, dtype=dtype)
self.length = 0 # committed tokens
self._pending = 0 # tokens written to scratch, not yet committed
# Monotone count of eviction events. `BufferPrefixKV` reconstructs a
# recorded step's prefix as `buffer[:upto]`, which is only that step's
# actual prefix if nothing has been evicted since it was recorded --
# eviction shifts surviving keys down and discards the oldest outright.
# The self-forcing trainer compares this counter against the value it
# stamped on each recorded block and refuses to take the gradient on a
# block whose prefix has since moved. See stream.trim_to_window.
self.evictions = 0
def reset(self):
self.length = 0
self._pending = 0
self.evictions = 0
def write(self, layer, k, v):
"""Write the current (uncommitted) chunk's K/V at [length, length+S)."""
s = k.shape[1]
if self.length + s > self.k.shape[1]:
raise RuntimeError(
f'K/V cache overflow: {self.length} committed + {s} pending > '
f'capacity {self.k.shape[1]}. Raise max_tokens or evict.')
self.k[layer, self.length:self.length + s].copy_(k[0])
self.v[layer, self.length:self.length + s].copy_(v[0])
self._pending = s
def context(self, layer, s):
"""View of past + current: ([1, length+s, n, d], same for v). No copy."""
end = self.length + s
return (self.k[layer, :end].unsqueeze(0),
self.v[layer, :end].unsqueeze(0))
def commit(self, s):
"""Promote the pending chunk to committed history."""
self.length += s
self._pending = 0
def evict_front(self, s, protect=0):
"""Drop `s` tokens from the front of the evictable region.
`protect` pins the first `protect` tokens permanently -- used to keep the
WORLD block resident while sliding a window over the event stream, which
is the behaviour the papers describe (W is persistent, events stream).
Eviction therefore removes the OLDEST EVENTS, not the world.
"""
if s <= 0:
return
evictable = self.length - protect
s = min(s, max(0, evictable))
if s == 0:
return
keep = self.length - protect - s # events surviving the eviction
if keep > 0:
src = protect + s
self.k[:, protect:protect + keep].copy_(self.k[:, src:src + keep])
self.v[:, protect:protect + keep].copy_(self.v[:, src:src + keep])
self.length = protect + keep
self.evictions += 1
@property
def num_tokens(self):
return self.length
def memory_bytes(self):
return self.k.numel() * self.k.element_size() * 2
|