Spaces:
Sleeping
Sleeping
| """Layer-addressable decoder with a static per-layer KV cache. | |
| The wrapper drives the submodules of a Hugging Face dense decoder (Qwen3, Llama) | |
| one layer range at a time. Every layer owns a KV buffer of size `max_len`; a | |
| token's K/V are written at its absolute position, so layers may hold different | |
| sets of positions while drafting (early exits) and verification can overwrite | |
| approximate entries in place. | |
| Attention for a chunk of `n` new rows at positions [s, s+n) reads K/V | |
| positions [0, s+n) of the same layer, with a causal mask inside the chunk. | |
| """ | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| import torch | |
| import torch.nn.functional as F | |
| def _rotate_half(x: torch.Tensor) -> torch.Tensor: | |
| x1, x2 = x.chunk(2, dim=-1) | |
| return torch.cat((-x2, x1), dim=-1) | |
| class ModelSpec: | |
| num_layers: int | |
| hidden_size: int | |
| vocab_size: int | |
| num_heads: int | |
| num_kv_heads: int | |
| head_dim: int | |
| has_qk_norm: bool | |
| def use_row_invariant_blas() -> None: | |
| """Select rocBLAS on ROCm: row-invariant for 1..64 rows on MI210 (job 893629), | |
| unlike hipBLASLt (o_proj breaks at 17 rows, q_proj at 24). No-op elsewhere.""" | |
| if torch.cuda.is_available() and torch.version.hip: | |
| torch.backends.cuda.preferred_blas_library("cublas") # maps to rocBLAS on ROCm | |
| class LayeredDecoder: | |
| """Run a frozen HF causal LM layer by layer with a positional KV cache.""" | |
| def __init__(self, hf_model, max_len: int, canonical_attention: bool = True, | |
| max_rows: int = 16): | |
| self.hf = hf_model | |
| # Canonical attention: every query row uses the one-query SDPA call over its | |
| # exact prefix, as in one-token AR. On MI210/bf16 the batched masked call | |
| # is not row-invariant while all projections and norms are (job 893599). | |
| self.canonical_attention = canonical_attention | |
| # GEMMs were verified row-invariant up to 16 rows (job 893599); larger | |
| # chunks are processed in causal sub-chunks of at most `max_rows` rows. | |
| self.max_rows = max_rows | |
| base = hf_model.model | |
| cfg = hf_model.config | |
| self.embed_tokens = base.embed_tokens | |
| self.layers = base.layers | |
| self.norm = base.norm | |
| self.lm_head = hf_model.lm_head | |
| self.rotary_emb = base.rotary_emb | |
| attn0 = self.layers[0].self_attn | |
| self.spec = ModelSpec( | |
| num_layers=len(self.layers), | |
| hidden_size=cfg.hidden_size, | |
| vocab_size=cfg.vocab_size, | |
| num_heads=cfg.num_attention_heads, | |
| num_kv_heads=cfg.num_key_value_heads, | |
| head_dim=attn0.head_dim, | |
| has_qk_norm=hasattr(attn0, "q_norm"), | |
| ) | |
| self.scaling = attn0.scaling | |
| self.max_len = max_len | |
| p = next(hf_model.parameters()) | |
| self.device, self.dtype = p.device, p.dtype | |
| shape = (self.spec.num_layers, self.spec.num_kv_heads, max_len, self.spec.head_dim) | |
| self.k_cache = torch.zeros(shape, device=self.device, dtype=self.dtype) | |
| self.v_cache = torch.zeros(shape, device=self.device, dtype=self.dtype) | |
| # Rotary tables for every position, computed once. | |
| pos = torch.arange(max_len, device=self.device).unsqueeze(0) | |
| dummy = torch.zeros(1, 1, self.spec.head_dim, device=self.device, dtype=self.dtype) | |
| cos, sin = self.rotary_emb(dummy, pos) | |
| self.cos, self.sin = cos[0], sin[0] # (max_len, head_dim) | |
| self.counters = {"layer_rows": 0, "kv_fill_rows": 0, "vocab_proj_rows": 0} | |
| # ------------------------------------------------------------------ helpers | |
| def reset_counters(self) -> None: | |
| for k in self.counters: | |
| self.counters[k] = 0 | |
| def embed(self, ids: torch.Tensor) -> torch.Tensor: | |
| """ids: (n,) -> (n, d)""" | |
| return self.embed_tokens(ids) | |
| def _rope(self, x: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: | |
| # x: (heads, n, head_dim); pos: (n,) absolute positions | |
| cos = self.cos.index_select(0, pos).unsqueeze(0) | |
| sin = self.sin.index_select(0, pos).unsqueeze(0) | |
| return x * cos + _rotate_half(x) * sin | |
| def _kv(self, attn, x_normed: torch.Tensor, pos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: | |
| n = x_normed.shape[0] | |
| hd = self.spec.head_dim | |
| k = attn.k_proj(x_normed).view(n, -1, hd) | |
| if self.spec.has_qk_norm: | |
| k = attn.k_norm(k) | |
| k = self._rope(k.transpose(0, 1), pos) | |
| v = attn.v_proj(x_normed).view(n, -1, hd).transpose(0, 1) | |
| return k, v | |
| def _positions(self, start: int, n: int, positions) -> tuple[torch.Tensor, list[int]]: | |
| if positions is None: | |
| plist = list(range(start, start + n)) | |
| else: | |
| plist = [int(p) for p in positions] | |
| assert len(plist) == n and all(a < b for a, b in zip(plist, plist[1:])), "ascending positions" | |
| return torch.tensor(plist, device=self.device), plist | |
| # ------------------------------------------------------------------ layers | |
| def run_layers(self, h: torch.Tensor, start: int, first: int, last: int, | |
| positions=None) -> torch.Tensor: | |
| """Apply layers [first, last) to hidden rows h (n, d). | |
| Rows sit at positions [start, start+n), or at the given ascending | |
| `positions` (rows of a draft block caught up to a deeper layer). K/V of | |
| every processed layer are written at those positions; row i attends to | |
| positions [0, pos_i] of the same layer. | |
| """ | |
| n = h.shape[0] | |
| pos, plist = self._positions(start, n, positions) | |
| if self.canonical_attention and n > self.max_rows: | |
| m = self.max_rows # causal: earlier rows never depend on later rows | |
| return torch.cat([self.run_layers(h[i:i + m], 0, first, last, plist[i:i + m]) | |
| for i in range(0, n, m)]) | |
| hd = self.spec.head_dim | |
| end = plist[-1] + 1 | |
| mask = None | |
| if n > 1 and not self.canonical_attention: | |
| mask = torch.arange(0, end, device=h.device).unsqueeze(0) <= pos.unsqueeze(1) # (n, end) | |
| for li in range(first, last): | |
| layer = self.layers[li] | |
| attn = layer.self_attn | |
| x = layer.input_layernorm(h) | |
| q = attn.q_proj(x).view(n, -1, hd) | |
| if self.spec.has_qk_norm: | |
| q = attn.q_norm(q) | |
| q = self._rope(q.transpose(0, 1), pos) | |
| k, v = self._kv(attn, x, pos) | |
| self.k_cache[li].index_copy_(1, pos, k) | |
| self.v_cache[li].index_copy_(1, pos, v) | |
| if n == 1 or not self.canonical_attention: | |
| o = F.scaled_dot_product_attention( | |
| q.unsqueeze(0), | |
| self.k_cache[li, :, :end].unsqueeze(0), | |
| self.v_cache[li, :, :end].unsqueeze(0), | |
| attn_mask=mask, | |
| scale=self.scaling, | |
| enable_gqa=True, | |
| )[0] | |
| else: | |
| o = torch.cat([F.scaled_dot_product_attention( | |
| q[:, i:i + 1].unsqueeze(0), | |
| self.k_cache[li, :, :p + 1].unsqueeze(0), | |
| self.v_cache[li, :, :p + 1].unsqueeze(0), | |
| scale=self.scaling, | |
| enable_gqa=True, | |
| )[0] for i, p in enumerate(plist)], dim=1) | |
| o = attn.o_proj(o.transpose(0, 1).reshape(n, -1)) | |
| h = h + o | |
| h = h + layer.mlp(layer.post_attention_layernorm(h)) | |
| self.counters["layer_rows"] += n * (last - first) | |
| return h | |
| def fill_kv_from_hidden(self, h: torch.Tensor, start: int, first: int, last: int) -> None: | |
| """CALM-style state copy (regime B): K/V of layers [first, last) computed from h itself. | |
| Used for tokens that exited early, so that deeper draft tokens can attend | |
| to them. Only affects draft quality, never verified outputs. | |
| """ | |
| pos, _ = self._positions(start, h.shape[0], None) | |
| for li in range(first, last): | |
| layer = self.layers[li] | |
| k, v = self._kv(layer.self_attn, layer.input_layernorm(h), pos) | |
| self.k_cache[li].index_copy_(1, pos, k) | |
| self.v_cache[li].index_copy_(1, pos, v) | |
| self.counters["kv_fill_rows"] += h.shape[0] * max(0, last - first) | |
| def final_logits(self, h: torch.Tensor) -> torch.Tensor: | |
| self.counters["vocab_proj_rows"] += h.shape[0] | |
| return self.lm_head(self.norm(h)).float() | |
| # ------------------------------------------------------------------ full passes | |
| def prefill(self, ids: torch.Tensor) -> torch.Tensor: | |
| """Process the whole prompt; returns last-position logits (V,). | |
| The prompt is shared by every decoding method, so it uses the fast batched | |
| attention (row invariance is only needed for generated rows). | |
| """ | |
| canon = self.canonical_attention | |
| self.canonical_attention = False | |
| try: | |
| h = self.run_layers(self.embed(ids), 0, 0, self.spec.num_layers) | |
| finally: | |
| self.canonical_attention = canon | |
| return self.final_logits(h[-1:])[0] | |
| def full_step(self, ids: torch.Tensor, start: int) -> torch.Tensor: | |
| """All layers on rows ids at positions [start, ...); returns logits (n, V).""" | |
| h = self.run_layers(self.embed(ids), start, 0, self.spec.num_layers) | |
| return self.final_logits(h) | |