| """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) |
| |
| 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 |
|
|
