import { kvBytesPerToken, DEFAULT_ARCHITECTURES } from "@agentium/core";

const llama70b = DEFAULT_ARCHITECTURES["llama-3.1-70b"];

kvBytesPerToken(llama70b, "bf16");  // 327,680 bytes (320 KB)
kvBytesPerToken(llama70b, "fp8");   // 163,840 bytes (160 KB)
kvBytesPerToken(llama70b, "int4");  //  81,920 bytes (80 KB)

import { kvCacheForContext } from "@agentium/core";

const result = kvCacheForContext(llama70b, 131_072, "bf16");
// { bytes: 42_949_672_960, gb: 40.0 }

const short = kvCacheForContext(llama70b, 4_096, "fp8");
// { bytes: 671_088_640, gb: 0.625 }

import { maxContextForMemory } from "@agentium/core";

maxContextForMemory(llama70b, 20, "bf16");  // ~65,536 tokens
maxContextForMemory(llama70b, 20, "fp8");   // ~131,072 tokens (2× with fp8)

import { weightMemory } from "@agentium/core";

weightMemory(llama70b, "bf16");  // 140 GB
weightMemory(llama70b, "int8");  //  70 GB
weightMemory(llama70b, "int4");  //  35 GB

import { kvCacheForContext, weightMemory, OVERHEAD_GB } from "@agentium/core";

const totalHbm = 80 * 2;  // 160 GB
const weights = weightMemory(llama70b, "bf16");  // 140 GB
const freeHbm = totalHbm - weights - OVERHEAD_GB;  // 15 GB
const kvPerSession = kvCacheForContext(llama70b, 4096, "fp8").gb;  // 0.625 GB
const sessions = Math.floor(freeHbm / kvPerSession);  // 24 sessions

const fullContext = kvCacheForContext(llama70b, 131_072, "bf16");
// 40 GB — this is why 128K context needs 4+ H100s