kv_shape_recompile_demo.py

gistfile1.txt

import torch
import torch.nn as nn
import torch.nn.functional as F
import time

torch._dynamo.config.verbose = True

class TinyAttention(nn.Module):
    def __init__(self, dim=64):
        super().__init__()
        self.q = nn.Linear(dim, dim)
        self.k = nn.Linear(dim, dim)
        self.v = nn.Linear(dim, dim)
        self.out = nn.Linear(dim, dim)

    def forward(self, x, kv_cache):
        # x: [B, 1, D]
        # kv_cache: [B, T, D]  <-- T changes every token
        q = self.q(x)
        k = self.k(kv_cache)
        v = self.v(kv_cache)

        scores = q @ k.transpose(-2, -1)
        scores = scores / (x.shape[-1] ** 0.5)
        attn = F.softmax(scores, dim=-1)

        y = attn @ v
        return self.out(y)


model = TinyAttention()
compiled_model = torch.compile(model, dynamic=False)

B = 1
D = 64
max_tokens = 32

for t in range(1, max_tokens + 1):
    x = torch.randn(B, 1, D)
    kv_cache = torch.randn(B, t, D)

    start = time.time()
    y = compiled_model(x, kv_cache)
    elapsed = time.time() - start

    print(f"step={t:02d}, kv_shape={tuple(kv_cache.shape)}, time={elapsed:.4f}s")