File size: 2,861 Bytes
c83c617 | 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 | """V4 functional single-token KV graph and local TorchScript validation."""
import math
import torch
from torch import nn
from vimeml.deployment.graph import trace_graph
CACHE_SHAPE = (7, 1, 6, 128, 64)
class KVStepGraphV4(nn.Module):
"""Explicit functional FP32 buffers; caller owns snapshots and rollback."""
def __init__(self, model):
super().__init__()
self.model = model.eval()
self.register_buffer("slots", torch.arange(128))
def forward(self, input_ids, position, key_cache, value_cache):
position = position.to(torch.int64)
hidden = self.model.token_embedding(input_ids) + self.model.position_embedding(position).unsqueeze(0)
insert = (self.slots == position[0]).reshape(1, 1, 128, 1)
valid = (self.slots <= position[0]).reshape(1, 1, 1, 128)
keys, values = [], []
for layer, block in enumerate(self.model.blocks):
qkv = block.attention.qkv(block.attention_norm(hidden)).reshape(1, 1, 3, 6, 64)
query, key, value = qkv.permute(2, 0, 3, 1, 4).unbind(0)
updated_key = torch.where(insert, key, key_cache[layer])
updated_value = torch.where(insert, value, value_cache[layer])
scores = torch.matmul(query, updated_key.transpose(-2, -1)) / math.sqrt(64)
probabilities = torch.softmax(scores.masked_fill(~valid, -torch.inf), dim=-1)
attended = torch.matmul(probabilities, updated_value).transpose(1, 2).reshape(1, 1, 384)
hidden = hidden + block.attention.projection(attended)
hidden = hidden + block.ffn(block.ffn_norm(hidden))
keys.append(updated_key)
values.append(updated_value)
return self.model.lm_head(self.model.final_norm(hidden)), torch.stack(keys), torch.stack(values)
@torch.inference_mode()
def trace_kv_graph(model):
graph = KVStepGraphV4(model).eval()
args = (torch.tensor([[2]], dtype=torch.int32), torch.tensor([0], dtype=torch.int32),
torch.zeros(CACHE_SHAPE), torch.zeros(CACHE_SHAPE))
traced = torch.jit.trace(graph, args)
# Different positions from the tracing example, including last cache slot.
ids = torch.randint(7, 8192, (1, 128), generator=torch.Generator().manual_seed(940))
cache = model.new_cache()
for position in (0, 7, 127):
if position > cache.length:
model.forward_cached(ids[:, cache.length:position], cache)
key, value = cache.key.clone(), cache.value.clone()
token = ids[:, position:position + 1].int()
reference = model.forward_cached(token, cache)
actual = traced(token, torch.tensor([position], dtype=torch.int32), key, value)
for a, b in zip(actual, (reference, cache.key, cache.value), strict=True):
torch.testing.assert_close(a, b, atol=3e-4, rtol=3e-4)
return traced
|