russell@unturf.com
8d4d60b421
ollama-0001: kvcache/causal.go buildMask() Except []int linear scan O(B*E).
For Gemma3 multi-image prompts, slices.Contains(opts.Except, i) is called
per batch token. With N images x 256 tokens/image, B and |except| both scale
as N*256 giving quadratic prefill cost. Fix: convert Except to map[int]struct{}
before the outer loop. 3.4x speedup at N=10 images. Unit test PASS.
dosbox-staging scanned for all 5 MOADs; all CLEAN. Linear scans operate on
fixed small palettes (16 CGA colors), hardcoded 3-item lists, or single-call
config paths. Mixer mutex consistent. No credential logging.
Updated defects/ollama/CLEAN.md to note the missed defect from prior scan.
100 lines
4 KiB
Java
100 lines
4 KiB
Java
import java.util.*;
|
|
|
|
/**
|
|
* Unit test for ollama-0001: kvcache.Causal.buildMask Except []int O(B*E) scan.
|
|
*
|
|
* Models the Ollama defect where buildMask() calls slices.Contains(opts.Except, i)
|
|
* inside a loop over all batch tokens (curBatchSize = B). Except is populated by
|
|
* Gemma3 for multimodal image token positions. For N images with 256 tokens each:
|
|
* |except| = N * 256, B = N * 256 + text_tokens ≈ N * 256
|
|
* Cost of Except check alone: O(B * E) = O((N*256)^2) — quadratic in image count.
|
|
*
|
|
* Fix: convert Except to a HashSet before the outer loop, giving O(1) per lookup
|
|
* and O(B + E) total cost for the Except component.
|
|
*/
|
|
public class OllamaKvcacheBuildmaskExceptTest {
|
|
|
|
// --- Defect: slices.Contains on []int Except, O(B*E) ---
|
|
// Isolates just the enabled[] computation (the Except scan), not the full mask fill.
|
|
static boolean[] computeEnabledLinear(int batchSize, int[] except) {
|
|
boolean[] enabled = new boolean[batchSize];
|
|
for (int i = 0; i < batchSize; i++) {
|
|
enabled[i] = !listContains(except, i);
|
|
}
|
|
return enabled;
|
|
}
|
|
|
|
static boolean listContains(int[] arr, int val) {
|
|
for (int v : arr) {
|
|
if (v == val) return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
// --- Fix: map[int]struct{} lookup, O(B+E) ---
|
|
static boolean[] computeEnabledHashSet(int batchSize, int[] except) {
|
|
Set<Integer> exceptSet = new HashSet<>(except.length * 2);
|
|
for (int idx : except) exceptSet.add(idx);
|
|
boolean[] enabled = new boolean[batchSize];
|
|
for (int i = 0; i < batchSize; i++) {
|
|
enabled[i] = !exceptSet.contains(i);
|
|
}
|
|
return enabled;
|
|
}
|
|
|
|
// Build an Except array simulating N images * tokensPerImage = 256 each
|
|
static int[] buildExcept(int numImages, int tokensPerImage) {
|
|
int[] except = new int[numImages * tokensPerImage];
|
|
for (int i = 0; i < except.length; i++) {
|
|
except[i] = i;
|
|
}
|
|
return except;
|
|
}
|
|
|
|
public static void main(String[] args) {
|
|
// Correctness: small case
|
|
int smallBatch = 16;
|
|
int[] smallExcept = {0, 1, 3, 5};
|
|
boolean[] enabledLinear = computeEnabledLinear(smallBatch, smallExcept);
|
|
boolean[] enabledHash = computeEnabledHashSet(smallBatch, smallExcept);
|
|
assert Arrays.equals(enabledLinear, enabledHash) : "Correctness failed";
|
|
System.out.println("Correctness OK: batch=" + smallBatch + " except=" + smallExcept.length);
|
|
|
|
// Performance: simulate 10-image Gemma3 prefill
|
|
// tokensPerImage=256, 10 images => batchSize=2560, except=2560
|
|
int numImages = 10;
|
|
int tokensPerImage = 256;
|
|
int batchSize = numImages * tokensPerImage;
|
|
int[] except = buildExcept(numImages, tokensPerImage);
|
|
|
|
// Correctness at full scale
|
|
boolean[] mL = computeEnabledLinear(batchSize, except);
|
|
boolean[] mH = computeEnabledHashSet(batchSize, except);
|
|
assert Arrays.equals(mL, mH) : "Correctness failed at full scale";
|
|
|
|
// Warmup
|
|
for (int w = 0; w < 3; w++) {
|
|
computeEnabledLinear(batchSize, except);
|
|
computeEnabledHashSet(batchSize, except);
|
|
}
|
|
|
|
int reps = 200;
|
|
long t0 = System.nanoTime();
|
|
for (int r = 0; r < reps; r++) computeEnabledLinear(batchSize, except);
|
|
long linearNs = System.nanoTime() - t0;
|
|
|
|
long t1 = System.nanoTime();
|
|
for (int r = 0; r < reps; r++) computeEnabledHashSet(batchSize, except);
|
|
long hashNs = System.nanoTime() - t1;
|
|
|
|
double ratio = (double) linearNs / hashNs;
|
|
System.out.printf("Gemma3 multi-image: %d images x %d tokens = B=%d |except|=%d%n",
|
|
numImages, tokensPerImage, batchSize, except.length);
|
|
System.out.printf(" Defect ([]int linear scan): %,d ns%n", linearNs);
|
|
System.out.printf(" Fix (HashSet O(1) lookup): %,d ns%n", hashNs);
|
|
System.out.printf(" speedup: %.1fx%n", ratio);
|
|
|
|
assert ratio > 3.0 : "Expected >3x speedup, got " + ratio;
|
|
System.out.println("PASS");
|
|
}
|
|
}
|