@sparkleideas/performance 3.0.0-alpha.7

package/benchmarks/attention/memory-efficiency.bench.ts

@@ -0,0 +1,569 @@
+/**
+ * Attention Memory Efficiency Benchmark
+ *
+ * Target: 50-75% memory reduction
+ *
+ * Measures memory efficiency of different attention implementations
+ * and optimization strategies.
+ */
+
+import { benchmark, BenchmarkRunner, formatTime, formatBytes } from '../../src/framework/benchmark.js';
+
+// ============================================================================
+// Memory Tracking
+// ============================================================================
+
+interface MemorySnapshot {
+  heapUsed: number;
+  heapTotal: number;
+  external: number;
+  arrayBuffers: number;
+  rss: number;
+}
+
+function takeMemorySnapshot(): MemorySnapshot {
+  const mem = process.memoryUsage();
+  return {
+    heapUsed: mem.heapUsed,
+    heapTotal: mem.heapTotal,
+    external: mem.external,
+    arrayBuffers: mem.arrayBuffers,
+    rss: mem.rss,
+  };
+}
+
+function calculateMemoryDelta(before: MemorySnapshot, after: MemorySnapshot): number {
+  return after.heapUsed - before.heapUsed;
+}
+
+// ============================================================================
+// Attention Implementations for Memory Testing
+// ============================================================================
+
+/**
+ * Standard attention - stores full attention matrix
+ */
+function standardAttention(
+  query: Float32Array,
+  key: Float32Array,
+  value: Float32Array,
+  seqLength: number,
+  headDim: number
+): { output: Float32Array; attentionMatrix: Float32Array } {
+  const scale = 1 / Math.sqrt(headDim);
+
+  // Full attention matrix - O(n^2) memory
+  const attentionMatrix = new Float32Array(seqLength * seqLength);
+
+  // Compute scores
+  for (let i = 0; i < seqLength; i++) {
+    for (let j = 0; j < seqLength; j++) {
+      let dot = 0;
+      for (let k = 0; k < headDim; k++) {
+        dot += query[i * headDim + k]! * key[j * headDim + k]!;
+      }
+      attentionMatrix[i * seqLength + j] = dot * scale;
+    }
+  }
+
+  // Softmax
+  for (let i = 0; i < seqLength; i++) {
+    let max = -Infinity;
+    for (let j = 0; j < seqLength; j++) {
+      max = Math.max(max, attentionMatrix[i * seqLength + j]!);
+    }
+
+    let sum = 0;
+    for (let j = 0; j < seqLength; j++) {
+      const exp = Math.exp(attentionMatrix[i * seqLength + j]! - max);
+      attentionMatrix[i * seqLength + j] = exp;
+      sum += exp;
+    }
+
+    for (let j = 0; j < seqLength; j++) {
+      attentionMatrix[i * seqLength + j]! /= sum;
+    }
+  }
+
+  // Output
+  const output = new Float32Array(seqLength * headDim);
+  for (let i = 0; i < seqLength; i++) {
+    for (let k = 0; k < headDim; k++) {
+      let sum = 0;
+      for (let j = 0; j < seqLength; j++) {
+        sum += attentionMatrix[i * seqLength + j]! * value[j * headDim + k]!;
+      }
+      output[i * headDim + k] = sum;
+    }
+  }
+
+  return { output, attentionMatrix };
+}
+
+/**
+ * Memory-efficient attention - no full matrix storage
+ */
+function memoryEfficientAttention(
+  query: Float32Array,
+  key: Float32Array,
+  value: Float32Array,
+  seqLength: number,
+  headDim: number
+): { output: Float32Array } {
+  const scale = 1 / Math.sqrt(headDim);
+  const output = new Float32Array(seqLength * headDim);
+
+  // Process row by row - O(n) memory for scores
+  const rowScores = new Float32Array(seqLength);
+
+  for (let i = 0; i < seqLength; i++) {
+    // Compute scores for this row
+    let max = -Infinity;
+    for (let j = 0; j < seqLength; j++) {
+      let dot = 0;
+      for (let k = 0; k < headDim; k++) {
+        dot += query[i * headDim + k]! * key[j * headDim + k]!;
+      }
+      rowScores[j] = dot * scale;
+      max = Math.max(max, rowScores[j]!);
+    }
+
+    // Softmax
+    let sum = 0;
+    for (let j = 0; j < seqLength; j++) {
+      rowScores[j] = Math.exp(rowScores[j]! - max);
+      sum += rowScores[j]!;
+    }
+    for (let j = 0; j < seqLength; j++) {
+      rowScores[j]! /= sum;
+    }
+
+    // Compute output for this row
+    for (let k = 0; k < headDim; k++) {
+      let val = 0;
+      for (let j = 0; j < seqLength; j++) {
+        val += rowScores[j]! * value[j * headDim + k]!;
+      }
+      output[i * headDim + k] = val;
+    }
+  }
+
+  return { output };
+}
+
+/**
+ * Chunked attention - process in blocks
+ */
+function chunkedAttention(
+  query: Float32Array,
+  key: Float32Array,
+  value: Float32Array,
+  seqLength: number,
+  headDim: number,
+  chunkSize: number = 64
+): { output: Float32Array } {
+  const scale = 1 / Math.sqrt(headDim);
+  const output = new Float32Array(seqLength * headDim);
+  const numChunks = Math.ceil(seqLength / chunkSize);
+
+  // Chunk buffer - O(chunkSize^2) memory
+  const chunkScores = new Float32Array(chunkSize * seqLength);
+  const rowMax = new Float32Array(chunkSize).fill(-Infinity);
+  const rowSum = new Float32Array(chunkSize).fill(0);
+
+  for (let ci = 0; ci < numChunks; ci++) {
+    const iStart = ci * chunkSize;
+    const iEnd = Math.min(iStart + chunkSize, seqLength);
+    const iSize = iEnd - iStart;
+
+    // Reset accumulators
+    rowMax.fill(-Infinity);
+    rowSum.fill(0);
+    output.fill(0, iStart * headDim, iEnd * headDim);
+
+    for (let cj = 0; cj < numChunks; cj++) {
+      const jStart = cj * chunkSize;
+      const jEnd = Math.min(jStart + chunkSize, seqLength);
+      const jSize = jEnd - jStart;
+
+      // Compute chunk scores
+      for (let i = 0; i < iSize; i++) {
+        for (let j = 0; j < jSize; j++) {
+          let dot = 0;
+          for (let k = 0; k < headDim; k++) {
+            dot += query[(iStart + i) * headDim + k]! * key[(jStart + j) * headDim + k]!;
+          }
+          chunkScores[i * seqLength + jStart + j] = dot * scale;
+        }
+      }
+
+      // Online softmax update
+      for (let i = 0; i < iSize; i++) {
+        const prevMax = rowMax[i]!;
+
+        // Find new max
+        for (let j = 0; j < jSize; j++) {
+          rowMax[i] = Math.max(rowMax[i]!, chunkScores[i * seqLength + jStart + j]!);
+        }
+
+        // Rescale previous
+        if (prevMax !== -Infinity && prevMax !== rowMax[i]) {
+          const rescale = Math.exp(prevMax - rowMax[i]!);
+          rowSum[i]! *= rescale;
+          for (let k = 0; k < headDim; k++) {
+            output[(iStart + i) * headDim + k]! *= rescale;
+          }
+        }
+
+        // Add new exponentials
+        for (let j = 0; j < jSize; j++) {
+          const exp = Math.exp(chunkScores[i * seqLength + jStart + j]! - rowMax[i]!);
+          chunkScores[i * seqLength + jStart + j] = exp;
+          rowSum[i]! += exp;
+        }
+
+        // Accumulate output
+        for (let k = 0; k < headDim; k++) {
+          for (let j = 0; j < jSize; j++) {
+            output[(iStart + i) * headDim + k]! +=
+              chunkScores[i * seqLength + jStart + j]! * value[(jStart + j) * headDim + k]!;
+          }
+        }
+      }
+    }
+
+    // Final normalization
+    for (let i = 0; i < iSize; i++) {
+      for (let k = 0; k < headDim; k++) {
+        output[(iStart + i) * headDim + k]! /= rowSum[i]!;
+      }
+    }
+  }
+
+  return { output };
+}
+
+// ============================================================================
+// Helper Functions
+// ============================================================================
+
+function generateRandomTensor(size: number): Float32Array {
+  const tensor = new Float32Array(size);
+  for (let i = 0; i < size; i++) {
+    tensor[i] = Math.random() * 2 - 1;
+  }
+  return tensor;
+}
+
+// ============================================================================
+// Benchmark Suite
+// ============================================================================
+
+export async function runMemoryEfficiencyBenchmarks(): Promise<void> {
+  const runner = new BenchmarkRunner('Attention Memory Efficiency');
+
+  console.log('\n--- Attention Memory Efficiency Benchmarks ---\n');
+
+  // Test configurations
+  const seqLengths = [128, 256, 512, 1024];
+  const headDim = 64;
+
+  // Memory scaling comparison
+  console.log('--- Memory Scaling by Sequence Length ---\n');
+
+  const memoryResults: Array<{
+    seqLength: number;
+    standard: number;
+    efficient: number;
+    chunked: number;
+    reduction: number;
+  }> = [];
+
+  for (const seqLength of seqLengths) {
+    console.log(`Sequence Length: ${seqLength}`);
+
+    const size = seqLength * headDim;
+    const query = generateRandomTensor(size);
+    const key = generateRandomTensor(size);
+    const value = generateRandomTensor(size);
+
+    // Standard attention memory
+    if (typeof global.gc === 'function') global.gc();
+    const standardBefore = takeMemorySnapshot();
+    const standardResult = standardAttention(query, key, value, seqLength, headDim);
+    const standardAfter = takeMemorySnapshot();
+    const standardMem = calculateMemoryDelta(standardBefore, standardAfter);
+    void standardResult;
+
+    // Memory-efficient attention
+    if (typeof global.gc === 'function') global.gc();
+    const efficientBefore = takeMemorySnapshot();
+    const efficientResult = memoryEfficientAttention(query, key, value, seqLength, headDim);
+    const efficientAfter = takeMemorySnapshot();
+    const efficientMem = calculateMemoryDelta(efficientBefore, efficientAfter);
+    void efficientResult;
+
+    // Chunked attention
+    if (typeof global.gc === 'function') global.gc();
+    const chunkedBefore = takeMemorySnapshot();
+    const chunkedResult = chunkedAttention(query, key, value, seqLength, headDim);
+    const chunkedAfter = takeMemorySnapshot();
+    const chunkedMem = calculateMemoryDelta(chunkedBefore, chunkedAfter);
+    void chunkedResult;
+
+    const reduction = ((standardMem - efficientMem) / standardMem) * 100;
+
+    memoryResults.push({
+      seqLength,
+      standard: standardMem,
+      efficient: efficientMem,
+      chunked: chunkedMem,
+      reduction,
+    });
+
+    console.log(`  Standard:  ${formatBytes(standardMem)}`);
+    console.log(`  Efficient: ${formatBytes(efficientMem)}`);
+    console.log(`  Chunked:   ${formatBytes(chunkedMem)}`);
+    console.log(`  Reduction: ${reduction.toFixed(1)}%`);
+    console.log('');
+  }
+
+  // Theoretical memory comparison
+  console.log('--- Theoretical Memory Analysis ---\n');
+
+  for (const seqLength of seqLengths) {
+    const bytesPerFloat = 4;
+
+    // Standard: stores full n x n attention matrix
+    const standardTheory = seqLength * seqLength * bytesPerFloat;
+
+    // Efficient: stores only one row at a time
+    const efficientTheory = seqLength * bytesPerFloat;
+
+    // Chunked: stores chunk x n scores
+    const chunkSize = 64;
+    const chunkedTheory = chunkSize * seqLength * bytesPerFloat;
+
+    console.log(`Seq ${seqLength}:`);
+    console.log(`  Standard:  ${formatBytes(standardTheory)} (n^2)`);
+    console.log(`  Efficient: ${formatBytes(efficientTheory)} (n)`);
+    console.log(`  Chunked:   ${formatBytes(chunkedTheory)} (chunk * n)`);
+    console.log(`  Theoretical reduction: ${((1 - efficientTheory / standardTheory) * 100).toFixed(1)}%`);
+    console.log('');
+  }
+
+  // Performance vs Memory tradeoff
+  console.log('--- Performance vs Memory Tradeoff ---\n');
+
+  const tradeoffConfig = { seqLength: 512, headDim: 64 };
+  const size = tradeoffConfig.seqLength * tradeoffConfig.headDim;
+  const q = generateRandomTensor(size);
+  const k = generateRandomTensor(size);
+  const v = generateRandomTensor(size);
+
+  // Standard performance
+  const standardPerfResult = await runner.run(
+    'standard-attention-perf',
+    async () => {
+      standardAttention(q, k, v, tradeoffConfig.seqLength, tradeoffConfig.headDim);
+    },
+    { iterations: 20 }
+  );
+
+  console.log(`Standard Performance: ${formatTime(standardPerfResult.mean)}`);
+
+  // Efficient performance
+  const efficientPerfResult = await runner.run(
+    'efficient-attention-perf',
+    async () => {
+      memoryEfficientAttention(q, k, v, tradeoffConfig.seqLength, tradeoffConfig.headDim);
+    },
+    { iterations: 20 }
+  );
+
+  console.log(`Memory-Efficient Performance: ${formatTime(efficientPerfResult.mean)}`);
+
+  // Chunked performance with different chunk sizes
+  const chunkSizes = [32, 64, 128, 256];
+
+  for (const chunkSize of chunkSizes) {
+    const chunkedPerfResult = await runner.run(
+      `chunked-attention-chunk${chunkSize}`,
+      async () => {
+        chunkedAttention(q, k, v, tradeoffConfig.seqLength, tradeoffConfig.headDim, chunkSize);
+      },
+      { iterations: 20 }
+    );
+
+    console.log(`Chunked (size=${chunkSize}): ${formatTime(chunkedPerfResult.mean)}`);
+  }
+
+  // Multi-head memory analysis
+  console.log('\n--- Multi-Head Memory Analysis ---\n');
+
+  const numHeads = [4, 8, 16, 32];
+  const mhaSeqLength = 256;
+
+  for (const heads of numHeads) {
+    const mhaSize = mhaSeqLength * headDim;
+
+    // Standard MHA memory
+    const standardMHAMem = mhaSeqLength * mhaSeqLength * 4 * heads; // attention matrices
+    const qkvMem = mhaSize * 4 * 3 * heads; // QKV storage
+
+    // GQA memory (shared KV)
+    const gqaKVHeads = heads / 4;
+    const gqaMem = mhaSeqLength * mhaSeqLength * 4 * heads + // attention matrices (same)
+      mhaSize * 4 * heads + // Q storage
+      mhaSize * 4 * 2 * gqaKVHeads; // shared KV
+
+    // MQA memory (single KV)
+    const mqaMem = mhaSeqLength * mhaSeqLength * 4 * heads + // attention matrices
+      mhaSize * 4 * heads + // Q storage
+      mhaSize * 4 * 2; // single KV
+
+    console.log(`${heads} heads:`);
+    console.log(`  Standard MHA: ${formatBytes(standardMHAMem + qkvMem)}`);
+    console.log(`  GQA (${gqaKVHeads} KV): ${formatBytes(gqaMem)}`);
+    console.log(`  MQA (1 KV):   ${formatBytes(mqaMem)}`);
+    console.log(`  MQA reduction: ${(((standardMHAMem + qkvMem) - mqaMem) / (standardMHAMem + qkvMem) * 100).toFixed(1)}%`);
+    console.log('');
+  }
+
+  // Summary
+  console.log('--- Summary ---\n');
+
+  console.log('Memory Reduction Achieved:');
+  for (const result of memoryResults) {
+    const targetMet = result.reduction >= 50;
+    console.log(
+      `  Seq ${result.seqLength}: ${result.reduction.toFixed(1)}% ${targetMet ? '(TARGET MET)' : ''}`
+    );
+  }
+
+  console.log('\nPerformance Comparison (seq=512):');
+  console.log(`  Standard: ${formatTime(standardPerfResult.mean)}`);
+  console.log(`  Efficient: ${formatTime(efficientPerfResult.mean)}`);
+
+  // Print full results
+  runner.printResults();
+}
+
+// ============================================================================
+// Memory Efficiency Optimization Strategies
+// ============================================================================
+
+export const memoryOptimizations = {
+  /**
+   * Online softmax computation
+   */
+  onlineSoftmax: {
+    description: 'Compute softmax in streaming fashion without storing all values',
+    expectedImprovement: 'O(n) instead of O(n^2) for softmax',
+    implementation: `
+      class OnlineSoftmax {
+        private max = -Infinity;
+        private sum = 0;
+        private count = 0;
+
+        add(value: number): void {
+          if (value > this.max) {
+            this.sum *= Math.exp(this.max - value);
+            this.max = value;
+          }
+          this.sum += Math.exp(value - this.max);
+          this.count++;
+        }
+
+        normalize(value: number): number {
+          return Math.exp(value - this.max) / this.sum;
+        }
+      }
+    `,
+  },
+
+  /**
+   * Gradient checkpointing
+   */
+  gradientCheckpointing: {
+    description: 'Recompute attention during backward pass instead of storing',
+    expectedImprovement: 'O(1) memory for activations',
+    implementation: `
+      function checkpointedAttention(q, k, v) {
+        const output = computeAttention(q, k, v);
+
+        function backward(gradOutput) {
+          // Recompute attention weights during backward
+          const attnWeights = recomputeAttention(q, k);
+          return computeGradients(gradOutput, attnWeights, q, k, v);
+        }
+
+        return { output, backward };
+      }
+    `,
+  },
+
+  /**
+   * Sparse attention patterns
+   */
+  sparseAttention: {
+    description: 'Only compute attention for relevant positions',
+    expectedImprovement: 'O(n * k) instead of O(n^2) where k << n',
+    implementation: `
+      function sparseAttention(q, k, v, pattern: 'local' | 'strided' | 'block') {
+        const sparseMask = generateSparsePattern(q.length, pattern);
+        return computeAttentionWithMask(q, k, v, sparseMask);
+      }
+    `,
+  },
+
+  /**
+   * Quantization
+   */
+  quantization: {
+    description: 'Use lower precision for attention computation',
+    expectedImprovement: '2-4x memory reduction',
+    implementation: `
+      function quantizedAttention(q, k, v) {
+        // Quantize to int8
+        const qInt8 = quantizeToInt8(q);
+        const kInt8 = quantizeToInt8(k);
+
+        // Compute in int8
+        const scores = computeInt8Attention(qInt8, kInt8);
+
+        // Dequantize for output
+        return dequantizeAndApply(scores, v);
+      }
+    `,
+  },
+
+  /**
+   * Memory pooling
+   */
+  memoryPooling: {
+    description: 'Reuse memory buffers across forward passes',
+    expectedImprovement: 'Eliminates allocation overhead',
+    implementation: `
+      class AttentionMemoryPool {
+        private scoreBuffer: Float32Array;
+        private outputBuffer: Float32Array;
+
+        forward(q, k, v) {
+          // Reuse pre-allocated buffers
+          computeScores(q, k, this.scoreBuffer);
+          applySoftmax(this.scoreBuffer);
+          computeOutput(this.scoreBuffer, v, this.outputBuffer);
+          return this.outputBuffer;
+        }
+      }
+    `,
+  },
+};
+
+// Run if executed directly
+if (import.meta.url === `file://${process.argv[1]}`) {
+  runMemoryEfficiencyBenchmarks().catch(console.error);
+}
+
+export default runMemoryEfficiencyBenchmarks;