| """KV-cached incremental decoding for the Dory port (dory/modeling.py). |
| |
| Every (physical attention layer, loop pass) owns its own key/value cache, like |
| vLLM's ``recurrent_kv_cache_mode="per_loop"``: pass 1 of the recurrent block |
| attends over pass-1 keys of earlier tokens, pass 2 over pass-2 keys, and so on. |
| |
| Caches are preallocated once per request: |
| |
| * full-attention layers: ``[B, kv_heads, capacity, head_dim]`` indexed by position; |
| * sliding-window layers: a ring buffer of ``window`` slots, slot = position % window. |
| The port's window keeps keys with ``0 <= i - j < window``, i.e. the current token |
| plus the previous ``window - 1`` tokens, which is exactly what the ring holds. |
| |
| The decode step always attends over the whole buffer with a boolean validity mask |
| (slot <= position), so its shapes are static and it can be captured in a CUDA graph. |
| RoPE is applied to keys before they are cached (as in the port), so no position |
| bookkeeping is needed at read time. Keys are stored after qk-norm; queries get it too. |
| |
| Only batch size 1 is exercised, but nothing assumes it except ``sample``. |
| """ |
|
|
| import math |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| from .modeling import DoryForCausalLM, apply_rope_neox, causal_mask, justnorm, rope_cos_sin, slerp_alpha |
|
|
|
|
| class KVCache: |
| """Per-(layer, loop) caches for one request.""" |
|
|
| def __init__(self, model: DoryForCausalLM, capacity: int, n_loops: int, batch: int = 1, device="cpu", |
| dtype=None): |
| cfg = model.config |
| p = next(model.parameters()) |
| dtype = dtype or p.dtype |
| self.cfg, self.capacity, self.device = cfg, capacity, torch.device(device) |
| self.window = cfg.sliding_window |
| self.pos = 0 |
| self.k, self.v = {}, {} |
| n_in, n_rec = cfg.n_input_layers, cfg.n_recurrent_layers |
| for i, lt in enumerate(cfg.layer_types): |
| if lt == "mlp": |
| continue |
| passes = range(n_loops) if n_in <= i < n_in + n_rec else range(1) |
| cap = self.window if lt == "sliding_attention" else capacity |
| for lp in passes: |
| shape = (batch, cfg.num_key_value_heads, cap, cfg.head_dim) |
| self.k[i, lp] = torch.zeros(shape, device=device, dtype=dtype) |
| self.v[i, lp] = torch.zeros(shape, device=device, dtype=dtype) |
| self.slots = torch.arange(max(capacity, self.window), device=device) |
| |
| pos = torch.arange(capacity, device=device)[None] |
| self.rope = {} |
| for lt in set(cfg.layer_types) - {"mlp"}: |
| c, s = rope_cos_sin(pos, cfg.head_dim, cfg.rope_parameters[lt]["rope_theta"], device) |
| self.rope[lt] = (c[0], s[0]) |
|
|
| def slot_of(self, layer_type, positions): |
| return positions % self.window if layer_type == "sliding_attention" else positions |
|
|
|
|
| def _qkv(mixer, x, cos, sin): |
| b, t, _ = x.shape |
| h, kv, d = mixer.num_heads, mixer.num_kv_heads, mixer.head_dim |
| q = mixer.q_proj(x).view(b, t, h, d) |
| k = mixer.k_proj(x).view(b, t, kv, d) |
| v = mixer.v_proj(x).view(b, t, kv, d) |
| gate = mixer.g_proj(x) |
| q, k = apply_rope_neox(q, cos, sin), apply_rope_neox(k, cos, sin) |
| sqk = (mixer.sqk.float() * mixer.sqk_factor).view(h, d) |
| q = (sqk.square() * justnorm(q).float()).to(q.dtype) |
| k = justnorm(k) |
| return q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), gate |
|
|
|
|
| def _finish(mixer, out, gate): |
| b, h, t, d = out.shape |
| out = out.transpose(1, 2).reshape(b, t, h * d) |
| sg = (mixer.s_gate.float() * mixer.s_gate_factor).square() |
| out = (out.float() * torch.sigmoid(gate.float() * sg)).to(out.dtype) |
| return mixer.o_proj(justnorm(out)) |
|
|
|
|
| def _attend(q, k, v, mask, d): |
| return F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=1.0 / math.sqrt(d), enable_gqa=True) |
|
|
|
|
| def attn_prefill(mixer, x, cache: KVCache, key, mask): |
| """x: [B, T, H] for positions 0..T-1 of an empty cache. Writes the cache, returns the mixer output.""" |
| lt = mixer.layer_type |
| t = x.shape[1] |
| positions = torch.arange(t, device=x.device) |
| cos, sin = cache.rope[lt] |
| q, k, v, gate = _qkv(mixer, x, cos[None, positions], sin[None, positions]) |
| kc, vc = cache.k[key], cache.v[key] |
| if lt == "sliding_attention" and t > cache.window: |
| keep = t - cache.window |
| slots = cache.slot_of(lt, positions[keep:]) |
| kc[:, :, slots], vc[:, :, slots] = k[:, :, keep:], v[:, :, keep:] |
| else: |
| slots = cache.slot_of(lt, positions) |
| kc[:, :, slots], vc[:, :, slots] = k, v |
| out = _attend(q, k, v, mask, mixer.head_dim) |
| return _finish(mixer, out, gate) |
|
|
|
|
| def attn_decode(mixer, x, cache: KVCache, key, pos, valid): |
| """x: [B, 1, H] at scalar-tensor position `pos`; `valid[lt]` is the bool key mask [1,1,1,cap].""" |
| lt = mixer.layer_type |
| cos, sin = cache.rope[lt] |
| idx = pos.view(1) |
| q, k, v, gate = _qkv(mixer, x, cos.index_select(0, idx)[None], sin.index_select(0, idx)[None]) |
| slot = cache.slot_of(lt, pos) |
| kc, vc = cache.k[key], cache.v[key] |
| kc.index_copy_(2, slot.view(1), k) |
| vc.index_copy_(2, slot.view(1), v) |
| out = _attend(q, kc, vc, valid[lt], mixer.head_dim) |
| return _finish(mixer, out, gate) |
|
|
|
|
| def _backbone(model, h, n_loops, run_layer): |
| """Shared layer schedule: input layers, n_loops recurrent passes, output layers.""" |
| cfg = model.config |
| layers = model.backbone.layers |
| n_in, n_rec = cfg.n_input_layers, cfg.n_recurrent_layers |
| for i in range(n_in): |
| h = run_layer(i, 0, h) |
| input_h = h |
| state = input_h |
| for lp in range(n_loops): |
| h = justnorm(state + input_h) |
| for i in range(n_in, n_in + n_rec): |
| h = run_layer(i, lp, h) |
| state = justnorm(h) |
| h = state |
| for i in range(n_in + n_rec, len(layers)): |
| h = run_layer(i, 0, h) |
| return h |
|
|
|
|
| def _layer(model, i, lp, h, attn): |
| layer = model.backbone.layers[i] |
| if layer.is_attention: |
| out = attn(layer.mixer, h, (i, lp)) |
| return slerp_alpha(h, out, layer.attn_alpha, layer.alpha_scale) |
| return slerp_alpha(h, layer.mixer(h), layer.mlp_alpha, layer.alpha_scale) |
|
|
|
|
| @torch.no_grad() |
| def prefill(model, input_ids, cache: KVCache, n_loops): |
| """Run the prompt (positions 0..T-1), fill the caches, return logits of the last token [B, V].""" |
| assert cache.pos == 0 |
| t = input_ids.shape[1] |
| dev = input_ids.device |
| masks = {lt: causal_mask(t, model.config.sliding_window if lt == "sliding_attention" else None, dev) |
| for lt in set(model.config.layer_types) - {"mlp"}} |
|
|
| def attn(mixer, h, key): |
| return attn_prefill(mixer, h, cache, key, masks[mixer.layer_type]) |
|
|
| h = model.backbone.embeddings(input_ids) |
| h = _backbone(model, h, n_loops, lambda i, lp, h: _layer(model, i, lp, h, attn)) |
| cache.pos = t |
| return model.logits(h[:, -1]) |
|
|
|
|
| @torch.no_grad() |
| def decode_step(model, token, cache: KVCache, n_loops, pos): |
| """One token [B, 1] at position `pos` (0-d long tensor). Returns logits [B, V].""" |
| slots = cache.slots |
| valid = {"full_attention": (slots[: cache.capacity] <= pos)[None, None, None], |
| "sliding_attention": (slots[: cache.window] <= pos)[None, None, None]} |
|
|
| def attn(mixer, h, key): |
| return attn_decode(mixer, h, cache, key, pos, valid) |
|
|
| h = model.backbone.embeddings(token) |
| h = _backbone(model, h, n_loops, lambda i, lp, h: _layer(model, i, lp, h, attn)) |
| return model.logits(h[:, -1]) |
|
|
|
|
| def sample(logits, temperature=1.0, top_p=0.95, generator=None): |
| """Nucleus sampling on [1, V] float logits; temperature 0 = greedy.""" |
| if temperature <= 0: |
| return logits.argmax(-1, keepdim=True) |
| probs = torch.softmax(logits / temperature, dim=-1) |
| if top_p < 1.0: |
| sp, si = probs.sort(descending=True, dim=-1) |
| cum = sp.cumsum(-1) |
| sp[cum - sp > top_p] = 0.0 |
| choice = torch.multinomial(sp / sp.sum(-1, keepdim=True), 1, generator=generator) |
| return si.gather(-1, choice) |
| return torch.multinomial(probs, 1, generator=generator) |
|
|
|
|
| class Decoder: |
| """Optional CUDA-graph wrapper around decode_step with static input/position buffers.""" |
|
|
| def __init__(self, model, cache, n_loops, use_graph=False): |
| self.model, self.cache, self.n_loops = model, cache, n_loops |
| dev = cache.device |
| self.tok = torch.zeros(1, 1, dtype=torch.long, device=dev) |
| self.pos = torch.zeros((), dtype=torch.long, device=dev) |
| self.graph = None |
| self.use_graph = use_graph and dev.type == "cuda" |
|
|
| def __call__(self, token, pos): |
| self.tok.fill_(int(token)) |
| self.pos.fill_(int(pos)) |
| if not self.use_graph: |
| return decode_step(self.model, self.tok, self.cache, self.n_loops, self.pos) |
| if self.graph is None: |
| try: |
| self._capture() |
| except Exception as e: |
| print(f"CUDA graph capture failed ({type(e).__name__}: {e}); decoding eagerly") |
| self.use_graph = False |
| return decode_step(self.model, self.tok, self.cache, self.n_loops, self.pos) |
| self.graph.replay() |
| return self.out |
|
|
| def _capture(self): |
| |
| |
| s = torch.cuda.Stream() |
| s.wait_stream(torch.cuda.current_stream()) |
| with torch.cuda.stream(s): |
| for _ in range(2): |
| decode_step(self.model, self.tok, self.cache, self.n_loops, self.pos) |
| torch.cuda.current_stream().wait_stream(s) |
| graph = torch.cuda.CUDAGraph() |
| with torch.cuda.graph(graph): |
| self.out = decode_step(self.model, self.tok, self.cache, self.n_loops, self.pos) |
| self.graph = graph |
|
|
|
|
| @torch.no_grad() |
| def generate_stream(model, input_ids, max_new_tokens, n_loops=2, temperature=1.0, top_p=0.95, |
| eos_ids=(2, 11), seed=None, use_graph=False, return_logits=False): |
| """Yield one generated token id (int) at a time; stops after an EOS id (which is yielded). |
| |
| With return_logits=True yields (token, logits[1, V]) pairs instead (for tests).""" |
| dev = input_ids.device |
| t = input_ids.shape[1] |
| cache = KVCache(model, capacity=t + max_new_tokens, n_loops=n_loops, device=dev) |
| gen = None |
| if seed is not None: |
| gen = torch.Generator(device=dev).manual_seed(seed) |
| logits = prefill(model, input_ids, cache, n_loops) |
| step = Decoder(model, cache, n_loops, use_graph=use_graph) |
| for i in range(max_new_tokens): |
| nxt = sample(logits, temperature, top_p, gen) |
| tok = int(nxt[0, 0]) |
| yield (tok, logits.clone()) if return_logits else tok |
| if tok in eos_ids: |
| break |
| if i + 1 < max_new_tokens: |
| logits = step(tok, t + i) |
|
|
