claude-flow
Version:
Ruflo - Enterprise AI agent orchestration for Claude Code. Deploy 60+ specialized agents in coordinated swarms with self-learning, fault-tolerant consensus, vector memory, and MCP integration
595 lines (516 loc) • 18.4 kB
text/typescript
/**
* @claude-flow/performance - Flash Attention Integration
*
* Integrates @ruvector/attention Flash Attention capabilities into V3 performance module.
* Provides optimized attention mechanisms with 2.49x-7.47x speedup targets.
*
* Features:
* - Flash Attention for memory-efficient processing
* - Automatic runtime selection (NAPI/WASM/JS)
* - Performance benchmarking and metrics
* - Speedup tracking and validation
*/
import { createRequire } from 'node:module';
export interface AttentionConfig {
dim: number;
numHeads?: number;
blockSize?: number;
}
interface AttentionRuntimeModule {
FlashAttention: new (dim: number, blockSize?: number) => {
compute(query: Float32Array, keys: Float32Array[], values: Float32Array[]): Float32Array;
};
DotProductAttention: new (dim: number) => {
compute(query: Float32Array, keys: Float32Array[], values: Float32Array[]): Float32Array;
computeWithMask?(
query: Float32Array,
keys: Float32Array[],
values: Float32Array[],
mask: Float32Array,
): Float32Array;
};
MultiHeadAttention: new (config: AttentionConfig) => {
compute(input: Float32Array): Float32Array;
};
LinearAttention: new (config: AttentionConfig) => {
compute(input: Float32Array): Float32Array;
};
}
const require = createRequire(import.meta.url);
const attentionRuntime = require('@ruvector/attention') as AttentionRuntimeModule;
export class FlashAttention {
private readonly impl: InstanceType<AttentionRuntimeModule['FlashAttention']>;
constructor(dim: number, blockSize: number = 64) {
this.impl = new attentionRuntime.FlashAttention(dim, blockSize);
}
/** @deprecated Use compute() — alias retained for backward compatibility with pre-interop callers */
computeRaw(query: Float32Array, keys: Float32Array[], values: Float32Array[]): Float32Array {
return this.impl.compute(query, keys, values);
}
compute(query: Float32Array, keys: Float32Array[], values: Float32Array[]): Float32Array {
return this.impl.compute(query, keys, values);
}
}
export class DotProductAttention {
private readonly impl: InstanceType<AttentionRuntimeModule['DotProductAttention']>;
constructor(dim: number) {
this.impl = new attentionRuntime.DotProductAttention(dim);
}
/** @deprecated Use compute() — alias retained for backward compatibility with pre-interop callers */
computeRaw(query: Float32Array, keys: Float32Array[], values: Float32Array[]): Float32Array {
return this.impl.compute(query, keys, values);
}
compute(query: Float32Array, keys: Float32Array[], values: Float32Array[]): Float32Array {
return this.impl.compute(query, keys, values);
}
}
export class MultiHeadAttention {
private impl: any;
constructor(config: AttentionConfig) {
this.impl = new attentionRuntime.MultiHeadAttention(config);
}
compute(input: Float32Array): Float32Array {
return this.impl.compute(input);
}
/** @deprecated Use compute() — alias retained for backward compatibility with pre-interop callers */
computeRaw(input: Float32Array): Float32Array {
return this.impl.compute(input);
}
}
export class LinearAttention {
private impl: any;
constructor(config: AttentionConfig) {
this.impl = new attentionRuntime.LinearAttention(config);
}
compute(input: Float32Array): Float32Array {
return this.impl.compute(input);
}
/** @deprecated Use compute() — alias retained for backward compatibility with pre-interop callers */
computeRaw(input: Float32Array): Float32Array {
return this.impl.compute(input);
}
}
export type ArrayInput = Float32Array | number[];
export type AttentionBenchmarkResult = BenchmarkResult;
// ============================================================================
// Types
// ============================================================================
export interface AttentionInput {
query: Float32Array | number[];
keys: Float32Array[] | number[][];
values: Float32Array[] | number[][];
dim?: number;
blockSize?: number;
}
export interface AttentionOutput {
result: Float32Array;
runtime: 'napi' | 'wasm' | 'js';
executionTimeMs: number;
memoryUsageBytes?: number;
}
export interface BenchmarkResult {
flashAttention: {
averageTimeMs: number;
opsPerSecond: number;
memoryUsageBytes?: number;
};
baseline: {
averageTimeMs: number;
opsPerSecond: number;
memoryUsageBytes?: number;
};
speedup: number;
meetsTarget: boolean; // true if speedup >= 2.49x
timestamp: Date;
}
export interface PerformanceMetrics {
totalOperations: number;
averageSpeedup: number;
peakSpeedup: number;
averageExecutionTimeMs: number;
totalMemorySavedBytes: number;
successRate: number; // % of operations meeting target
// Memory tracking metrics
baselineMemoryBytes: number;
optimizedMemoryBytes: number;
memorySavedBytes: number;
memorySavedPercent: number;
peakMemoryBytes: number;
}
// ============================================================================
// Flash Attention Optimizer
// ============================================================================
export class FlashAttentionOptimizer {
private flashAttention: FlashAttention;
private baselineAttention: DotProductAttention;
private metrics: {
operations: number;
totalSpeedup: number;
peakSpeedup: number;
totalExecutionTime: number;
successfulOperations: number;
// Memory tracking
totalBaselineMemory: number;
totalOptimizedMemory: number;
peakMemory: number;
};
constructor(
private readonly dim: number = 512,
private readonly blockSize: number = 64
) {
this.flashAttention = new FlashAttention(dim, blockSize);
this.baselineAttention = new DotProductAttention(dim);
this.metrics = {
operations: 0,
totalSpeedup: 0,
peakSpeedup: 0,
totalExecutionTime: 0,
successfulOperations: 0,
totalBaselineMemory: 0,
totalOptimizedMemory: 0,
peakMemory: 0,
};
}
/**
* Optimize attention computation using Flash Attention
* @param input - Query, keys, and values for attention computation
* @returns Optimized attention output with performance metrics
*/
optimize(input: AttentionInput): AttentionOutput {
const startTime = performance.now();
const startMemory = this.getMemoryUsage();
// Convert inputs if needed
const query = this.ensureFloat32Array(input.query);
const keys = input.keys.map(k => this.ensureFloat32Array(k));
const values = input.values.map(v => this.ensureFloat32Array(v));
// Use synchronous Flash Attention with raw Float32Arrays
const result = this.flashAttention.computeRaw(query, keys, values);
const executionTimeMs = performance.now() - startTime;
const endMemory = this.getMemoryUsage();
const memoryUsageBytes = endMemory - startMemory;
// Update metrics
this.metrics.operations++;
this.metrics.totalExecutionTime += executionTimeMs;
return {
result,
runtime: this.detectRuntime(),
executionTimeMs,
memoryUsageBytes: memoryUsageBytes > 0 ? memoryUsageBytes : undefined,
};
}
/**
* Benchmark Flash Attention vs baseline attention
* @returns Comprehensive benchmark results with speedup metrics
*/
benchmark(): BenchmarkResult {
const dim = this.dim;
const numKeys = 100;
const iterations = 1000;
// Create test data
const query = new Float32Array(dim);
const keys = Array.from({ length: numKeys }, () => new Float32Array(dim));
const values = Array.from({ length: numKeys }, () => new Float32Array(dim));
// Fill with random data
for (let i = 0; i < dim; i++) {
query[i] = Math.random();
}
for (let i = 0; i < numKeys; i++) {
for (let j = 0; j < dim; j++) {
keys[i][j] = Math.random();
values[i][j] = Math.random();
}
}
// Force garbage collection if available for accurate memory measurement
this.forceGC();
// Measure baseline memory usage
const baselineMemoryBefore = this.getMemoryUsage();
let baselinePeakMemory = baselineMemoryBefore;
// Benchmark baseline (DotProduct) - run first to establish baseline memory
const baselineStart = performance.now();
for (let i = 0; i < iterations; i++) {
this.baselineAttention.computeRaw(query, keys, values);
// Sample memory periodically (every 100 iterations to reduce overhead)
if (i % 100 === 0) {
const currentMemory = this.getMemoryUsage();
if (currentMemory > baselinePeakMemory) {
baselinePeakMemory = currentMemory;
}
}
}
const baselineEnd = performance.now();
const baselineMemoryAfter = this.getMemoryUsage();
const baselineTimeMs = baselineEnd - baselineStart;
const baselineAvgMs = baselineTimeMs / iterations;
const baselineMemoryUsed = Math.max(0, baselinePeakMemory - baselineMemoryBefore);
// Force garbage collection before Flash Attention benchmark
this.forceGC();
// Measure Flash Attention memory usage
const flashMemoryBefore = this.getMemoryUsage();
let flashPeakMemory = flashMemoryBefore;
// Benchmark Flash Attention
const flashStart = performance.now();
for (let i = 0; i < iterations; i++) {
this.flashAttention.computeRaw(query, keys, values);
// Sample memory periodically
if (i % 100 === 0) {
const currentMemory = this.getMemoryUsage();
if (currentMemory > flashPeakMemory) {
flashPeakMemory = currentMemory;
}
}
}
const flashEnd = performance.now();
const flashTimeMs = flashEnd - flashStart;
const flashAvgMs = flashTimeMs / iterations;
const flashMemoryUsed = Math.max(0, flashPeakMemory - flashMemoryBefore);
const speedup = baselineAvgMs / flashAvgMs;
const meetsTarget = speedup >= 2.49; // Minimum target: 2.49x
// Update peak speedup
if (speedup > this.metrics.peakSpeedup) {
this.metrics.peakSpeedup = speedup;
}
// Update memory tracking metrics
this.metrics.totalBaselineMemory += baselineMemoryUsed;
this.metrics.totalOptimizedMemory += flashMemoryUsed;
if (flashPeakMemory > this.metrics.peakMemory) {
this.metrics.peakMemory = flashPeakMemory;
}
this.metrics.totalSpeedup += speedup;
if (meetsTarget) {
this.metrics.successfulOperations++;
}
return {
flashAttention: {
averageTimeMs: flashAvgMs,
opsPerSecond: 1000 / flashAvgMs,
memoryUsageBytes: flashMemoryUsed,
},
baseline: {
averageTimeMs: baselineAvgMs,
opsPerSecond: 1000 / baselineAvgMs,
memoryUsageBytes: baselineMemoryUsed,
},
speedup,
meetsTarget,
timestamp: new Date(),
};
}
/**
* Get current speedup factor from accumulated metrics
* @returns Average speedup factor across all operations
*/
getSpeedup(): number {
if (this.metrics.operations === 0) {
return 0;
}
return this.metrics.totalSpeedup / this.metrics.operations;
}
/**
* Get comprehensive performance metrics
* @returns Detailed performance statistics
*/
getMetrics(): PerformanceMetrics {
const avgSpeedup = this.getSpeedup();
// Calculate memory savings
const baselineMemory = this.metrics.totalBaselineMemory;
const optimizedMemory = this.metrics.totalOptimizedMemory;
const memorySaved = Math.max(0, baselineMemory - optimizedMemory);
const memorySavedPercent =
baselineMemory > 0 ? (memorySaved / baselineMemory) * 100 : 0;
return {
totalOperations: this.metrics.operations,
averageSpeedup: avgSpeedup,
peakSpeedup: this.metrics.peakSpeedup,
averageExecutionTimeMs:
this.metrics.operations > 0
? this.metrics.totalExecutionTime / this.metrics.operations
: 0,
totalMemorySavedBytes: memorySaved,
successRate:
this.metrics.operations > 0
? (this.metrics.successfulOperations / this.metrics.operations) * 100
: 0,
// Memory tracking metrics
baselineMemoryBytes: baselineMemory,
optimizedMemoryBytes: optimizedMemory,
memorySavedBytes: memorySaved,
memorySavedPercent: memorySavedPercent,
peakMemoryBytes: this.metrics.peakMemory,
};
}
/**
* Reset all metrics
*/
resetMetrics(): void {
this.metrics = {
operations: 0,
totalSpeedup: 0,
peakSpeedup: 0,
totalExecutionTime: 0,
successfulOperations: 0,
totalBaselineMemory: 0,
totalOptimizedMemory: 0,
peakMemory: 0,
};
}
/**
* Ensure input is Float32Array for optimal performance
*/
private ensureFloat32Array(input: ArrayInput): Float32Array {
if (input instanceof Float32Array) {
return input;
}
return new Float32Array(input);
}
/**
* Detect which runtime is being used
*/
private detectRuntime(): 'napi' | 'wasm' | 'js' {
// Check if NAPI bindings are available
try {
if (typeof process !== 'undefined' && process.versions && 'napi' in process.versions) {
return 'napi';
}
} catch {
// Not in Node.js environment
}
// Check for WebAssembly support
if (typeof globalThis !== 'undefined' && 'WebAssembly' in globalThis) {
return 'wasm';
}
// Fallback to pure JS
return 'js';
}
/**
* Get current memory usage in bytes
*/
private getMemoryUsage(): number {
if (typeof process !== 'undefined' && process.memoryUsage) {
return process.memoryUsage().heapUsed;
}
return 0;
}
/**
* Force garbage collection if available (requires --expose-gc flag)
* This helps get more accurate memory measurements
*/
private forceGC(): void {
if (typeof global !== 'undefined' && typeof (global as any).gc === 'function') {
(global as any).gc();
}
}
/**
* Benchmark memory usage across multiple dimensions
* Validates the 50-75% memory reduction target
* @param dimensions - Array of dimensions to test (default: [128, 256, 512, 1024])
* @returns Memory profiling results for each dimension
*/
benchmarkMemory(
dimensions: number[] = [128, 256, 512, 1024]
): {
dimension: number;
baselineMemoryBytes: number;
optimizedMemoryBytes: number;
memorySavedBytes: number;
memorySavedPercent: number;
meetsTarget: boolean; // true if 50-75% reduction achieved
}[] {
const results: {
dimension: number;
baselineMemoryBytes: number;
optimizedMemoryBytes: number;
memorySavedBytes: number;
memorySavedPercent: number;
meetsTarget: boolean;
}[] = [];
for (const dim of dimensions) {
const numKeys = 100;
const iterations = 100; // Fewer iterations for memory profiling
// Create test data
const query = new Float32Array(dim);
const keys = Array.from({ length: numKeys }, () => new Float32Array(dim));
const values = Array.from({ length: numKeys }, () => new Float32Array(dim));
// Fill with random data
for (let i = 0; i < dim; i++) {
query[i] = Math.random();
}
for (let i = 0; i < numKeys; i++) {
for (let j = 0; j < dim; j++) {
keys[i][j] = Math.random();
values[i][j] = Math.random();
}
}
// Create temporary instances for this dimension
const flashAttention = new FlashAttention(dim, this.blockSize);
const baselineAttention = new DotProductAttention(dim);
// Measure baseline memory
this.forceGC();
const baselineMemBefore = this.getMemoryUsage();
let baselinePeak = baselineMemBefore;
for (let i = 0; i < iterations; i++) {
baselineAttention.computeRaw(query, keys, values);
if (i % 10 === 0) {
const curr = this.getMemoryUsage();
if (curr > baselinePeak) baselinePeak = curr;
}
}
const baselineMemoryUsed = Math.max(0, baselinePeak - baselineMemBefore);
// Measure Flash Attention memory
this.forceGC();
const flashMemBefore = this.getMemoryUsage();
let flashPeak = flashMemBefore;
for (let i = 0; i < iterations; i++) {
flashAttention.computeRaw(query, keys, values);
if (i % 10 === 0) {
const curr = this.getMemoryUsage();
if (curr > flashPeak) flashPeak = curr;
}
}
const flashMemoryUsed = Math.max(0, flashPeak - flashMemBefore);
const memorySaved = Math.max(0, baselineMemoryUsed - flashMemoryUsed);
const memorySavedPercent =
baselineMemoryUsed > 0 ? (memorySaved / baselineMemoryUsed) * 100 : 0;
// Target: 50-75% memory reduction
const meetsTarget = memorySavedPercent >= 50 && memorySavedPercent <= 100;
results.push({
dimension: dim,
baselineMemoryBytes: baselineMemoryUsed,
optimizedMemoryBytes: flashMemoryUsed,
memorySavedBytes: memorySaved,
memorySavedPercent: memorySavedPercent,
meetsTarget: meetsTarget,
});
// Update global metrics
this.metrics.totalBaselineMemory += baselineMemoryUsed;
this.metrics.totalOptimizedMemory += flashMemoryUsed;
if (flashPeak > this.metrics.peakMemory) {
this.metrics.peakMemory = flashPeak;
}
}
return results;
}
}
// ============================================================================
// Convenience Factory Functions
// ============================================================================
/**
* Create a Flash Attention optimizer with default settings
* @param dim - Dimension of attention vectors (default: 512)
* @param blockSize - Block size for Flash Attention (default: 64)
* @returns Configured FlashAttentionOptimizer instance
*/
export function createFlashAttentionOptimizer(
dim: number = 512,
blockSize: number = 64
): FlashAttentionOptimizer {
return new FlashAttentionOptimizer(dim, blockSize);
}
/**
* Quick benchmark of Flash Attention performance
* @param dim - Dimension to test (default: 512)
* @returns Benchmark results with speedup metrics
*/
export function quickBenchmark(dim: number = 512): BenchmarkResult {
const optimizer = createFlashAttentionOptimizer(dim);
return optimizer.benchmark();
}