kiku-agent 1.2.0 → 1.3.0

package/src/lib/vram/estimator.js

@@ -0,0 +1,464 @@
+export const GIB = 1024 ** 3;
+
+
+export const QUANT_BPW = Object.freeze({
+  fp32: 32.0,
+  f32: 32.0,
+  fp16: 16.0,
+  f16: 16.0,
+  bf16: 16.0,
+  int8: 8.0,
+  q8_0: 8.0,
+  q8: 8.0,
+  q6_k: 6.56,
+  q5_k_m: 5.5,
+  q5_k_s: 5.2,
+  q4_k_m: 4.5,
+  mxfp4: 4.25,
+  q4_k_s: 4.2,
+  q4_0: 4.2,
+  q4_1: 4.4,
+  iq4_xs: 4.25,
+  q3_k_l: 3.6,
+  q3_k_m: 3.4,
+  q3_k_s: 3.2,
+  q2_k: 2.6,
+  iq2_xxs: 2.1,
+});
+
+const QUANT_ALIASES = Object.freeze({
+  fp32: 'fp32', f32: 'fp32', float32: 'fp32',
+  fp16: 'fp16', f16: 'fp16', bf16: 'bf16', float16: 'fp16',
+  int8: 'int8', q8: 'q8_0', q80: 'q8_0', q8_0: 'q8_0',
+  q6k: 'q6_k', q6_k: 'q6_k',
+  q5km: 'q5_k_m', q5_k_m: 'q5_k_m', q5ks: 'q5_k_s', q5_k_s: 'q5_k_s',
+  q4km: 'q4_k_m', q4_k_m: 'q4_k_m', mxfp4: 'mxfp4', q4ks: 'q4_k_s', q4_k_s: 'q4_k_s',
+  q40: 'q4_0', q4_0: 'q4_0', q41: 'q4_1', q4_1: 'q4_1',
+  iq4xs: 'iq4_xs', iq4_xs: 'iq4_xs',
+  q3kl: 'q3_k_l', q3_k_l: 'q3_k_l', q3km: 'q3_k_m', q3_k_m: 'q3_k_m', q3ks: 'q3_k_s', q3_k_s: 'q3_k_s',
+  q2k: 'q2_k', q2_k: 'q2_k',
+  iq2xxs: 'iq2_xxs', iq2_xxs: 'iq2_xxs',
+});
+
+
+export const ACTIVATION_RANGE_GB = Object.freeze([0.15, 0.35]);
+export const DEFAULT_ACTIVATION_GB = 0.25;
+
+export const OVERHEAD_RANGE_GB = Object.freeze([0.6, 0.8]);
+export const DEFAULT_OVERHEAD_GB = 0.7;
+
+export const VISION_TOKENS_RANGE = Object.freeze([576, 2048]);
+
+
+
+export const ALLOCATION_THRESHOLDS = Object.freeze({ green: 0.85, yellow: 1.00 });
+
+
+export const SWA_ARCH_PROFILES = Object.freeze({
+  gemma2: { windowTokens: 4096, patternLocal: 1, patternGlobal: 1 },
+  gemma3: { windowTokens: 1024, patternLocal: 5, patternGlobal: 1 },
+  gemma3n: { windowTokens: 512, patternLocal: 5, patternGlobal: 1 },
+  gemma4: { windowTokens: 1024, patternLocal: 5, patternGlobal: 1 },
+});
+
+export function knownSwaProfile(arch) {
+  const key = String(arch || '').toLowerCase();
+  if (!key) return null;
+  if (SWA_ARCH_PROFILES[key]) return { ...SWA_ARCH_PROFILES[key] };
+  const prefix = Object.keys(SWA_ARCH_PROFILES)
+    .filter((k) => key.startsWith(k))
+    .sort((a, b) => b.length - a.length)[0];
+  return prefix ? { ...SWA_ARCH_PROFILES[prefix] } : null;
+}
+
+export function normalizeSwaProfile(swa) {
+  if (!swa || typeof swa !== 'object') return null;
+  const windowTokens = Math.round(Number(swa.windowTokens));
+  if (!Number.isFinite(windowTokens) || windowTokens <= 0) return null;
+  const patternLocal = Math.max(0, Math.round(Number(swa.patternLocal ?? 5)) || 0);
+  const patternGlobal = Math.max(0, Math.round(Number(swa.patternGlobal ?? 1)) || 0);
+  if (patternLocal + patternGlobal <= 0) return null;
+  return { windowTokens, patternLocal, patternGlobal };
+}
+
+export function normalizeQuantName(quant) {
+  if (quant === null || quant === undefined) return null;
+  if (typeof quant === 'number' && Number.isFinite(quant)) return Number(quant);
+  const compact = String(quant).toLowerCase().replace(/[^a-z0-9]/g, '');
+  if (!compact) return null;
+  if (QUANT_ALIASES[compact]) return QUANT_ALIASES[compact];
+  if (QUANT_BPW[compact]) return compact;
+  return compact;
+}
+
+
+export function parseBpw(quant) {
+  if (quant === null || quant === undefined || quant === '') return null;
+  if (typeof quant === 'number' && Number.isFinite(quant)) return quant;
+  const normalized = normalizeQuantName(quant);
+  if (typeof normalized === 'number') return normalized;
+  if (normalized && QUANT_BPW[normalized] !== undefined) return QUANT_BPW[normalized];
+  return null;
+}
+
+export function bytesToGb(bytes) {
+  return Number(bytes) / GIB;
+}
+
+export function gbToBytes(gb) {
+  return Number(gb) * GIB;
+}
+
+export function formatGb(bytes, digits = 2) {
+  if (bytes === null || bytes === undefined || !Number.isFinite(Number(bytes))) return '—';
+  return `${(Number(bytes) / GIB).toFixed(digits)} GB`;
+}
+
+
+export function estimateWeightsBytes({ paramCount = null, paramsB = null, quant = 'q4_k_m', weightBytesOverride = null } = {}) {
+  const bpw = parseBpw(quant);
+  if (weightBytesOverride !== null && weightBytesOverride !== undefined && Number.isFinite(Number(weightBytesOverride)) && Number(weightBytesOverride) > 0) {
+    let count = paramCount;
+    if (count === null || count === undefined) {
+      count = paramsB !== null && paramsB !== undefined && Number.isFinite(Number(paramsB))
+        ? Math.round(Number(paramsB) * 1e9)
+        : null;
+    }
+    return { bytes: Math.round(Number(weightBytesOverride)), bpw, paramCount: count, source: 'override' };
+  }
+  let count = paramCount;
+  if (count === null || count === undefined) {
+    count = paramsB !== null && paramsB !== undefined && Number.isFinite(Number(paramsB))
+      ? Math.round(Number(paramsB) * 1e9)
+      : null;
+  }
+  if (!count || count <= 0 || bpw === null) return { bytes: null, bpw, paramCount: count };
+  return { bytes: Math.round(count * bpw / 8), bpw, paramCount: count };
+}
+
+
+export function computeKvBytesPerToken({ layers, kvHeads, headDim, cacheBits = 16 } = {}) {
+  if (!layers || !kvHeads || !headDim) return null;
+  return 2 * layers * kvHeads * headDim * (Number(cacheBits || 16) / 8);
+}
+
+export function estimateKvCacheBytes({
+  layers = null,
+  kvHeads = null,
+  headDim = null,
+  contextLength = null,
+  batchSize = 1,
+  cacheBits = 16,
+  swa = null,
+} = {}) {
+  const perLayerPerToken = (layers && kvHeads && headDim)
+    ? 2 * kvHeads * headDim * (Number(cacheBits || 16) / 8)
+    : null;
+  const bytesPerToken = perLayerPerToken !== null ? perLayerPerToken * layers : null;
+  const ctx = Number(contextLength);
+  const profile = normalizeSwaProfile(swa);
+  const kvModel = profile ? 'swa-hybrid' : 'full-attention';
+  if (!perLayerPerToken || !ctx || ctx <= 0) {
+    return {
+      bytes: null,
+      bytesPerToken,
+      contextLength: Number.isFinite(ctx) ? ctx : null,
+      kvModel,
+      swa: null,
+    };
+  }
+  const batch = Math.max(1, Number(batchSize) || 1);
+  if (profile) {
+    const pattern = profile.patternLocal + profile.patternGlobal;
+    const globalLayers = Math.max(1, Math.ceil((layers * profile.patternGlobal) / pattern));
+    const localLayers = Math.max(0, layers - globalLayers);
+    const window = Math.min(profile.windowTokens, ctx);
+    return {
+      bytes: Math.round(perLayerPerToken * (globalLayers * ctx + localLayers * window) * batch),
+      bytesPerToken,
+      contextLength: ctx,
+      kvModel,
+      swa: { ...profile, globalLayers, localLayers },
+    };
+  }
+  return {
+    bytes: Math.round(bytesPerToken * ctx * batch),
+    bytesPerToken,
+    contextLength: ctx,
+    kvModel,
+    swa: null,
+  };
+}
+
+
+export function estimateActivationBytes({
+  batchSize = 1,
+  hiddenSize = null,
+  queryHeads = null,
+  kvHeads = null,
+  activationBits = 16,
+} = {}) {
+  if (!hiddenSize) {
+    return { analyticBytes: null, reservedBytes: Math.round(DEFAULT_ACTIVATION_GB * GIB) };
+  }
+  const gqaFactor = queryHeads && kvHeads ? 5 * (queryHeads / kvHeads) : 5;
+  const batch = Math.max(1, Number(batchSize) || 1);
+  const analyticBytes = Math.round(
+    batch * hiddenSize * (34 + gqaFactor) * (Number(activationBits || 16) / 8)
+  );
+  return { analyticBytes, reservedBytes: Math.round(DEFAULT_ACTIVATION_GB * GIB) };
+}
+
+/**
+ * Full footprint model. Unknown components are returned as null and excluded
+ * from the total, with a warning explaining what is missing.
+ *
+ * @param {object} input
+ * @param {string}  [input.modelId]              Display / registry identifier.
+ * @param {number}  [input.paramCount|paramsB]   Total parameters (MoE: total loaded params).
+ * @param {string}  [input.quant='q4_k_m']       Weight quantization.
+ * @param {number}  [input.contextLength]        Runtime sequence limit (text tokens).
+ * @param {number}  [input.batchSize=1]
+ * @param {number}  [input.cacheBits=16]         KV cache precision (FP16=16, Q8_0=8).
+ * @param {number}  [input.layers|kvHeads|queryHeads|headDim|hiddenSize]
+ * @param {object}  [input.vision]               { images, tokensPerImage, encoderParamsB }
+ * @param {number}  [input.activationGb]         Override static activation reservation.
+ * @param {number}  [input.overheadGb]           Override driver/runtime overhead.
+ * @returns {object} Breakdown in bytes + binary GB, warnings included.
+ */
+export function estimateVramFootprint(input = {}) {
+  const warnings = [];
+  const quant = input.quant ? normalizeQuantName(input.quant) : 'q4_k_m';
+
+  const weights = estimateWeightsBytes({
+    paramCount: input.paramCount,
+    paramsB: input.paramsB,
+    quant,
+    weightBytesOverride: input.weightBytesOverride,
+  });
+
+  const vision = input.vision || null;
+ 
+  let visionBytes = 0;
+  if (vision && (vision.encoderParamsB || vision.encoderParams)) {
+    const vpB = Number(vision.encoderParamsB ?? Number(vision.encoderParams) / 1e9);
+    if (Number.isFinite(vpB) && vpB > 0) {
+      visionBytes = Math.round(vpB * 1e9 * (weights.bpw ?? 16) / 8);
+    }
+  }
+
+  const textCtx = Number(input.contextLength) || 0;
+  let tokensPerImage = vision && vision.tokensPerImage ? Number(vision.tokensPerImage) : 0;
+  if (tokensPerImage && (tokensPerImage < VISION_TOKENS_RANGE[0] || tokensPerImage > VISION_TOKENS_RANGE[1])) {
+    tokensPerImage = Math.min(VISION_TOKENS_RANGE[1], Math.max(VISION_TOKENS_RANGE[0], tokensPerImage));
+  }
+  const images = vision && vision.images ? Math.max(0, Math.round(Number(vision.images))) : 0;
+  const effectiveContextLength = textCtx + (tokensPerImage ? images * tokensPerImage : 0);
+
+  const kv = estimateKvCacheBytes({
+    layers: input.layers,
+    kvHeads: input.kvHeads,
+    headDim: input.headDim,
+    contextLength: effectiveContextLength || null,
+    batchSize: input.batchSize,
+    cacheBits: input.cacheBits,
+    swa: input.swa,
+  });
+
+  const activation = estimateActivationBytes({
+    batchSize: input.batchSize,
+    hiddenSize: input.hiddenSize,
+    queryHeads: input.queryHeads,
+    kvHeads: input.kvHeads,
+    activationBits: 16,
+  });
+  const activationBytes = input.activationGb !== undefined && input.activationGb !== null
+    ? Math.round(Number(input.activationGb) * GIB)
+    : activation.reservedBytes;
+  const overheadBytes = input.overheadGb !== undefined && input.overheadGb !== null
+    ? Math.round(Number(input.overheadGb) * GIB)
+    : Math.round(DEFAULT_OVERHEAD_GB * GIB);
+
+  if (weights.bytes === null) {
+    warnings.push('Parameter count unknown — weights component excluded from total.');
+  }
+  if (kv.bytes === null) {
+    warnings.push('Architecture geometry incomplete (layers/kv_heads/head_dim) — KV cache component excluded.');
+  }
+  if (!textCtx) {
+    warnings.push('Context length not provided — KV cache computed against 0 tokens.');
+  }
+  if (vision && images > 0 && !tokensPerImage) {
+    warnings.push('Vision input detected but tokens-per-image unknown — context correction skipped.');
+  }
+
+  const totalBytes = weights.bytes === null
+    ? null
+    : (weights.bytes ?? 0) + visionBytes + (kv.bytes ?? 0) + activationBytes + overheadBytes;
+
+  const gb = (b) => (b === null || b === undefined ? null : Number((b / GIB).toFixed(4)));
+
+  return {
+    modelId: input.modelId || null,
+    quant,
+    quantBpw: weights.bpw,
+    paramCount: weights.paramCount,
+    batchSize: Math.max(1, Number(input.batchSize) || 1),
+    cacheBits: Number(input.cacheBits) || 16,
+    layers: input.layers ?? null,
+    kvHeads: input.kvHeads ?? null,
+    queryHeads: input.queryHeads ?? null,
+    headDim: input.headDim ?? null,
+    hiddenSize: input.hiddenSize ?? null,
+    contextLength: textCtx || null,
+    effectiveContextLength: effectiveContextLength || null,
+    vision: vision && (images > 0 || visionBytes > 0)
+      ? { images, tokensPerImage: tokensPerImage || null, visionEncoderBytes: visionBytes || null }
+      : null,
+    weightsBytes: weights.bytes,
+    visionEncoderBytes: visionBytes || null,
+    kvCacheBytes: kv.bytes,
+    kvBytesPerToken: kv.bytesPerToken,
+    kvModel: kv.kvModel,
+    swa: kv.swa,
+    activationBytes,
+    activationAnalyticBytes: activation.analyticBytes,
+    activationSource: input.activationGb !== undefined && input.activationGb !== null ? 'override' : 'static-reservation',
+    overheadBytes,
+    totalBytes,
+    gb: {
+      weights: gb(weights.bytes),
+      visionEncoder: visionBytes ? gb(visionBytes) : null,
+      kvCache: gb(kv.bytes),
+      activation: gb(activationBytes),
+      overhead: gb(overheadBytes),
+      total: gb(totalBytes),
+      kvBytesPerToken: kv.bytesPerToken,
+    },
+    warnings,
+  };
+}
+
+
+export function classifyAllocation({ totalBytes, availableBytes } = {}) {
+  if (!totalBytes || totalBytes <= 0 || !availableBytes || availableBytes <= 0) {
+    return {
+      zone: 'UNKNOWN',
+      usagePct: null,
+      headroomBytes: null,
+      message: 'Available VRAM or model weight parameters could not be determined — feasibility check skipped.',
+      requiresForce: false,
+    };
+  }
+  const usagePct = Number(((totalBytes / availableBytes) * 100).toFixed(1));
+  const headroomBytes = availableBytes - totalBytes;
+  if (totalBytes <= availableBytes * ALLOCATION_THRESHOLDS.green) {
+    return {
+      zone: 'GREEN',
+      usagePct,
+      headroomBytes,
+      message: 'Feasible: execution permitted without memory pressure.',
+      requiresForce: false,
+    };
+  }
+  if (totalBytes <= availableBytes * ALLOCATION_THRESHOLDS.yellow) {
+    return {
+      zone: 'YELLOW',
+      usagePct,
+      headroomBytes,
+      message: 'Marginal: fits in physical VRAM but leaves limited headroom for peak prefill.',
+      requiresForce: false,
+    };
+  }
+  return {
+    zone: 'RED',
+    usagePct,
+    headroomBytes,
+    message: headroomBytes >= 0
+      ? 'Critical: allocation at the physical VRAM limit.'
+      : 'Out of memory: allocation exceeds physical VRAM.',
+    requiresForce: true,
+  };
+}
+
+
+export const QUANT_LADDER = Object.freeze(['fp16', 'q8_0', 'q6_k', 'q5_k_m', 'q4_k_m', 'mxfp4', 'q4_k_s', 'q4_0', 'q3_k_m', 'q2_k']);
+
+
+export function recommendAdjustments(estimate, availableBytes, { targetZone = 'GREEN' } = {}) {
+  if (!estimate || !estimate.totalBytes || !availableBytes || availableBytes <= 0) return [];
+  const threshold = targetZone === 'GREEN' ? ALLOCATION_THRESHOLDS.green : ALLOCATION_THRESHOLDS.yellow;
+  const suggestions = [];
+
+  const kvBytes = estimate.kvCacheBytes ?? 0;
+  const fixedBytes = estimate.totalBytes - kvBytes;
+  const ctx = estimate.effectiveContextLength || estimate.contextLength;
+  if (ctx && kvBytes > 0) {
+    const budget = availableBytes * threshold - fixedBytes;
+    if (budget > 0) {
+   
+      const kvFor = (c) => estimateKvCacheBytes({
+        layers: estimate.layers,
+        kvHeads: estimate.kvHeads,
+        headDim: estimate.headDim,
+        contextLength: c,
+        batchSize: estimate.batchSize,
+        cacheBits: estimate.cacheBits,
+        swa: estimate.swa,
+      }).bytes ?? 0;
+      const step = 1024;
+      let lo = step;
+      let hi = Math.floor(ctx / step) * step;
+      let best = 0;
+      while (lo <= hi) {
+        const mid = Math.floor((lo + hi) / 2 / step) * step;
+        if (kvFor(mid) <= budget) {
+          best = mid;
+          lo = mid + step;
+        } else {
+          hi = mid - step;
+        }
+      }
+      if (best >= step && best < ctx) {
+        const projected = Math.round(fixedBytes + kvFor(best));
+        suggestions.push({
+          kind: 'context',
+          description: `Set --ctx ${best.toLocaleString('en-US')} to reduce footprint to ${formatGb(projected)}.`,
+          projectedBytes: projected,
+        });
+      }
+    }
+  }
+
+  if (kvBytes > 0 && Number(estimate.cacheBits) === 16) {
+    const projected = estimate.totalBytes - Math.round(kvBytes / 2);
+    if (projected <= availableBytes * threshold) {
+      suggestions.push({
+        kind: 'kv-quant',
+        description: `Switch the KV cache to Q8_0 (--cache-type-k q8_0 --cache-type-v q8_0) for a ${formatGb(projected)} footprint.`,
+        projectedBytes: Math.round(projected),
+      });
+    }
+  }
+
+  const currentIndex = QUANT_LADDER.indexOf(estimate.quant);
+  if (estimate.weightsBytes && currentIndex !== -1) {
+    for (let i = currentIndex + 1; i < QUANT_LADDER.length; i++) {
+      const bpw = parseBpw(QUANT_LADDER[i]);
+      if (!bpw || !estimate.paramCount) continue;
+      const projectedWeights = Math.round(estimate.paramCount * bpw / 8);
+      const projected = projectedWeights + (estimate.kvCacheBytes ?? 0) +
+        estimate.activationBytes + estimate.overheadBytes;
+      if (projected <= availableBytes * threshold) {
+        suggestions.push({
+          kind: 'quantization',
+          description: `Switch to ${QUANT_LADDER[i]} (${bpw} bpw) for a ${formatGb(projected)} footprint.`,
+          projectedBytes: Math.round(projected),
+        });
+        break;
+      }
+    }
+  }
+
+  return suggestions;
+}