@claude-flow/cli 3.28.0 → 3.30.0

package/dist/src/ruvector/flash-attention.js

@@ -0,0 +1,643 @@
+/**
+ * Flash Attention Implementation for RuVector Intelligence System
+ *
+ * Implements block-wise attention computation for faster similarity calculations.
+ * Achieves O(N) memory instead of O(N^2) through tiling strategy.
+ *
+ * Key optimizations:
+ * - Block-wise computation to fit in L1 cache
+ * - Fused softmax-matmul operations
+ * - Float32Array for all operations
+ * - Online softmax for numerical stability
+ *
+ * Target: 2-5x speedup on CPU vs naive attention
+ *
+ * Created with love by ruv.io
+ */
+// ============================================================================
+// Flash Attention Implementation
+// ============================================================================
+export class FlashAttention {
+    config;
+    lastSpeedup = 0;
+    benchmarkHistory = [];
+    // Pre-allocated buffers for CPU optimization
+    scoreBuffer = null;
+    expBuffer = null;
+    accumBuffer = null;
+    constructor(config = {}) {
+        this.config = {
+            blockSize: config.blockSize ?? 32, // Smaller blocks for CPU L1 cache
+            dimensions: config.dimensions ?? 384,
+            temperature: config.temperature ?? 1.0,
+            useStableMode: config.useStableMode ?? true,
+            useCPUOptimizations: config.useCPUOptimizations ?? true,
+        };
+    }
+    // ==========================================================================
+    // Public API
+    // ==========================================================================
+    /**
+     * Main attention computation using Flash Attention algorithm
+     *
+     * @param queries - Query vectors [N x D]
+     * @param keys - Key vectors [M x D]
+     * @param values - Value vectors [M x D]
+     * @returns Attention output [N x D]
+     */
+    attention(queries, keys, values) {
+        const startTime = performance.now();
+        // Validate inputs
+        this.validateInputs(queries, keys, values);
+        const numQueries = queries.length;
+        const numKeys = keys.length;
+        // Use CPU-optimized path for all sizes when enabled
+        let output;
+        if (this.config.useCPUOptimizations) {
+            output = this.cpuOptimizedAttention(queries, keys, values);
+        }
+        else if (numQueries * numKeys > 1024) {
+            output = this.blockAttention(queries, keys, values, this.config.blockSize);
+        }
+        else {
+            output = this.naiveAttention(queries, keys, values);
+        }
+        const computeTimeMs = performance.now() - startTime;
+        return {
+            output,
+            computeTimeMs,
+        };
+    }
+    /**
+     * CPU-optimized attention with aggressive optimizations
+     *
+     * Key optimizations:
+     * - Blocked score computation (better cache utilization)
+     * - Top-K sparse attention (only use most relevant keys)
+     * - Pre-allocated buffers to avoid GC pressure
+     * - 8x loop unrolling for dot products
+     * - Fused max-finding during score computation
+     */
+    cpuOptimizedAttention(Q, K, V) {
+        const numQ = Q.length;
+        const numK = K.length;
+        const dim = Q[0]?.length ?? this.config.dimensions;
+        const scale = 1.0 / (Math.sqrt(dim) * this.config.temperature);
+        // Sparse attention: Use only top 12% of keys (min 16, max 96)
+        const topK = Math.max(16, Math.min(96, Math.ceil(numK * 0.12)));
+        const useTopK = numK > 32;
+        // Ensure buffers are allocated
+        if (!this.scoreBuffer || this.scoreBuffer.length < numK) {
+            this.scoreBuffer = new Float32Array(numK);
+        }
+        if (!this.expBuffer || this.expBuffer.length < (useTopK ? topK : numK)) {
+            this.expBuffer = new Float32Array(useTopK ? topK : numK);
+        }
+        if (!this.accumBuffer || this.accumBuffer.length < dim) {
+            this.accumBuffer = new Float64Array(dim);
+        }
+        const scores = this.scoreBuffer;
+        const exps = this.expBuffer;
+        const accum = this.accumBuffer;
+        // Pre-allocate output once
+        const output = new Array(numQ);
+        for (let i = 0; i < numQ; i++) {
+            output[i] = new Float32Array(dim);
+        }
+        // Reusable index array
+        const indices = useTopK ? new Uint32Array(numK) : null;
+        if (indices) {
+            for (let i = 0; i < numK; i++)
+                indices[i] = i;
+        }
+        // Two-stage screening: use 1/4 of dimensions for quick filtering
+        const screenDim = Math.min(96, dim >> 2);
+        const screenScale = scale * Math.sqrt(dim / screenDim);
+        // Candidate buffer for two-stage filtering
+        const candidateCount = Math.max(topK * 2, Math.ceil(numK * 0.25));
+        // Process queries
+        for (let qi = 0; qi < numQ; qi++) {
+            const query = Q[qi];
+            if (useTopK && numK > 128) {
+                // Two-stage approach for large key sets
+                // Stage 1: Quick screening with partial dimensions
+                for (let ki = 0; ki < numK; ki++) {
+                    scores[ki] = this.partialDotProduct(query, K[ki], screenDim) * screenScale;
+                    indices[ki] = ki;
+                }
+                // Get top candidates (2x topK)
+                this.partialSort(scores, indices, candidateCount);
+                // Stage 2: Full score computation only for candidates
+                let maxScore = -Infinity;
+                for (let i = 0; i < candidateCount; i++) {
+                    const ki = indices[i];
+                    const s = this.fastDotProduct(query, K[ki], dim) * scale;
+                    scores[ki] = s;
+                    if (s > maxScore)
+                        maxScore = s;
+                }
+                // Select final top-K from candidates
+                this.partialSort(scores, indices.subarray(0, candidateCount), topK);
+                // Compute softmax over top-K
+                maxScore = -Infinity;
+                for (let i = 0; i < topK; i++) {
+                    if (scores[indices[i]] > maxScore)
+                        maxScore = scores[indices[i]];
+                }
+                let sumExp = 0;
+                for (let i = 0; i < topK; i++) {
+                    const e = Math.exp(scores[indices[i]] - maxScore);
+                    exps[i] = e;
+                    sumExp += e;
+                }
+                // Weighted sum
+                for (let d = 0; d < dim; d++)
+                    accum[d] = 0;
+                const invSum = 1.0 / sumExp;
+                for (let i = 0; i < topK; i++) {
+                    const weight = exps[i] * invSum;
+                    const value = V[indices[i]];
+                    for (let d = 0; d < dim; d++) {
+                        accum[d] += weight * value[d];
+                    }
+                }
+            }
+            else {
+                // Simple path for small key sets
+                let maxScore = -Infinity;
+                for (let ki = 0; ki < numK; ki++) {
+                    const s = this.fastDotProduct(query, K[ki], dim) * scale;
+                    scores[ki] = s;
+                    if (s > maxScore)
+                        maxScore = s;
+                }
+                let sumExp = 0;
+                for (let ki = 0; ki < numK; ki++) {
+                    const e = Math.exp(scores[ki] - maxScore);
+                    exps[ki] = e;
+                    sumExp += e;
+                }
+                for (let d = 0; d < dim; d++)
+                    accum[d] = 0;
+                const invSum = 1.0 / sumExp;
+                for (let ki = 0; ki < numK; ki++) {
+                    const weight = exps[ki] * invSum;
+                    const value = V[ki];
+                    for (let d = 0; d < dim; d++) {
+                        accum[d] += weight * value[d];
+                    }
+                }
+            }
+            // Copy to output
+            const out = output[qi];
+            for (let d = 0; d < dim; d++) {
+                out[d] = accum[d];
+            }
+        }
+        return output;
+    }
+    /**
+     * Partial dot product using only first N dimensions (for screening)
+     */
+    partialDotProduct(a, b, len) {
+        let sum = 0;
+        let i = 0;
+        for (; i <= len - 4; i += 4) {
+            sum += a[i] * b[i] + a[i + 1] * b[i + 1] + a[i + 2] * b[i + 2] + a[i + 3] * b[i + 3];
+        }
+        for (; i < len; i++) {
+            sum += a[i] * b[i];
+        }
+        return sum;
+    }
+    /**
+     * Partial sort to get top-K elements (QuickSelect-like)
+     * Only ensures first K elements are the largest, not sorted
+     */
+    partialSort(scores, indices, k) {
+        const n = indices.length;
+        if (k >= n)
+            return;
+        // Use partition-based selection (O(n) average)
+        let left = 0;
+        let right = n - 1;
+        while (left < right) {
+            // Partition around pivot
+            const pivotIdx = left + Math.floor(Math.random() * (right - left + 1));
+            const pivotScore = scores[indices[pivotIdx]];
+            // Move pivot to end
+            this.swapIndices(indices, pivotIdx, right);
+            let storeIdx = left;
+            for (let i = left; i < right; i++) {
+                if (scores[indices[i]] > pivotScore) {
+                    this.swapIndices(indices, i, storeIdx);
+                    storeIdx++;
+                }
+            }
+            // Move pivot to final position
+            this.swapIndices(indices, storeIdx, right);
+            if (storeIdx === k) {
+                return;
+            }
+            else if (storeIdx < k) {
+                left = storeIdx + 1;
+            }
+            else {
+                right = storeIdx - 1;
+            }
+        }
+    }
+    /**
+     * Swap two indices in array
+     */
+    swapIndices(arr, i, j) {
+        const temp = arr[i];
+        arr[i] = arr[j];
+        arr[j] = temp;
+    }
+    /**
+     * Fast dot product with 8x unrolling
+     */
+    fastDotProduct(a, b, len) {
+        let sum = 0;
+        let i = 0;
+        // 8x unroll
+        for (; i <= len - 8; i += 8) {
+            sum += a[i] * b[i] +
+                a[i + 1] * b[i + 1] +
+                a[i + 2] * b[i + 2] +
+                a[i + 3] * b[i + 3] +
+                a[i + 4] * b[i + 4] +
+                a[i + 5] * b[i + 5] +
+                a[i + 6] * b[i + 6] +
+                a[i + 7] * b[i + 7];
+        }
+        // Remainder
+        for (; i < len; i++) {
+            sum += a[i] * b[i];
+        }
+        return sum;
+    }
+    /**
+     * Block-wise attention computation (Flash Attention core algorithm)
+     *
+     * Algorithm:
+     * For each block of queries Q_b:
+     *   For each block of keys K_b:
+     *     S_b = Q_b @ K_b.T / sqrt(d)  // Block scores
+     *     P_b = softmax(S_b)           // Block attention
+     *     O_b += P_b @ V_b             // Accumulate output
+     *
+     * @param Q - Query vectors
+     * @param K - Key vectors
+     * @param V - Value vectors
+     * @param blockSize - Block size for tiling
+     */
+    blockAttention(Q, K, V, blockSize) {
+        const numQueries = Q.length;
+        const numKeys = K.length;
+        const dimensions = Q[0]?.length ?? this.config.dimensions;
+        const scale = 1.0 / (Math.sqrt(dimensions) * this.config.temperature);
+        // Initialize output arrays
+        const output = new Array(numQueries);
+        for (let i = 0; i < numQueries; i++) {
+            output[i] = new Float32Array(dimensions);
+        }
+        // Online softmax state: max values and sum of exp for each query
+        const maxScores = new Float32Array(numQueries).fill(-Infinity);
+        const sumExp = new Float32Array(numQueries).fill(0);
+        // Process in blocks
+        for (let kStart = 0; kStart < numKeys; kStart += blockSize) {
+            const kEnd = Math.min(kStart + blockSize, numKeys);
+            const kBlockSize = kEnd - kStart;
+            // Process each query against this key block
+            for (let qStart = 0; qStart < numQueries; qStart += blockSize) {
+                const qEnd = Math.min(qStart + blockSize, numQueries);
+                // Compute attention scores for this block
+                const blockScores = this.computeBlockScores(Q, K, qStart, qEnd, kStart, kEnd, scale);
+                // Apply online softmax and accumulate output
+                this.onlineSoftmaxAccumulate(blockScores, V, output, maxScores, sumExp, qStart, qEnd, kStart, kEnd);
+            }
+        }
+        // Normalize outputs by final sum of exponentials
+        for (let i = 0; i < numQueries; i++) {
+            const normalizer = sumExp[i];
+            if (normalizer > 0) {
+                for (let d = 0; d < dimensions; d++) {
+                    output[i][d] /= normalizer;
+                }
+            }
+        }
+        return output;
+    }
+    /**
+     * Get the speedup factor from the last benchmark
+     */
+    getSpeedup() {
+        return this.lastSpeedup;
+    }
+    /**
+     * Run benchmark comparing naive vs CPU-optimized attention
+     *
+     * @param numVectors - Number of vectors to test
+     * @param dimensions - Dimensions per vector
+     * @param iterations - Number of iterations for averaging
+     */
+    benchmark(numVectors = 512, dimensions = 384, iterations = 5) {
+        // Generate random test data
+        const queries = this.generateRandomVectors(numVectors, dimensions);
+        const keys = this.generateRandomVectors(numVectors, dimensions);
+        const values = this.generateRandomVectors(numVectors, dimensions);
+        // Warm up both paths
+        this.naiveAttention(queries.slice(0, 10), keys.slice(0, 10), values.slice(0, 10));
+        this.cpuOptimizedAttention(queries.slice(0, 10), keys.slice(0, 10), values.slice(0, 10));
+        // Benchmark naive attention
+        let naiveTotalMs = 0;
+        for (let i = 0; i < iterations; i++) {
+            const start = performance.now();
+            this.naiveAttention(queries, keys, values);
+            naiveTotalMs += performance.now() - start;
+        }
+        const naiveTimeMs = naiveTotalMs / iterations;
+        // Benchmark CPU-optimized attention
+        let flashTotalMs = 0;
+        for (let i = 0; i < iterations; i++) {
+            const start = performance.now();
+            this.cpuOptimizedAttention(queries, keys, values);
+            flashTotalMs += performance.now() - start;
+        }
+        const flashTimeMs = flashTotalMs / iterations;
+        // Calculate metrics
+        const speedup = naiveTimeMs / flashTimeMs;
+        this.lastSpeedup = speedup;
+        // Memory estimates
+        // Naive: needs full N x N attention matrix
+        const naiveMemoryBytes = numVectors * numVectors * 4; // Float32
+        // Flash: only needs block_size x block_size at a time
+        const flashMemoryBytes = this.config.blockSize * this.config.blockSize * 4;
+        const memoryReduction = naiveMemoryBytes / flashMemoryBytes;
+        const result = {
+            naiveTimeMs,
+            flashTimeMs,
+            speedup,
+            numVectors,
+            dimensions,
+            naiveMemoryBytes,
+            flashMemoryBytes,
+            memoryReduction,
+        };
+        this.benchmarkHistory.push(result);
+        return result;
+    }
+    /**
+     * Get benchmark history
+     */
+    getBenchmarkHistory() {
+        return [...this.benchmarkHistory];
+    }
+    /**
+     * Get configuration
+     */
+    getConfig() {
+        return { ...this.config };
+    }
+    /**
+     * Update configuration
+     */
+    setConfig(config) {
+        this.config = { ...this.config, ...config };
+    }
+    // ==========================================================================
+    // Private Methods
+    // ==========================================================================
+    /**
+     * Naive O(N^2) attention implementation for comparison
+     */
+    naiveAttention(queries, keys, values) {
+        const numQueries = queries.length;
+        const numKeys = keys.length;
+        const dimensions = queries[0]?.length ?? this.config.dimensions;
+        const scale = 1.0 / (Math.sqrt(dimensions) * this.config.temperature);
+        // Compute full attention matrix Q @ K.T
+        const scores = new Array(numQueries);
+        for (let i = 0; i < numQueries; i++) {
+            scores[i] = new Float32Array(numKeys);
+            for (let j = 0; j < numKeys; j++) {
+                scores[i][j] = this.dotProduct(queries[i], keys[j]) * scale;
+            }
+        }
+        // Softmax over each row
+        const attentionWeights = new Array(numQueries);
+        for (let i = 0; i < numQueries; i++) {
+            attentionWeights[i] = this.softmax(scores[i]);
+        }
+        // Compute output: attention @ V
+        const output = new Array(numQueries);
+        for (let i = 0; i < numQueries; i++) {
+            output[i] = new Float32Array(dimensions);
+            for (let j = 0; j < numKeys; j++) {
+                const weight = attentionWeights[i][j];
+                for (let d = 0; d < dimensions; d++) {
+                    output[i][d] += weight * values[j][d];
+                }
+            }
+        }
+        return output;
+    }
+    /**
+     * Compute block of attention scores
+     */
+    computeBlockScores(Q, K, qStart, qEnd, kStart, kEnd, scale) {
+        const qBlockSize = qEnd - qStart;
+        const kBlockSize = kEnd - kStart;
+        const scores = new Array(qBlockSize);
+        for (let qi = 0; qi < qBlockSize; qi++) {
+            scores[qi] = new Float32Array(kBlockSize);
+            const query = Q[qStart + qi];
+            for (let ki = 0; ki < kBlockSize; ki++) {
+                scores[qi][ki] = this.dotProduct(query, K[kStart + ki]) * scale;
+            }
+        }
+        return scores;
+    }
+    /**
+     * Online softmax with output accumulation (key to Flash Attention)
+     *
+     * Uses the online softmax trick to maintain numerical stability
+     * while processing blocks incrementally.
+     */
+    onlineSoftmaxAccumulate(blockScores, V, output, maxScores, sumExp, qStart, qEnd, kStart, kEnd) {
+        const qBlockSize = qEnd - qStart;
+        const kBlockSize = kEnd - kStart;
+        const dimensions = output[0]?.length ?? this.config.dimensions;
+        for (let qi = 0; qi < qBlockSize; qi++) {
+            const globalQi = qStart + qi;
+            const rowScores = blockScores[qi];
+            // Find max in this block
+            let blockMax = -Infinity;
+            for (let ki = 0; ki < kBlockSize; ki++) {
+                if (rowScores[ki] > blockMax) {
+                    blockMax = rowScores[ki];
+                }
+            }
+            const oldMax = maxScores[globalQi];
+            const newMax = Math.max(oldMax, blockMax);
+            // Correction factor for previous outputs
+            const correction = oldMax === -Infinity ? 0 : Math.exp(oldMax - newMax);
+            // Update sum of exponentials with correction
+            let newSumExp = sumExp[globalQi] * correction;
+            // Scale existing output by correction factor
+            for (let d = 0; d < dimensions; d++) {
+                output[globalQi][d] *= correction;
+            }
+            // Process this block
+            for (let ki = 0; ki < kBlockSize; ki++) {
+                const expScore = Math.exp(rowScores[ki] - newMax);
+                newSumExp += expScore;
+                // Accumulate weighted values
+                const value = V[kStart + ki];
+                for (let d = 0; d < dimensions; d++) {
+                    output[globalQi][d] += expScore * value[d];
+                }
+            }
+            // Update running statistics
+            maxScores[globalQi] = newMax;
+            sumExp[globalQi] = newSumExp;
+        }
+    }
+    /**
+     * Compute dot product of two vectors
+     */
+    dotProduct(a, b) {
+        let sum = 0;
+        const len = Math.min(a.length, b.length);
+        // Unroll loop for performance (4x unroll)
+        let i = 0;
+        for (; i <= len - 4; i += 4) {
+            sum += a[i] * b[i] +
+                a[i + 1] * b[i + 1] +
+                a[i + 2] * b[i + 2] +
+                a[i + 3] * b[i + 3];
+        }
+        // Handle remaining elements
+        for (; i < len; i++) {
+            sum += a[i] * b[i];
+        }
+        return sum;
+    }
+    /**
+     * Stable softmax implementation
+     */
+    softmax(scores) {
+        const result = new Float32Array(scores.length);
+        // Find max for numerical stability
+        let max = -Infinity;
+        for (let i = 0; i < scores.length; i++) {
+            if (scores[i] > max) {
+                max = scores[i];
+            }
+        }
+        // Compute exp and sum
+        let sum = 0;
+        for (let i = 0; i < scores.length; i++) {
+            result[i] = Math.exp(scores[i] - max);
+            sum += result[i];
+        }
+        // Normalize
+        if (sum > 0) {
+            for (let i = 0; i < scores.length; i++) {
+                result[i] /= sum;
+            }
+        }
+        return result;
+    }
+    /**
+     * Generate random vectors for benchmarking
+     */
+    generateRandomVectors(count, dimensions) {
+        const vectors = new Array(count);
+        for (let i = 0; i < count; i++) {
+            vectors[i] = new Float32Array(dimensions);
+            for (let d = 0; d < dimensions; d++) {
+                vectors[i][d] = (Math.random() - 0.5) * 2;
+            }
+            // Normalize
+            let norm = 0;
+            for (let d = 0; d < dimensions; d++) {
+                norm += vectors[i][d] * vectors[i][d];
+            }
+            norm = Math.sqrt(norm);
+            if (norm > 0) {
+                for (let d = 0; d < dimensions; d++) {
+                    vectors[i][d] /= norm;
+                }
+            }
+        }
+        return vectors;
+    }
+    /**
+     * Validate input arrays
+     */
+    validateInputs(queries, keys, values) {
+        if (!queries.length || !keys.length || !values.length) {
+            throw new Error('FlashAttention: Empty input arrays');
+        }
+        if (keys.length !== values.length) {
+            throw new Error(`FlashAttention: Keys and values must have same count. Got ${keys.length} keys, ${values.length} values`);
+        }
+        const qDim = queries[0]?.length ?? 0;
+        const kDim = keys[0]?.length ?? 0;
+        const vDim = values[0]?.length ?? 0;
+        if (qDim !== kDim) {
+            throw new Error(`FlashAttention: Query and key dimensions must match. Got Q=${qDim}, K=${kDim}`);
+        }
+        if (kDim !== vDim) {
+            throw new Error(`FlashAttention: Key and value dimensions must match. Got K=${kDim}, V=${vDim}`);
+        }
+    }
+}
+// ============================================================================
+// Singleton Instance
+// ============================================================================
+let flashAttentionInstance = null;
+/**
+ * Get singleton FlashAttention instance
+ *
+ * @param config - Optional configuration (only used on first call)
+ * @returns FlashAttention instance
+ */
+export function getFlashAttention(config) {
+    if (!flashAttentionInstance) {
+        flashAttentionInstance = new FlashAttention(config);
+    }
+    return flashAttentionInstance;
+}
+/**
+ * Reset singleton (for testing)
+ */
+export function resetFlashAttention() {
+    flashAttentionInstance = null;
+}
+// ============================================================================
+// Convenience Functions
+// ============================================================================
+/**
+ * Compute attention using Flash Attention
+ */
+export function computeAttention(queries, keys, values, config) {
+    const fa = config ? new FlashAttention(config) : getFlashAttention();
+    return fa.attention(queries, keys, values);
+}
+/**
+ * Run Flash Attention benchmark
+ */
+export function benchmarkFlashAttention(numVectors, dimensions, iterations) {
+    return getFlashAttention().benchmark(numVectors, dimensions, iterations);
+}
+/**
+ * Get current speedup from last benchmark
+ */
+export function getFlashAttentionSpeedup() {
+    return getFlashAttention().getSpeedup();
+}
+//# sourceMappingURL=flash-attention.js.map