fix: performance improvements

Author: giladgd

src/gguf/insights/GgufInsights.ts

@@ -457,15 +457,6 @@ export class GgufInsights {
         const nHeadKv: number | number[] = this._ggufFileInfo.architectureMetadata.attention?.head_count_kv ?? nHead;
         const nEmbdHeadV = this._ggufFileInfo.architectureMetadata.attention?.value_length ?? ((nHead == 0) ? 0 : nEmbd / nHead);
 
-        const nHeadKvArray = Array(layers).fill(nHead);
-        if (Array.isArray(nHeadKv)) {
-            for (let i = 0; i < layers; i++) {
-                if (nHeadKv[i] !== 0)
-                    nHeadKvArray[i] = nHeadKv[i];
-            }
-        } else
-            nHeadKvArray.fill(nHeadKv);
-
         const ssmDConv = this._ggufFileInfo.architectureMetadata.ssm?.conv_kernel ?? 0;
         const ssmDInner = this._ggufFileInfo.architectureMetadata.ssm?.inner_size ?? 0;
         const modelNEmbdKS = (this._ggufFileInfo.architectureMetadata.wkv?.head_size ?? 0) !== 0
@@ -479,9 +470,15 @@ export class GgufInsights {
 
         let totalElementsK = 0;
         let totalElementsV = 0;
-        for (let i = 0; i < nHeadKvArray.length; i++) {
-            const nEmbdKGqa = nEmbdHeadK * nHeadKvArray[i]!;
-            const nEmbdVGqa = nEmbdHeadV * nHeadKvArray[i]!;
+        for (let i = 0; i < layers; i++) {
+            const nHeadKvArrayItem: number = (typeof nHeadKv === "number")
+                ? nHeadKv
+                : nHeadKv[i] !== 0
+                    ? nHeadKv[i]
+                    : nHead;
+
+            const nEmbdKGqa = nEmbdHeadK * nHeadKvArrayItem;
+            const nEmbdVGqa = nEmbdHeadV * nHeadKvArrayItem;
 
             const totalNEmbdKGqa = nEmbdKGqa + modelNEmbdKS;
             const totalNEmbdVGqa = nEmbdVGqa + modelNEmbdVS;

src/gguf/insights/utils/resolveModelGpuLayersOption.ts

@@ -239,22 +239,78 @@ function findMaxPossibleContextSizeForVram({gpuLayers, ggufInsights, vram, isEmb
 }) {
     const maxContextSize = getDefaultModelContextSize({trainContextSize: ggufInsights.trainContextSize});
 
-    for (let contextSize = maxContextSize; contextSize >= minAllowedContextSizeInCalculations; contextSize--) {
-        const contextVram = ggufInsights.estimateContextResourceRequirements({
-            contextSize,
-            batchSize: getDefaultContextBatchSize({contextSize, sequences: 1}),
-            modelGpuLayers: gpuLayers,
-            sequences: 1,
-            isEmbeddingContext,
-            flashAttention
-        }).gpuVram;
-
-        if (contextVram <= vram)
-            return {
+    return findMaxValidValue({
+        maxValue: maxContextSize,
+        minValue: minAllowedContextSizeInCalculations,
+        minStep: 1,
+        test(contextSize) {
+            const contextVram = ggufInsights.estimateContextResourceRequirements({
                 contextSize,
-                vram: contextVram
-            };
+                batchSize: getDefaultContextBatchSize({contextSize, sequences: 1}),
+                modelGpuLayers: gpuLayers,
+                sequences: 1,
+                isEmbeddingContext,
+                flashAttention
+            }).gpuVram;
+
+            if (contextVram <= vram)
+                return {
+                    contextSize,
+                    vram: contextVram
+                };
+
+            return null;
+        }
+    });
+}
+
+function findMaxValidValue<T>({
+    maxValue,
+    minValue,
+    minStep = 1,
+    test
+}: {
+    maxValue: number,
+    minValue: number,
+    minStep?: number,
+    test(value: number): T | null
+}): T | null {
+    let step = -Math.max(minStep, Math.floor((maxValue - minValue) / 4));
+    let bestValue: null | {value: number, result: T} = null;
+
+    for (let value = maxValue; value >= minValue;) {
+        const result = test(value);
+        if (result != null) {
+            if (bestValue == null || value >= bestValue.value) {
+                bestValue = {value: value, result: result};
+
+                if (step === -minStep)
+                    break;
+                else if (step < 0)
+                    step = Math.max(minStep, Math.floor(-step / 2));
+            }
+        } else if (bestValue != null && value < bestValue.value) {
+            value = bestValue.value;
+            step = Math.max(minStep, Math.floor(Math.abs(step) / 2));
+            continue;
+        } else if (step > 0)
+            step = -Math.max(minStep, Math.floor(step / 2));
+
+        if (value === minValue && step === -minStep)
+            break;
+
+        value += step;
+        if (value < minValue) {
+            value = minValue;
+            step = Math.max(minStep, Math.floor(Math.abs(step) / 2));
+        } else if (value > maxValue) {
+            value = maxValue;
+            step = -Math.max(minStep, Math.floor(Math.abs(step) / 2));
+        }
     }
 
+    if (bestValue != null)
+        return bestValue.result;
+
     return null;
 }