width = 64
bytes_per_value = 4

attention_values_per_token = 2 * width
recurrent_state_values = width * width

for tokens in [128, 4096]:
    attention_kb = tokens * attention_values_per_token * bytes_per_value // 1024
    recurrent_kb = recurrent_state_values * bytes_per_value // 1024

    print(f"{tokens} tokens")
    print("  attention KV:", attention_kb, "KB")
    print("  recurrent state:", recurrent_kb, "KB")

# Output:
"""
128 tokens
  attention KV: 64 KB
  recurrent state: 16 KB
4096 tokens
  attention KV: 2048 KB
  recurrent state: 16 KB
"""