@simulatte/doppler 0.1.8 → 0.2.0

package/src/inference/kv-cache/base.js

@@ -1,1083 +0,0 @@
-
-
-import { getDevice } from '../../gpu/device.js';
-import { allowReadback } from '../../gpu/perf-guards.js';
-import { log } from '../../debug/index.js';
-import { readBuffer } from '../../memory/buffer-pool.js';
-import {
-  isContiguousLayer,
-  isPagedLayer,
-  f32ToF16Array,
-  f16ToF32Array,
-} from './types.js';
-
-// ============================================================================
-// KVCache Class
-// ============================================================================
-
-export class KVCache {
-  
-  constructor(config) {
-    this.numLayers = config.numLayers;
-    
-    this.numHeads = config.numHeads;
-    
-    this.headDim = config.headDim;
-    if (!Number.isFinite(config.maxSeqLen) || config.maxSeqLen <= 0) {
-      throw new Error('KVCache requires a valid maxSeqLen.');
-    }
-    if (config.useGPU == null) {
-      throw new Error('KVCache requires useGPU to be set.');
-    }
-    if (!config.layout) {
-      throw new Error('KVCache requires a layout.');
-    }
-    if (!Number.isFinite(config.pageSize) || config.pageSize <= 0) {
-      throw new Error('KVCache requires a valid pageSize.');
-    }
-    if (!config.kvDtype) {
-      throw new Error('KVCache requires kvDtype.');
-    }
-
-    this.maxSeqLen = config.maxSeqLen;
-    
-    this.useGPU = config.useGPU;
-    
-    this.layout = config.layout;
-    
-    this.pageSize = config.pageSize;
-    
-    this.kvDtype = config.kvDtype;
-    
-    this.bytesPerElem = this.kvDtype === 'f16' ? 2 : 4;
-
-    // Size of one KV pair per position
-    
-    this.kvSize = this.numHeads * this.headDim;
-
-    // Initialize layer caches
-    
-    this.layers = new Array(this.numLayers);
-    
-    this.currentSeqLen = 0;
-    
-    this.totalTokensSeen = 0;
-
-    // Memory usage tracking
-    
-    this.memoryUsage = 0;
-
-    // GPU context (set externally)
-    
-    this.gpuContext = null;
-
-    // Initialize storage
-    this._initializeStorage();
-  }
-
-  // ==========================================================================
-  // Storage Initialization
-  // ==========================================================================
-
-  
-  _initializeStorage() {
-    if (this.layout === 'paged') {
-      this._initializePagedStorage();
-    } else {
-      this._initializeContiguousStorage();
-    }
-  }
-
-  
-  _initializeContiguousStorage() {
-    const sizePerLayer = this.maxSeqLen * this.kvSize;
-    const bytesPerLayer = sizePerLayer * this.bytesPerElem * 2; // K + V
-
-    const device = this.useGPU ? getDevice() : null;
-
-    for (let l = 0; l < this.numLayers; l++) {
-      if (device && this.useGPU) {
-        // GPU-native storage
-        const keysGPU = device.createBuffer({
-          label: `kv_cache_keys_layer_${l}`,
-          size: sizePerLayer * this.bytesPerElem,
-          usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
-        });
-        const valuesGPU = device.createBuffer({
-          label: `kv_cache_values_layer_${l}`,
-          size: sizePerLayer * this.bytesPerElem,
-          usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
-        });
-
-        this.layers[l] = {
-          keysGPU,
-          valuesGPU,
-          // Use empty CPU arrays in GPU mode - saves ~2.7GB for 9B models
-          // CPU shadows are only allocated on-demand if migrateFromGPU() is called
-          keys: new Float32Array(0),
-          values: new Float32Array(0),
-          seqLen: 0
-        };
-      } else {
-        // CPU-only storage
-        this.layers[l] = {
-          keys: new Float32Array(sizePerLayer),
-          values: new Float32Array(sizePerLayer),
-          keysGPU: null,
-          valuesGPU: null,
-          seqLen: 0
-        };
-      }
-      this.memoryUsage += bytesPerLayer;
-    }
-  }
-
-  
-  _initializePagedStorage() {
-    const numPages = Math.ceil(this.maxSeqLen / this.pageSize);
-    const device = this.useGPU ? getDevice() : null;
-    const sizePerLayer = this.maxSeqLen * this.kvSize;
-    const bytesPerLayer = sizePerLayer * this.bytesPerElem;
-    const pageTableBytes = numPages * 4;
-
-    for (let l = 0; l < this.numLayers; l++) {
-      if (device && this.useGPU) {
-        const keysGPU = device.createBuffer({
-          label: `kv_cache_keys_paged_layer_${l}`,
-          size: bytesPerLayer,
-          usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
-        });
-        const valuesGPU = device.createBuffer({
-          label: `kv_cache_values_paged_layer_${l}`,
-          size: bytesPerLayer,
-          usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
-        });
-        const pageTable = new Uint32Array(numPages);
-        for (let i = 0; i < numPages; i++) {
-          pageTable[i] = i;
-        }
-        const pageTableGPU = device.createBuffer({
-          label: `kv_cache_page_table_layer_${l}`,
-          size: pageTableBytes,
-          usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
-        });
-        device.queue.writeBuffer(pageTableGPU, 0, pageTable);
-
-        this.layers[l] = {
-          keyPages: new Array(numPages).fill(null),
-          valuePages: new Array(numPages).fill(null),
-          allocatedPages: 0,
-          seqLen: 0,
-          keysGPU,
-          valuesGPU,
-          pageTable,
-          pageTableGPU,
-        };
-        this.memoryUsage += (bytesPerLayer * 2) + pageTableBytes;
-      } else {
-        this.layers[l] = {
-          keyPages: new Array(numPages).fill(null),
-          valuePages: new Array(numPages).fill(null),
-          allocatedPages: 0,
-          seqLen: 0,
-          keysGPU: null,
-          valuesGPU: null,
-          pageTable: null,
-          pageTableGPU: null,
-        };
-      }
-    }
-  }
-
-  // ==========================================================================
-  // Paged Storage Helpers
-  // ==========================================================================
-
-  
-  _allocatePage() {
-    const pageElements = this.pageSize * this.kvSize;
-    const page = new Float32Array(pageElements);
-    this.memoryUsage += pageElements * 4;
-    return page;
-  }
-
-  
-  _getPageLocation(pos) {
-    const pageIdx = Math.floor(pos / this.pageSize);
-    const offset = (pos % this.pageSize) * this.kvSize;
-    return { pageIdx, offset };
-  }
-
-  
-  _ensurePagesAllocated(layerIdx, pos) {
-    if (this.layout !== 'paged') return;
-
-    const layer =  (this.layers[layerIdx]);
-    const neededPage = Math.floor(pos / this.pageSize);
-
-    for (let p = layer.allocatedPages; p <= neededPage; p++) {
-      if (!layer.keyPages[p]) {
-        layer.keyPages[p] = this._allocatePage();
-        layer.valuePages[p] = this._allocatePage();
-        layer.allocatedPages = p + 1;
-      }
-    }
-  }
-
-  // ==========================================================================
-  // Update Methods
-  // ==========================================================================
-
-  
-  _assertLayerIndex(layerIdx) {
-    if (!Number.isInteger(layerIdx) || layerIdx < 0 || layerIdx >= this.numLayers) {
-      throw new Error(`KVCache layer index out of range: ${layerIdx}`);
-    }
-  }
-
-  
-  _assertStartPos(startPos) {
-    if (!Number.isInteger(startPos) || startPos < 0) {
-      throw new Error('KVCache startPos must be a non-negative integer.');
-    }
-  }
-
-  
-  _resolveTokenCount(keys, values) {
-    const bytesPerToken = this.kvSize * this.bytesPerElem;
-    const keysIsGPU = keys instanceof GPUBuffer;
-    const valuesIsGPU = values instanceof GPUBuffer;
-
-    if (keysIsGPU || valuesIsGPU) {
-      if (!keysIsGPU || !valuesIsGPU) {
-        throw new Error('KVCache update requires both keys and values as GPU buffers.');
-      }
-      if (!Number.isFinite(keys.size) || !Number.isFinite(values.size)) {
-        throw new Error('KVCache update requires GPU buffer sizes.');
-      }
-      if (keys.size !== values.size) {
-        throw new Error('KVCache update requires matching GPU buffer sizes.');
-      }
-      if (keys.size % bytesPerToken !== 0) {
-        throw new Error('KVCache update GPU buffer size must align to kvSize.');
-      }
-      return { numNewTokens: keys.size / bytesPerToken, isGPU: true };
-    }
-
-    if (!Number.isFinite(keys?.length) || !Number.isFinite(values?.length)) {
-      throw new Error('KVCache update requires typed array inputs.');
-    }
-    if (keys.length !== values.length) {
-      throw new Error('KVCache update requires matching keys and values lengths.');
-    }
-    if (keys.length % this.kvSize !== 0) {
-      throw new Error('KVCache update array length must align to kvSize.');
-    }
-    return { numNewTokens: keys.length / this.kvSize, isGPU: false };
-  }
-
-  
-  update(
-    layerIdx,
-    keys,
-    values,
-    startPos = this.currentSeqLen
-  ) {
-    this._assertLayerIndex(layerIdx);
-    this._assertStartPos(startPos);
-
-    const { numNewTokens } = this._resolveTokenCount(keys, values);
-    if (!Number.isInteger(numNewTokens) || numNewTokens < 0) {
-      throw new Error('KVCache update requires a non-negative integer token count.');
-    }
-    if (numNewTokens === 0) {
-      return;
-    }
-
-    if (startPos + numNewTokens > this.maxSeqLen) {
-      throw new Error(
-        `Cache overflow: ${startPos + numNewTokens} > ${this.maxSeqLen}`
-      );
-    }
-
-    const layer = this.layers[layerIdx];
-
-    if (this.layout === 'paged') {
-      if (keys instanceof GPUBuffer || values instanceof GPUBuffer) {
-        throw new Error('Paged layout does not support GPU buffer inputs');
-      }
-      this._updatePaged(layerIdx,  (layer), keys, values, startPos, numNewTokens);
-    } else {
-      this._updateContiguous( (layer), keys, values, startPos, numNewTokens);
-    }
-
-    layer.seqLen = Math.max(layer.seqLen, startPos + numNewTokens);
-    this.totalTokensSeen = Math.max(this.totalTokensSeen, startPos + numNewTokens);
-
-    // Update global sequence length if this is the last layer
-    if (layerIdx === this.numLayers - 1) {
-      this.currentSeqLen = Math.max(this.currentSeqLen, startPos + numNewTokens);
-    }
-  }
-
-  
-  updateFromGPU(
-    layerIdx,
-    keysBuffer,
-    valuesBuffer,
-    startPos,
-    numTokens
-  ) {
-    this._assertLayerIndex(layerIdx);
-    this._assertStartPos(startPos);
-    if (!Number.isInteger(numTokens) || numTokens < 0) {
-      throw new Error('KVCache updateFromGPU requires a non-negative integer token count.');
-    }
-    if (numTokens === 0) {
-      return;
-    }
-
-    const layer =  (this.layers[layerIdx]);
-    const device = getDevice();
-
-    if (!device || !layer.keysGPU) {
-      throw new Error('GPU cache not initialized');
-    }
-
-    if (startPos + numTokens > this.maxSeqLen) {
-      throw new Error(
-        `Cache overflow: ${startPos + numTokens} > ${this.maxSeqLen}`
-      );
-    }
-
-    const byteOffset = startPos * this.kvSize * this.bytesPerElem;
-    const byteSize = numTokens * this.kvSize * this.bytesPerElem;
-    if (byteSize > keysBuffer.size || byteSize > valuesBuffer.size) {
-      throw new Error('KVCache updateFromGPU buffer size is smaller than requested write.');
-    }
-
-    // Copy directly from source buffers to cache buffers
-    const encoder = device.createCommandEncoder({ label: 'kv_cache_update' });
-    encoder.copyBufferToBuffer(keysBuffer, 0, layer.keysGPU, byteOffset, byteSize);
-    encoder.copyBufferToBuffer(valuesBuffer, 0, layer.valuesGPU, byteOffset, byteSize);
-    device.queue.submit([encoder.finish()]);
-
-    if (this.layout === 'paged') {
-      const neededPages = Math.ceil((startPos + numTokens) / this.pageSize);
-      if (Number.isFinite(layer.allocatedPages)) {
-        layer.allocatedPages = Math.max(layer.allocatedPages, neededPages);
-      } else {
-        layer.allocatedPages = neededPages;
-      }
-    }
-
-    layer.seqLen = Math.max(layer.seqLen, startPos + numTokens);
-    this.totalTokensSeen = Math.max(this.totalTokensSeen, startPos + numTokens);
-
-    if (layerIdx === this.numLayers - 1) {
-      this.currentSeqLen = Math.max(this.currentSeqLen, startPos + numTokens);
-    }
-  }
-
-  
-  recordUpdateFromGPU(
-    recorder,
-    layerIdx,
-    keysBuffer,
-    valuesBuffer,
-    startPos,
-    numTokens
-  ) {
-    this._assertLayerIndex(layerIdx);
-    this._assertStartPos(startPos);
-    if (!Number.isInteger(numTokens) || numTokens < 0) {
-      throw new Error('KVCache recordUpdateFromGPU requires a non-negative integer token count.');
-    }
-    if (numTokens === 0) {
-      return;
-    }
-
-    const encoder = recorder.getEncoder();
-    const layer =  (this.layers[layerIdx]);
-
-    if (!layer.keysGPU) {
-      throw new Error('GPU cache not initialized');
-    }
-
-    if (startPos + numTokens > this.maxSeqLen) {
-      throw new Error(
-        `Cache overflow: ${startPos + numTokens} > ${this.maxSeqLen}`
-      );
-    }
-
-    const byteOffset = startPos * this.kvSize * this.bytesPerElem;
-    const byteSize = numTokens * this.kvSize * this.bytesPerElem;
-    if (byteSize > keysBuffer.size || byteSize > valuesBuffer.size) {
-      throw new Error('KVCache recordUpdateFromGPU buffer size is smaller than requested write.');
-    }
-
-    // Record copy operations to the provided encoder (no submit)
-    encoder.copyBufferToBuffer(keysBuffer, 0, layer.keysGPU, byteOffset, byteSize);
-    encoder.copyBufferToBuffer(valuesBuffer, 0, layer.valuesGPU, byteOffset, byteSize);
-
-    if (this.layout === 'paged') {
-      const neededPages = Math.ceil((startPos + numTokens) / this.pageSize);
-      if (Number.isFinite(layer.allocatedPages)) {
-        layer.allocatedPages = Math.max(layer.allocatedPages, neededPages);
-      } else {
-        layer.allocatedPages = neededPages;
-      }
-    }
-
-    // Update seqLen metadata (this happens immediately, copies happen when encoder is submitted)
-    layer.seqLen = Math.max(layer.seqLen, startPos + numTokens);
-    this.totalTokensSeen = Math.max(this.totalTokensSeen, startPos + numTokens);
-
-    if (layerIdx === this.numLayers - 1) {
-      this.currentSeqLen = Math.max(this.currentSeqLen, startPos + numTokens);
-    }
-  }
-
-  
-  _updateContiguous(
-    layer,
-    keys,
-    values,
-    startPos,
-    numNewTokens
-  ) {
-    const offset = startPos * this.kvSize;
-    const device = getDevice();
-
-    // Handle GPU buffer inputs
-    if (keys instanceof GPUBuffer) {
-      // For GPU inputs, copy to GPU cache directly
-      if (layer.keysGPU && device) {
-        const byteOffset = offset * this.bytesPerElem;
-        const byteSize = numNewTokens * this.kvSize * this.bytesPerElem;
-        const encoder = device.createCommandEncoder({ label: 'kv_update_gpu' });
-        encoder.copyBufferToBuffer(keys, 0, layer.keysGPU, byteOffset, byteSize);
-        encoder.copyBufferToBuffer( (values), 0, layer.valuesGPU, byteOffset, byteSize);
-        device.queue.submit([encoder.finish()]);
-      }
-      return;
-    }
-
-    // CPU path
-    layer.keys.set(keys, offset);
-    layer.values.set( (values), offset);
-
-    // Also update GPU if available
-    if (layer.keysGPU && device) {
-      const byteOffset = offset * this.bytesPerElem;
-      if (this.kvDtype === 'f16') {
-        const keysF16 = f32ToF16Array( (keys));
-        const valuesF16 = f32ToF16Array( (values));
-        device.queue.writeBuffer(layer.keysGPU, byteOffset, keysF16);
-        device.queue.writeBuffer(layer.valuesGPU, byteOffset, valuesF16);
-      } else {
-        device.queue.writeBuffer(layer.keysGPU, byteOffset,  (keys));
-        device.queue.writeBuffer(layer.valuesGPU, byteOffset,  (values));
-      }
-    }
-  }
-
-  
-  _updatePaged(
-    layerIdx,
-    layer,
-    keys,
-    values,
-    startPos,
-    numNewTokens
-  ) {
-    const device = getDevice();
-    if (layer.keysGPU && layer.valuesGPU && device) {
-      const byteOffset = startPos * this.kvSize * this.bytesPerElem;
-      if (this.kvDtype === 'f16') {
-        const keysF16 = f32ToF16Array( (keys));
-        const valuesF16 = f32ToF16Array( (values));
-        device.queue.writeBuffer(layer.keysGPU, byteOffset, keysF16);
-        device.queue.writeBuffer(layer.valuesGPU, byteOffset, valuesF16);
-      } else {
-        device.queue.writeBuffer(layer.keysGPU, byteOffset,  (keys));
-        device.queue.writeBuffer(layer.valuesGPU, byteOffset,  (values));
-      }
-    }
-
-    for (let t = 0; t < numNewTokens; t++) {
-      const pos = startPos + t;
-      this._ensurePagesAllocated(layerIdx, pos);
-
-      const { pageIdx, offset } = this._getPageLocation(pos);
-      const srcOffset = t * this.kvSize;
-
-      layer.keyPages[pageIdx].set(
-        keys.subarray(srcOffset, srcOffset + this.kvSize),
-        offset
-      );
-      layer.valuePages[pageIdx].set(
-        values.subarray(srcOffset, srcOffset + this.kvSize),
-        offset
-      );
-    }
-  }
-
-  // ==========================================================================
-  // Get Methods
-  // ==========================================================================
-
-  
-  get(layerIdx, startPos = 0, endPos) {
-    this._assertLayerIndex(layerIdx);
-    this._assertStartPos(startPos);
-    const layer = this.layers[layerIdx];
-    const actualEndPos = endPos ?? layer.seqLen;
-    if (!Number.isInteger(actualEndPos) || actualEndPos < startPos) {
-      throw new Error('KVCache get requires endPos >= startPos.');
-    }
-    if (actualEndPos > layer.seqLen) {
-      throw new Error('KVCache get range exceeds cached sequence length.');
-    }
-
-    if (this.layout === 'paged' && this.useGPU) {
-      throw new Error('Paged GPU cache does not support CPU readback via get().');
-    }
-
-    if (this.layout === 'paged') {
-      return this._getPaged( (layer), startPos, actualEndPos);
-    } else {
-      return this._getContiguous( (layer), startPos, actualEndPos);
-    }
-  }
-
-  
-  getKeyCache(layerIdx) {
-    const layer = this.layers[layerIdx];
-    if (isContiguousLayer(layer)) {
-      return layer.keysGPU || layer.keys;
-    }
-    if (this.layout === 'paged') {
-      return layer.keysGPU || null;
-    }
-    return null;
-  }
-
-  
-  getValueCache(layerIdx) {
-    const layer = this.layers[layerIdx];
-    if (isContiguousLayer(layer)) {
-      return layer.valuesGPU || layer.values;
-    }
-    if (this.layout === 'paged') {
-      return layer.valuesGPU || null;
-    }
-    return null;
-  }
-
-  
-  getGPUBuffers(layerIdx) {
-    const layer = this.layers[layerIdx];
-
-    if (this.layout === 'paged') {
-      if (!layer.keysGPU || !layer.valuesGPU || !layer.pageTableGPU) {
-        return null;
-      }
-      return {
-        keysGPU: layer.keysGPU,
-        valuesGPU: layer.valuesGPU,
-        seqLen: layer.seqLen,
-        pageTableGPU: layer.pageTableGPU,
-        pageSize: this.pageSize,
-      };
-    }
-
-    if (!isContiguousLayer(layer) || !layer.keysGPU || !layer.valuesGPU) {
-      return null;
-    }
-
-    return {
-      keysGPU: layer.keysGPU,
-      valuesGPU: layer.valuesGPU,
-      seqLen: layer.seqLen,
-    };
-  }
-
-  
-  hasGPUCache() {
-    const firstLayer = this.layers[0];
-    if (!this.useGPU) return false;
-    if (this.layout === 'paged') {
-      return firstLayer.keysGPU != null && firstLayer.valuesGPU != null;
-    }
-    return isContiguousLayer(firstLayer) && firstLayer.keysGPU != null;
-  }
-
-  
-  _getContiguous(
-    layer,
-    startPos,
-    endPos
-  ) {
-    const startOffset = startPos * this.kvSize;
-    const endOffset = endPos * this.kvSize;
-
-    return {
-      keys: layer.keys.subarray(startOffset, endOffset),
-      values: layer.values.subarray(startOffset, endOffset)
-    };
-  }
-
-  
-  _getPaged(
-    layer,
-    startPos,
-    endPos
-  ) {
-    const length = (endPos - startPos) * this.kvSize;
-    const keys = new Float32Array(length);
-    const values = new Float32Array(length);
-
-    let destOffset = 0;
-    for (let pos = startPos; pos < endPos; pos++) {
-      const { pageIdx, offset } = this._getPageLocation(pos);
-      const keyPage = layer.keyPages[pageIdx];
-      const valuePage = layer.valuePages[pageIdx];
-      if (!keyPage || !valuePage) {
-        destOffset += this.kvSize;
-        continue;
-      }
-
-      keys.set(keyPage.subarray(offset, offset + this.kvSize), destOffset);
-      values.set(valuePage.subarray(offset, offset + this.kvSize), destOffset);
-
-      destOffset += this.kvSize;
-    }
-
-    return { keys, values };
-  }
-
-  // ==========================================================================
-  // Cache Management
-  // ==========================================================================
-
-  
-  clear() {
-    this.currentSeqLen = 0;
-    this.totalTokensSeen = 0;
-
-    for (let l = 0; l < this.numLayers; l++) {
-      const layer = this.layers[l];
-      layer.seqLen = 0;
-
-      if (this.layout === 'paged') {
-        // Don't deallocate pages, just reset length
-        // Pages will be reused
-      } else {
-        // Zero out contiguous arrays
-        const contiguousLayer =  (layer);
-        contiguousLayer.keys.fill(0);
-        contiguousLayer.values.fill(0);
-      }
-    }
-  }
-
-  
-  clone() {
-    if (this.layout === 'paged') {
-      const cloned = new KVCache({
-        numLayers: this.numLayers,
-        numHeads: this.numHeads,
-        headDim: this.headDim,
-        maxSeqLen: this.maxSeqLen,
-        useGPU: this.useGPU,
-        layout: 'paged',
-        pageSize: this.pageSize,
-        kvDtype: this.kvDtype
-      });
-
-      cloned.currentSeqLen = this.currentSeqLen;
-      cloned.totalTokensSeen = this.totalTokensSeen;
-
-      const device = this.useGPU ? getDevice() : null;
-      if (this.useGPU && !device) {
-        throw new Error('GPU device not initialized');
-      }
-
-      for (let l = 0; l < this.numLayers; l++) {
-        const src = this.layers[l];
-        const dst = cloned.layers[l];
-
-        const inferredPages = Math.ceil(src.seqLen / this.pageSize);
-        const allocatedPages = Number.isFinite(src.allocatedPages)
-          ? Math.max(src.allocatedPages, inferredPages)
-          : inferredPages;
-        dst.allocatedPages = allocatedPages;
-        dst.seqLen = src.seqLen;
-
-        if (this.useGPU && src.keysGPU && src.valuesGPU && dst.keysGPU && dst.valuesGPU) {
-          const usedBytes = src.seqLen * this.kvSize * this.bytesPerElem;
-          if (usedBytes > 0) {
-            const encoder = device.createCommandEncoder({ label: `kv_cache_clone_paged_${l}` });
-            encoder.copyBufferToBuffer(src.keysGPU, 0, dst.keysGPU, 0, usedBytes);
-            encoder.copyBufferToBuffer(src.valuesGPU, 0, dst.valuesGPU, 0, usedBytes);
-            device.queue.submit([encoder.finish()]);
-          }
-        }
-
-        if (src.pageTable && dst.pageTable && dst.pageTableGPU && device) {
-          dst.pageTable.set(src.pageTable);
-          device.queue.writeBuffer(dst.pageTableGPU, 0, src.pageTable);
-        }
-
-        if (!this.useGPU) {
-          for (let p = 0; p < allocatedPages; p++) {
-            const keyPage = src.keyPages?.[p];
-            const valuePage = src.valuePages?.[p];
-            if (!keyPage || !valuePage) {
-              continue;
-            }
-            const nextKeyPage = cloned._allocatePage();
-            nextKeyPage.set(keyPage);
-            const nextValuePage = cloned._allocatePage();
-            nextValuePage.set(valuePage);
-            dst.keyPages[p] = nextKeyPage;
-            dst.valuePages[p] = nextValuePage;
-          }
-        }
-      }
-
-      return cloned;
-    }
-
-    const cloned = new KVCache({
-      numLayers: this.numLayers,
-      numHeads: this.numHeads,
-      headDim: this.headDim,
-      maxSeqLen: this.maxSeqLen,
-      useGPU: false, // Always clone to CPU
-      layout: 'contiguous', // Simplify for clone
-      pageSize: this.pageSize,
-      kvDtype: this.kvDtype
-    });
-
-    cloned.currentSeqLen = this.currentSeqLen;
-
-    for (let l = 0; l < this.numLayers; l++) {
-      const { keys, values } = this.get(l);
-      const clonedLayer =  (cloned.layers[l]);
-      clonedLayer.keys.set(keys);
-      clonedLayer.values.set(values);
-      clonedLayer.seqLen = this.layers[l].seqLen;
-    }
-
-    return cloned;
-  }
-
-  
-  truncate(length) {
-    if (length >= this.currentSeqLen) return;
-
-    this.currentSeqLen = length;
-    this.totalTokensSeen = Math.min(this.totalTokensSeen, length);
-    for (let l = 0; l < this.numLayers; l++) {
-      this.layers[l].seqLen = Math.min(this.layers[l].seqLen, length);
-    }
-  }
-
-  
-  getMemoryStats() {
-    const theoretical = this.numLayers * 2 * this.maxSeqLen * this.kvSize * this.bytesPerElem;
-    const actual = this.memoryUsage;
-    const used = this.numLayers * 2 * this.currentSeqLen * this.kvSize * this.bytesPerElem;
-
-    return {
-      theoretical: theoretical,
-      allocated: actual,
-      used: used,
-      efficiency: used / actual,
-      seqLen: this.currentSeqLen,
-      maxSeqLen: this.maxSeqLen,
-      layout: this.layout
-    };
-  }
-
-  // ==========================================================================
-  // GPU Migration
-  // ==========================================================================
-
-  
-  setGPUContext(gpuContext) {
-    this.gpuContext = gpuContext;
-
-    // Migrate existing cache to GPU if we have data
-    if (this.currentSeqLen > 0 && gpuContext?.device) {
-      this._migrateToGPU(gpuContext.device);
-    }
-  }
-
-  _destroyGpuBuffer(buffer) {
-    if (!buffer) return;
-    try {
-      buffer.destroy();
-    } catch {
-      // Ignore already-destroyed buffers during rollback.
-    }
-  }
-
-  _snapshotLayerGpuState(layer) {
-    return {
-      keysGPU: layer.keysGPU ?? null,
-      valuesGPU: layer.valuesGPU ?? null,
-      pageTableGPU: layer.pageTableGPU ?? null,
-    };
-  }
-
-  _rollbackMigratedLayers(snapshots) {
-    for (let l = 0; l < this.numLayers; l++) {
-      const layer = this.layers[l];
-      const snapshot = snapshots[l];
-      if (!snapshot) continue;
-
-      if (layer.keysGPU && layer.keysGPU !== snapshot.keysGPU) {
-        this._destroyGpuBuffer(layer.keysGPU);
-      }
-      if (layer.valuesGPU && layer.valuesGPU !== snapshot.valuesGPU) {
-        this._destroyGpuBuffer(layer.valuesGPU);
-      }
-      if (layer.pageTableGPU && layer.pageTableGPU !== snapshot.pageTableGPU) {
-        this._destroyGpuBuffer(layer.pageTableGPU);
-      }
-
-      layer.keysGPU = snapshot.keysGPU;
-      layer.valuesGPU = snapshot.valuesGPU;
-      if ('pageTableGPU' in layer) {
-        layer.pageTableGPU = snapshot.pageTableGPU;
-      }
-    }
-  }
-
-  
-  _migrateToGPU(device) {
-    const snapshots = this.layers.map((layer) => this._snapshotLayerGpuState(layer));
-
-    try {
-    if (this.layout === 'paged') {
-      log.info('KVCache', `Migrating ${this.currentSeqLen} positions to GPU (paged)...`);
-      const numPages = Math.ceil(this.maxSeqLen / this.pageSize);
-      const sizePerLayer = this.maxSeqLen * this.kvSize;
-      const bytesPerLayer = sizePerLayer * this.bytesPerElem;
-      const pageTableBytes = numPages * 4;
-
-      for (let l = 0; l < this.numLayers; l++) {
-        const layer =  (this.layers[l]);
-        let keysGPU = null;
-        let valuesGPU = null;
-        let pageTableGPU = null;
-        try {
-          keysGPU = device.createBuffer({
-            label: `kv_cache_keys_paged_layer_${l}`,
-            size: bytesPerLayer,
-            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
-          });
-          valuesGPU = device.createBuffer({
-            label: `kv_cache_values_paged_layer_${l}`,
-            size: bytesPerLayer,
-            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
-          });
-
-          if (!layer.pageTable) {
-            layer.pageTable = new Uint32Array(numPages);
-            for (let i = 0; i < numPages; i++) {
-              layer.pageTable[i] = i;
-            }
-          }
-          pageTableGPU = device.createBuffer({
-            label: `kv_cache_page_table_layer_${l}`,
-            size: pageTableBytes,
-            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
-          });
-          device.queue.writeBuffer(pageTableGPU, 0, layer.pageTable);
-
-          const allocatedPages = layer.allocatedPages ?? 0;
-          if (allocatedPages > 0) {
-            const pageElems = this.pageSize * this.kvSize;
-            const pageBytes = pageElems * this.bytesPerElem;
-            for (let p = 0; p < allocatedPages; p++) {
-              const keyPage = layer.keyPages?.[p];
-              const valuePage = layer.valuePages?.[p];
-              if (!keyPage || !valuePage) continue;
-              const byteOffset = p * pageBytes;
-              if (this.kvDtype === 'f16') {
-                const keysF16 = f32ToF16Array(keyPage);
-                const valuesF16 = f32ToF16Array(valuePage);
-                device.queue.writeBuffer(keysGPU, byteOffset, keysF16);
-                device.queue.writeBuffer(valuesGPU, byteOffset, valuesF16);
-              } else {
-                device.queue.writeBuffer(keysGPU, byteOffset, keyPage);
-                device.queue.writeBuffer(valuesGPU, byteOffset, valuePage);
-              }
-            }
-          }
-
-          this._destroyGpuBuffer(layer.keysGPU);
-          this._destroyGpuBuffer(layer.valuesGPU);
-          this._destroyGpuBuffer(layer.pageTableGPU);
-          layer.keysGPU = keysGPU;
-          layer.valuesGPU = valuesGPU;
-          layer.pageTableGPU = pageTableGPU;
-        } catch (error) {
-          this._destroyGpuBuffer(keysGPU);
-          this._destroyGpuBuffer(valuesGPU);
-          this._destroyGpuBuffer(pageTableGPU);
-          throw error;
-        }
-      }
-
-      this.useGPU = true;
-      log.info('KVCache', 'Paged migration complete');
-      return;
-    }
-
-    log.info('KVCache', `Migrating ${this.currentSeqLen} positions to GPU...`);
-    const sizePerLayer = this.maxSeqLen * this.kvSize;
-    const bytesPerLayer = sizePerLayer * this.bytesPerElem;
-
-    for (let l = 0; l < this.numLayers; l++) {
-      const layer =  (this.layers[l]);
-      let keysGPU = null;
-      let valuesGPU = null;
-      try {
-        keysGPU = device.createBuffer({
-          label: `kv_cache_keys_layer_${l}`,
-          size: bytesPerLayer,
-          usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
-        });
-        valuesGPU = device.createBuffer({
-          label: `kv_cache_values_layer_${l}`,
-          size: bytesPerLayer,
-          usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
-        });
-
-        // Upload existing CPU data to GPU
-        const usedElems = layer.seqLen * this.kvSize;
-        const usedSize = usedElems * this.bytesPerElem;
-        if (usedSize > 0) {
-          if (this.kvDtype === 'f16') {
-            const keysF16 = f32ToF16Array(layer.keys.subarray(0, usedElems));
-            const valuesF16 = f32ToF16Array(layer.values.subarray(0, usedElems));
-            device.queue.writeBuffer(keysGPU, 0, keysF16);
-            device.queue.writeBuffer(valuesGPU, 0, valuesF16);
-          } else {
-            device.queue.writeBuffer(
-              keysGPU,
-              0,
-              layer.keys.buffer,
-              layer.keys.byteOffset,
-              usedSize
-            );
-            device.queue.writeBuffer(
-              valuesGPU,
-              0,
-              layer.values.buffer,
-              layer.values.byteOffset,
-              usedSize
-            );
-          }
-        }
-
-        this._destroyGpuBuffer(layer.keysGPU);
-        this._destroyGpuBuffer(layer.valuesGPU);
-        layer.keysGPU = keysGPU;
-        layer.valuesGPU = valuesGPU;
-      } catch (error) {
-        this._destroyGpuBuffer(keysGPU);
-        this._destroyGpuBuffer(valuesGPU);
-        throw error;
-      }
-    }
-
-    this.useGPU = true;
-    log.info('KVCache', 'Migration complete');
-    } catch (error) {
-      this._rollbackMigratedLayers(snapshots);
-      throw error;
-    }
-  }
-
-  
-  async syncToCPU() {
-    if (!this.useGPU || this.layout === 'paged') return;
-    if (!allowReadback('kv-cache.syncToCPU')) return;
-
-    const device = getDevice();
-    if (!device) return;
-
-    const sizePerLayer = this.maxSeqLen * this.kvSize;
-
-    for (let l = 0; l < this.numLayers; l++) {
-      const layer =  (this.layers[l]);
-      if (!layer.keysGPU || !layer.valuesGPU) continue;
-
-      const usedSize = layer.seqLen * this.kvSize * this.bytesPerElem;
-      if (usedSize === 0) continue;
-
-      // Allocate CPU arrays on-demand if empty (lazy allocation from GPU mode)
-      if (layer.keys.length === 0) {
-        layer.keys = new Float32Array(sizePerLayer);
-        layer.values = new Float32Array(sizePerLayer);
-      }
-
-      const [keysBytes, valuesBytes] = await Promise.all([
-        readBuffer(layer.keysGPU, usedSize),
-        readBuffer(layer.valuesGPU, usedSize),
-      ]);
-
-      if (this.kvDtype === 'f16') {
-        const keysRaw = new Uint16Array(keysBytes);
-        const valuesRaw = new Uint16Array(valuesBytes);
-        const keysData = f16ToF32Array(keysRaw);
-        const valuesData = f16ToF32Array(valuesRaw);
-        layer.keys.set(keysData);
-        layer.values.set(valuesData);
-      } else {
-        const keysData = new Float32Array(keysBytes);
-        const valuesData = new Float32Array(valuesBytes);
-        layer.keys.set(keysData);
-        layer.values.set(valuesData);
-      }
-    }
-  }
-
-  
-  destroy() {
-    for (let l = 0; l < this.numLayers; l++) {
-      const layer = this.layers[l];
-      if (this.layout === 'paged') {
-        if (layer.keysGPU) {
-          layer.keysGPU.destroy();
-          layer.keysGPU = null;
-        }
-        if (layer.valuesGPU) {
-          layer.valuesGPU.destroy();
-          layer.valuesGPU = null;
-        }
-        if (layer.pageTableGPU) {
-          layer.pageTableGPU.destroy();
-          layer.pageTableGPU = null;
-        }
-        continue;
-      }
-      if (isContiguousLayer(layer)) {
-        if (layer.keysGPU) {
-          layer.keysGPU.destroy();
-          layer.keysGPU = null;
-        }
-        if (layer.valuesGPU) {
-          layer.valuesGPU.destroy();
-          layer.valuesGPU = null;
-        }
-      }
-    }
-  }
-}