| """Differentiable fixed-prefix execution; one working query after the first pass.""" |
| import torch |
| from torch import Tensor, nn |
| from .memory import PromptMemory |
| from torch.utils.checkpoint import checkpoint |
| from transformers.models.qwen2.modeling_qwen2 import apply_rotary_pos_emb, eager_attention_forward |
| from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS |
|
|
|
|
| def pick(x: Tensor, positions: list[int]) -> Tensor: |
| return torch.stack([x[row, position:position + 1] for row, position in enumerate(positions)]) |
|
|
|
|
| def layer_compute(layer: nn.Module, hidden: Tensor, embeddings: tuple[Tensor, Tensor], |
| mask: Tensor | None, positions: list[int] | None = None, |
| cached: tuple[Tensor, Tensor] | None = None) -> tuple[Tensor, Tensor, Tensor]: |
| """Pure operation so non-reentrant checkpointing preserves prefix gradients.""" |
| attention = layer.self_attn |
| residual = hidden |
| x = layer.input_layernorm(hidden) |
| shape = (*x.shape[:-1], -1, attention.head_dim) |
| q = attention.q_proj(x).view(shape).transpose(1, 2) |
| k = attention.k_proj(x).view(shape).transpose(1, 2) |
| v = attention.v_proj(x).view(shape).transpose(1, 2) |
| q, k = apply_rotary_pos_emb(q, k, *embeddings) |
| if cached is not None: |
| keys, values = cached |
| active = (torch.arange(keys.shape[-2], device=x.device)[None, :] == |
| torch.tensor(positions, device=x.device)[:, None])[:, None, :, None] |
| k = torch.where(active, k, keys) |
| v = torch.where(active, v, values) |
| interface = ALL_ATTENTION_FUNCTIONS.get_interface(attention.config._attn_implementation, eager_attention_forward) |
| output, _ = interface(attention, q, k, v, mask, |
| dropout=attention.attention_dropout if layer.training else 0.0, |
| scaling=attention.scaling, sliding_window=None) |
| output = output.reshape(*x.shape[:-1], -1).contiguous() |
| hidden = residual + attention.o_proj(output) |
| hidden = hidden + layer.mlp(layer.post_attention_layernorm(hidden)) |
| return hidden, k, v |
|
|
|
|
| def fixed_prefix_layer(layer: nn.Module, hidden: Tensor, embeddings: tuple[Tensor, Tensor], |
| mask: Tensor | None, memory: PromptMemory, index: int, |
| use_checkpoint: bool) -> Tensor: |
| """Reuse prefix K/V without detaching it or caching across model forwards.""" |
| positions = memory.write_mask[..., 0].long().argmax(-1).tolist() |
| cached = None |
| first = memory.output is None |
| if not first: |
| hidden = pick(hidden, positions) |
| embeddings = tuple(pick(e.expand(memory.write_mask.shape[0], -1, -1), positions) for e in embeddings) |
| cached = memory.keys_values[index] |
| if mask is None: |
| keep = torch.arange(cached[0].shape[-2], device=hidden.device)[None, :] <= torch.tensor(positions, device=hidden.device)[:, None] |
| mask = torch.zeros_like(keep, dtype=hidden.dtype).masked_fill(~keep, torch.finfo(hidden.dtype).min)[:, None, None, :] |
| else: |
| mask = torch.stack([mask[row if mask.shape[0] > 1 else 0, :, pos:pos + 1, :] for row, pos in enumerate(positions)]) |
| def compute(x, cos, sin, *kv): |
| return layer_compute(layer, x, (cos, sin), mask, positions, tuple(kv) if kv else None) |
| args = (hidden, *embeddings, *(cached or ())) |
| result, keys, values = (checkpoint(compute, *args, use_reentrant=False) |
| if use_checkpoint and torch.is_grad_enabled() else compute(*args)) |
| if first: |
| memory.keys_values.append((keys, values)) |
| memory.layer_outputs.append(result) |
| return result |
| return torch.where(memory.write_mask, result, memory.layer_outputs[index]) |
|
|
