estimateModelMemory function

ModelMemoryEstimate? estimateModelMemory({ required String architecture, required Map<String, int> metadata, required int modelFileSizeBytes, int draftFileSizeBytes = 0, })

Builds a ModelMemoryEstimate from GGUF numeric metadata (keyed as <architecture>.<suffix>) and artifact file sizes.

Returns null when the metadata lacks the required hyperparameters (block count, embedding length, head count) — callers then fall back to loading with the spec's context size and skipping dynamic budgeting.

Implementation

ModelMemoryEstimate? estimateModelMemory({
  required String architecture,
  required Map<String, int> metadata,
  required int modelFileSizeBytes,
  int draftFileSizeBytes = 0,
}) {
  int? value(String suffix) => metadata['$architecture.$suffix'];

  final layerCount = value('block_count');
  final embeddingLength = value('embedding_length');
  final headCount = value('attention.head_count');
  if (layerCount == null || layerCount <= 0) return null;
  if (embeddingLength == null || embeddingLength <= 0) return null;
  if (headCount == null || headCount <= 0) return null;

  // Missing KV head count (e.g. stored per-layer as an array) falls back
  // to full multi-head attention — an over-estimate, which is the safe
  // direction for a budget.
  final kvHeadCount = value('attention.head_count_kv') ?? headCount;
  final headDimension = embeddingLength ~/ headCount;
  final keyLength = value('attention.key_length') ?? headDimension;
  final valueLength = value('attention.value_length') ?? headDimension;

  const kvBytesPerElement = 2;
  final kvBytesPerToken =
      layerCount * kvHeadCount * (keyLength + valueLength) * kvBytesPerElement;

  return ModelMemoryEstimate(
    weightsBytes: modelFileSizeBytes + draftFileSizeBytes,
    kvBytesPerToken: kvBytesPerToken,
    fixedOverheadBytes:
        _baseOverheadBytes + embeddingLength * _computeOverheadPerEmbedding,
    trainedContextTokens: value('context_length'),
  );
}