webinfer 0.0.2 → 0.0.3

package/README.md

@@ -11,44 +11,34 @@ npm install webinfer
 ## Usage
 
 ```typescript
-import { WebInferDevice, Tensor, matmul, flashAttention } from 'webinfer';
+import { WebInferDevice, attention } from 'webinfer';
 
-// Initialize WebGPU device
 const device = await WebInferDevice.create();
 
-// Matrix multiplication
-const a = Tensor.rand(device, [1024, 1024]);
-const b = Tensor.rand(device, [1024, 1024]);
-const c = await matmul(device, a, b);
+// Single decode attention
+const q = new Float32Array(32 * 128);  // [num_qo_heads, head_dim]
+const k = new Float32Array(2048 * 32 * 128);  // [kv_len, num_kv_heads, head_dim]
+const v = new Float32Array(2048 * 32 * 128);
 
-// Read result back to CPU
-const result = await c.toArray();
+const output = await attention(device, { q, k, v });
 ```
 
-## Operations
+## API
 
-| Category | Operations |
-|----------|------------|
-| **Core** | matmul, flashAttention |
-| **Normalization** | rmsNorm, layerNorm |
-| **Activations** | gelu, silu, relu, softmax |
-| **Position** | rope (rotary embeddings) |
-| **Sampling** | topK, topP, sample |
-| **Quantization** | INT4/INT8 quantized matmul |
-| **Model Loading** | SafeTensors, GGUF |
+| Category | Exports |
+|----------|---------|
+| **Attention** | `attention`, `BatchAttention`, `AttentionKernel`, `cascadedAttention` |
+| **KV Cache** | `PagedKVCache`, `BlockManager`, `pagedAttention` |
+| **Patterns** | `buildCausalMask`, `buildSlidingWindowMask`, `buildBlockSparseCSR` |
+| **Sampling** | `topKSamplingFromProbs`, `topPSamplingFromProbs`, `minPSamplingFromProbs`, `topKTopPSamplingFromLogits` |
+| **Normalization** | `rmsNorm`, `layerNorm`, `fusedAddRmsNorm`, `gemmaRmsNorm` |
+| **Core** | `matmul`, `rope`, `gelu`, `silu`, `softmax` |
 
-## Requirements
-
-- Browser with WebGPU support (Chrome 113+, Edge 113+)
-- Or Node.js with `@aspect-build/aspect-cli` for server-side WebGPU
-
-## Benchmarks
+## Release
 
 ```bash
-git clone https://github.com/guan404ming/webinfer
-cd webinfer
-bun install
-bun run bench
+bun run build
+npm publish
 ```
 
 ## License

package/dist/attention/block-sparse/patterns/tree.d.ts

@@ -0,0 +1,65 @@
+/**
+ * Tree Attention Pattern
+ * Used in speculative decoding (Medusa, EAGLE) and tree-based generation
+ */
+import { type BlockSparseCSR } from "../format.ts";
+/**
+ * Tree node structure
+ */
+export interface TreeNode {
+    /** Token position in sequence */
+    position: number;
+    /** Parent position (-1 for root) */
+    parent: number;
+    /** Depth in tree (0 for root) */
+    depth: number;
+}
+/**
+ * Tree attention configuration
+ */
+export interface TreeAttentionConfig {
+    /** Total sequence length including prompt */
+    seqLen: number;
+    /** Prompt length (prefix that all tokens attend to) */
+    promptLen: number;
+    /** Tree structure for speculative tokens */
+    tree: TreeNode[];
+    /** Block size for sparse format */
+    blockSize?: number;
+}
+/**
+ * Build tree attention mask
+ *
+ * In tree attention:
+ * - All tokens attend to the prompt (positions 0 to promptLen-1)
+ * - Tree tokens attend to their ancestors in the tree
+ * - Maintains causal property within the tree structure
+ */
+export declare function buildTreeMask(config: TreeAttentionConfig): BlockSparseCSR;
+/**
+ * Build a simple chain tree (linear speculation)
+ * Each token depends on the previous one
+ */
+export declare function buildChainTree(numSpecTokens: number): TreeNode[];
+/**
+ * Build a wide tree (parallel speculation)
+ * All speculative tokens depend only on the prompt
+ */
+export declare function buildWideTree(numSpecTokens: number): TreeNode[];
+/**
+ * Build a binary tree for speculation
+ */
+export declare function buildBinaryTree(depth: number): TreeNode[];
+/**
+ * Build Medusa-style tree
+ * Multiple heads predict tokens at different positions
+ */
+export declare function buildMedusaTree(numHeads: number, tokensPerHead: number): TreeNode[];
+/**
+ * Calculate tree sparsity ratio
+ */
+export declare function getTreeSparsity(config: TreeAttentionConfig): number;
+/**
+ * Validate tree structure
+ */
+export declare function validateTree(tree: TreeNode[]): boolean;

package/dist/attention/cascaded-inference.d.ts

@@ -0,0 +1,29 @@
+/**
+ * Cascaded Inference - TDR-safe attention for very long sequences
+ * Splits attention computation across multiple passes with browser yields
+ */
+import type { WebInferDevice } from "../core/device.ts";
+import { Tensor } from "../core/tensor.ts";
+export interface CascadedAttentionConfig {
+    numHeads: number;
+    headDim: number;
+    seqLen: number;
+    scale?: number;
+    causal?: boolean;
+}
+/**
+ * Cascaded Attention - Safe for very long sequences
+ * Uses Split-K strategy to prevent TDR (GPU timeout)
+ *
+ * @param device WebInfer device
+ * @param q Query tensor [seqLen, numHeads, headDim]
+ * @param k Key tensor [seqLen, numHeads, headDim]
+ * @param v Value tensor [seqLen, numHeads, headDim]
+ * @param config Attention configuration
+ * @param onProgress Optional progress callback
+ */
+export declare function cascadedAttention(device: WebInferDevice, q: Tensor, k: Tensor, v: Tensor, config: CascadedAttentionConfig, onProgress?: (chunk: number, total: number) => void): Promise<Tensor>;
+/**
+ * CPU reference implementation for verification
+ */
+export declare function cascadedAttentionCPU(q: Float32Array, k: Float32Array, v: Float32Array, seqLen: number, numHeads: number, headDim: number, causal?: boolean): Float32Array;

package/dist/attention/index.d.ts

@@ -1,9 +1,118 @@
 /**
  * Attention Module Exports
  */
-export { flashAttention, attentionCPU, type AttentionConfig, } from "./flash-attention.ts";
+import type { WebInferDevice } from "../core/device.ts";
+import { Tensor } from "../core/tensor.ts";
+import type { BlockSparseCSR } from "./block-sparse/format.ts";
+import { PagedKVCache } from "./paged-kv/page-table.ts";
+/**
+ * Tensor-like input types
+ */
+export type TensorLike = Tensor | Float32Array | {
+    data: Float32Array;
+    shape: number[];
+};
+/**
+ * Attention options
+ */
+export interface AttentionOptions {
+    q: TensorLike;
+    k: TensorLike;
+    v: TensorLike;
+    causal?: boolean;
+    scale?: number;
+    window?: number;
+    mask?: BlockSparseCSR;
+    returnLse?: boolean;
+}
+/**
+ * Attention result with optional LSE
+ */
+export interface AttentionResultWithLse {
+    output: Tensor;
+    lse: Tensor;
+}
+/**
+ * Simple attention function
+ */
+export declare function attention(device: WebInferDevice, options: AttentionOptions): Promise<Tensor>;
+export declare function attention(device: WebInferDevice, options: AttentionOptions & {
+    returnLse: true;
+}): Promise<AttentionResultWithLse>;
+/**
+ * BatchAttention configuration
+ */
+export interface BatchAttentionConfig {
+    numHeads: number;
+    headDim: number;
+    numKvHeads?: number;
+    maxBatchSize?: number;
+    maxSeqLen?: number;
+}
+/**
+ * Prefill input
+ */
+export interface PrefillInput {
+    queries: TensorLike[];
+    keys: TensorLike[];
+    values: TensorLike[];
+    causal?: boolean;
+    window?: number;
+}
+/**
+ * Decode input
+ */
+export interface DecodeInput {
+    query: TensorLike;
+    kvCache: PagedKVCache;
+    seqIds: number[];
+}
+/**
+ * BatchAttention - Batched attention for prefill and decode
+ */
+export declare class BatchAttention {
+    private device;
+    private config;
+    constructor(device: WebInferDevice, config: BatchAttentionConfig);
+    /**
+     * Prefill: Process variable-length sequences
+     */
+    prefill(input: PrefillInput): Promise<Tensor[]>;
+    /**
+     * Decode: Single token per sequence with KV cache
+     */
+    decode(input: DecodeInput): Promise<Tensor>;
+    getConfig(): Required<BatchAttentionConfig>;
+    dispose(): void;
+}
+/**
+ * AttentionKernel configuration
+ */
+export interface AttentionKernelConfig {
+    numHeads: number;
+    headDim: number;
+    causal?: boolean;
+    blockSize?: number;
+}
+/**
+ * AttentionKernel - Low-level compiled kernel
+ */
+export declare class AttentionKernel {
+    private device;
+    private config;
+    private constructor();
+    static compile(device: WebInferDevice, config: AttentionKernelConfig): Promise<AttentionKernel>;
+    execute(input: {
+        q: Tensor;
+        k: Tensor;
+        v: Tensor;
+    }): Promise<Tensor>;
+    dispose(): void;
+}
+export { flashAttention, type AttentionConfig } from "./flash-attention.ts";
 export { buildBlockSparseCSR, getSparsityRatio, estimateMemorySavings, type BlockSparseCSR, type AttentionPattern, } from "./block-sparse/format.ts";
 export { buildCausalMask, getCausalSparsity } from "./block-sparse/patterns/causal.ts";
 export { buildSlidingWindowMask, buildCausalSlidingWindowMask, getSlidingWindowSparsity, } from "./block-sparse/patterns/sliding.ts";
-export { AttentionScheduler, type ChunkPlan } from "./scheduler.ts";
-export { PagedKVCache, type PagedKVCacheConfig, type SequenceEntry, BlockManager, ContinuousBatchScheduler, type BlockManagerConfig, type AllocationPolicy, type AllocationRequest, } from "./paged-kv/index.ts";
+export { PagedKVCache, type PagedKVCacheConfig, type SequenceEntry, type DefragmentResult, BlockManager, ContinuousBatchScheduler, type BlockManagerConfig, type AllocationPolicy, type AllocationRequest, } from "./paged-kv/index.ts";
+export { pagedAttention, appendToPagedCache, type PagedAttentionConfig, type PagedAttentionInput, } from "./paged-attention.ts";
+export { cascadedAttention, type CascadedAttentionConfig, } from "./cascaded-inference.ts";

package/dist/attention/paged-attention.d.ts

@@ -0,0 +1,40 @@
+/**
+ * Paged Attention Implementation for WebGPU
+ * Efficient attention computation using paged KV cache (vLLM-style)
+ */
+import type { WebInferDevice } from "../core/device.ts";
+import { Tensor } from "../core/tensor.ts";
+import type { PagedKVCache } from "./paged-kv/page-table.ts";
+export interface PagedAttentionConfig {
+    numHeads: number;
+    headDim: number;
+    scale?: number;
+}
+export interface PagedAttentionInput {
+    query: Tensor;
+    kvCache: PagedKVCache;
+    seqIds: number[];
+    positions: number[];
+}
+/**
+ * Paged Attention forward pass for decoding
+ * Computes attention against paged KV cache for single-token queries
+ *
+ * @param device WebInfer device
+ * @param input Paged attention input (query, kv cache, sequence info)
+ * @param config Attention configuration
+ */
+export declare function pagedAttention(device: WebInferDevice, input: PagedAttentionInput, config: PagedAttentionConfig): Promise<Tensor>;
+/**
+ * Append new KV to paged cache
+ */
+export declare function appendToPagedCache(device: WebInferDevice, kvCache: PagedKVCache, seqId: number, key: Tensor, // [numHeads, headDim]
+value: Tensor): Promise<void>;
+/**
+ * CPU reference implementation for verification
+ */
+export declare function pagedAttentionCPU(q: Float32Array, // [batchSize, numHeads, headDim]
+keyCache: Float32Array, // [maxPages, pageSize, numHeads, headDim]
+valueCache: Float32Array, pageTable: number[][], // [batchSize][pages]
+seqLens: number[], // [batchSize]
+numHeads: number, headDim: number, pageSize: number, maxPages: number): Float32Array;

package/dist/attention/paged-kv/index.d.ts

@@ -1,5 +1,5 @@
 /**
  * PagedKV Module Exports
  */
-export { PagedKVCache, type PagedKVCacheConfig, type SequenceEntry, } from "./page-table.ts";
-export { BlockManager, ContinuousBatchScheduler, type BlockManagerConfig, type AllocationPolicy, type AllocationRequest, } from "./block-manager.ts";
+export { type AllocationPolicy, type AllocationRequest, BlockManager, type BlockManagerConfig, ContinuousBatchScheduler, } from "./block-manager.ts";
+export { type DefragmentResult, PagedKVCache, type PagedKVCacheConfig, type SequenceEntry, } from "./page-table.ts";

package/dist/attention/paged-kv/page-table.d.ts

@@ -96,4 +96,70 @@ export declare class PagedKVCache {
      * Dispose GPU resources
      */
     dispose(): void;
+    /**
+     * Allocate a new sequence (v2 alias)
+     */
+    alloc(initialLength?: number): number;
+    /**
+     * Append KV to sequence (v2 API)
+     */
+    append(seqId: number, kv: {
+        key: Float32Array;
+        value: Float32Array;
+        layer?: number;
+    }): void;
+    /**
+     * Batch append KV for all layers (v2 API)
+     */
+    appendBatch(seqId: number, kv: {
+        keys: Float32Array;
+        values: Float32Array;
+    }): void;
+    /**
+     * Free sequence (v2 alias)
+     */
+    free(seqId: number): void;
+    /**
+     * Get stats (v2 API)
+     */
+    stats(): {
+        usedPages: number;
+        freePages: number;
+        fragmentation: number;
+    };
+    /**
+     * Check if defrag needed (v2 alias)
+     */
+    needsDefrag(threshold?: number): boolean;
+    /**
+     * Check if defragmentation is needed
+     * Returns fragmentation ratio (0 = no fragmentation, 1 = fully fragmented)
+     */
+    getFragmentationRatio(): number;
+    /**
+     * Check if defragmentation would be beneficial
+     */
+    needsDefragmentation(threshold?: number): boolean;
+    /**
+     * Defragment the KV cache by compacting pages
+     * Returns the number of pages moved
+     */
+    defragment(): Promise<DefragmentResult>;
+    /**
+     * Move a single page from one location to another
+     */
+    private movePage;
+    /**
+     * Rebuild the free page list from scratch
+     */
+    private rebuildFreePageList;
+}
+/**
+ * Result of defragmentation operation
+ */
+export interface DefragmentResult {
+    pagesMoved: number;
+    durationMs: number;
+    fragmentationBefore: number;
+    fragmentationAfter: number;
 }

package/dist/core/tdr.d.ts

@@ -0,0 +1,114 @@
+/**
+ * TDR (Timeout Detection and Recovery) Prevention
+ * Handles browser-specific GPU timeout limits and graceful degradation
+ */
+/**
+ * Browser detection result
+ */
+export interface BrowserInfo {
+    name: "chrome" | "safari" | "firefox" | "edge" | "unknown";
+    version: number;
+    isMobile: boolean;
+    hasWebGPU: boolean;
+}
+/**
+ * TDR configuration for each browser
+ */
+export interface TDRConfig {
+    timeoutMs: number;
+    safetyMargin: number;
+    maxChunkSize: number;
+    supportsTimestampQuery: boolean;
+}
+/**
+ * Detect browser information
+ */
+export declare function detectBrowser(): BrowserInfo;
+/**
+ * Get TDR configuration for current browser
+ */
+export declare function getTDRConfig(browser?: BrowserInfo): TDRConfig;
+/**
+ * Graceful degradation options
+ */
+export interface DegradationOptions {
+    /** Use CPU fallback when GPU fails */
+    enableCpuFallback: boolean;
+    /** Reduce precision to f16 when memory is tight */
+    enablePrecisionReduction: boolean;
+    /** Auto-chunk large sequences */
+    enableAutoChunking: boolean;
+    /** Maximum retries before falling back */
+    maxRetries: number;
+}
+/**
+ * Error types for graceful handling
+ */
+export type WebGPUErrorType = "device_lost" | "out_of_memory" | "validation" | "timeout" | "unknown";
+/**
+ * Classify WebGPU error
+ */
+export declare function classifyError(error: unknown): WebGPUErrorType;
+/**
+ * TDR-safe execution wrapper
+ */
+export declare class TDRGuard {
+    private browser;
+    private config;
+    private options;
+    private lastExecutionTime;
+    constructor(options?: Partial<DegradationOptions>);
+    /**
+     * Get safe execution time limit
+     */
+    getSafeTimeLimit(): number;
+    /**
+     * Check if operation might cause TDR
+     */
+    mightTimeout(estimatedMs: number): boolean;
+    /**
+     * Calculate chunks needed for safe execution
+     */
+    calcChunks(estimatedTotalMs: number): number;
+    /**
+     * Yield to browser main thread
+     */
+    yield(): Promise<void>;
+    /**
+     * Execute with TDR protection
+     */
+    execute<T>(fn: () => Promise<T>, options?: {
+        estimatedMs?: number;
+        onRetry?: (attempt: number, error: unknown) => void;
+        fallback?: () => T | Promise<T>;
+    }): Promise<T>;
+    /**
+     * Get browser info
+     */
+    getBrowserInfo(): BrowserInfo;
+    /**
+     * Get TDR config
+     */
+    getTDRConfig(): TDRConfig;
+    /**
+     * Get last execution time
+     */
+    getLastExecutionTime(): number;
+}
+/**
+ * Check WebGPU support and capabilities
+ */
+export declare function checkWebGPUSupport(): Promise<{
+    supported: boolean;
+    reason?: string;
+    adapter?: GPUAdapter;
+    limits?: GPUSupportedLimits;
+}>;
+/**
+ * Create device with graceful degradation
+ */
+export declare function createDeviceWithFallback(adapter: GPUAdapter, options?: {
+    requiredFeatures?: GPUFeatureName[];
+    requiredLimits?: Record<string, number>;
+    onFallback?: (reason: string) => void;
+}): Promise<GPUDevice>;

package/dist/index.d.ts

@@ -5,18 +5,20 @@
 export { WebInferDevice, type DeviceInfo } from "./core/device.ts";
 export { Tensor, type DType } from "./core/tensor.ts";
 export { BufferPool } from "./core/buffer-pool.ts";
+export { TDRGuard, detectBrowser, getTDRConfig, classifyError, checkWebGPUSupport, createDeviceWithFallback, type BrowserInfo, type TDRConfig, type DegradationOptions, type WebGPUErrorType, } from "./core/tdr.ts";
+export { attention, BatchAttention, AttentionKernel, type AttentionOptions, type AttentionResultWithLse, type BatchAttentionConfig, type PrefillInput, type DecodeInput, type TensorLike, type AttentionKernelConfig, } from "./attention/index.ts";
 export { KernelCache, type CacheStats } from "./jit/kernel-cache.ts";
 export { WGSLCompiler, type MatMulConfig } from "./jit/compiler.ts";
-export { matmul, matmulCPU, getMatMulCacheStats } from "./ops/matmul.ts";
-export { layerNorm, layerNormCPU, rmsNorm, rmsNormCPU, } from "./ops/normalization.ts";
-export { rope, ropeCPU, computeRoPEFrequencies, type RoPEConfig, } from "./ops/rope.ts";
-export { gelu, geluCPU, geluExactCPU, silu, siluCPU, relu, reluCPU, sigmoidCPU, } from "./ops/activations.ts";
-export { softmaxGPU, softmaxCPU, logSoftmaxCPU } from "./ops/softmax.ts";
-export { add, addCPU, mul, mulCPU, scale, scaleCPU, addScalarCPU, fmaCPU, } from "./ops/elementwise.ts";
-export { embedding, embeddingCPU, batchedEmbeddingCPU, } from "./ops/embedding.ts";
-export { transpose2D, transpose2DCPU, transposeCPU, reshapeCPU, permuteCPU, } from "./ops/reshape.ts";
-export { quantizeToInt8, quantizeToInt4, dequantizeInt8, dequantizeInt4, quantizationError, getMemorySavings, qmatmulInt8CPU, qmatmulInt4CPU, qmatmulInt8BlockCPU, estimateQMatMulFlops, estimateQMatMulBandwidth, type QuantConfig, type QuantizedTensor, } from "./quantization/index.ts";
-export { flashAttention, attentionCPU, type AttentionConfig, buildBlockSparseCSR, getSparsityRatio, estimateMemorySavings, type BlockSparseCSR, type AttentionPattern, buildCausalMask, getCausalSparsity, buildSlidingWindowMask, buildCausalSlidingWindowMask, getSlidingWindowSparsity, AttentionScheduler, type ChunkPlan, PagedKVCache, type PagedKVCacheConfig, type SequenceEntry, BlockManager, ContinuousBatchScheduler, type BlockManagerConfig, type AllocationPolicy, type AllocationRequest, } from "./attention/index.ts";
-export { topK, topKCPU, topKFilter, topPFilter, topPFilterCPU, sample, sampleCPU, sampleGreedy, sampleFromProbs, softmax, applyRepetitionPenalty, type SamplingConfig, } from "./sampling/index.ts";
+export { matmul, getMatMulCacheStats } from "./ops/matmul.ts";
+export { layerNorm, rmsNorm, fusedAddRmsNorm, gemmaRmsNorm, gemmaFusedAddRmsNorm, } from "./ops/normalization.ts";
+export { rope, computeRoPEFrequencies, type RoPEConfig } from "./ops/rope.ts";
+export { gelu, silu, relu } from "./ops/activations.ts";
+export { softmaxGPU } from "./ops/softmax.ts";
+export { add, mul, scale } from "./ops/elementwise.ts";
+export { embedding } from "./ops/embedding.ts";
+export { transpose2D } from "./ops/reshape.ts";
+export { quantizeToInt8, quantizeToInt4, dequantizeInt8, dequantizeInt4, quantizationError, getMemorySavings, estimateQMatMulFlops, estimateQMatMulBandwidth, type QuantConfig, type QuantizedTensor, } from "./quantization/index.ts";
+export { flashAttention, type AttentionConfig, buildBlockSparseCSR, getSparsityRatio, estimateMemorySavings, type BlockSparseCSR, type AttentionPattern, buildCausalMask, getCausalSparsity, buildSlidingWindowMask, buildCausalSlidingWindowMask, getSlidingWindowSparsity, PagedKVCache, type PagedKVCacheConfig, type SequenceEntry, type DefragmentResult, BlockManager, ContinuousBatchScheduler, type BlockManagerConfig, type AllocationPolicy, type AllocationRequest, pagedAttention, appendToPagedCache, type PagedAttentionConfig, type PagedAttentionInput, cascadedAttention, type CascadedAttentionConfig, } from "./attention/index.ts";
+export { topK, topKFilter, topPFilter, sample, sampleGreedy, sampleFromProbs, softmax, applyRepetitionPenalty, minPSamplingFromProbs, topKSamplingFromProbs, topPSamplingFromProbs, topKTopPSamplingFromProbs, topKTopPSamplingFromLogits, topPRenormProbs, topKRenormProbs, topKMaskLogits, type SamplingConfig, } from "./sampling/index.ts";
 export { type ModelFormat, type SafetensorsDType, GGUFQuantType, GGUFMetadataValueType, type TensorInfo, type SafetensorsHeader, type ModelMetadata, type LoadedModel, type LoadOptions, parseSafetensorsHeader, loadSafetensors, loadSafetensorsFromUrl, isSafetensors, parseGGUFHeader, loadGGUF, loadGGUFFromUrl, loadGGUFTensor, isGGUF, loadModel, } from "./model/index.ts";
 export { type ModelConfig, type InferenceConfig, type GenerationConfig, type GenerationResult, type StreamToken, type FinishReason, type ForwardResult, DEFAULT_GENERATION_CONFIG, normalizeGenerationConfig, InferenceEngine, generate, generateStream, greedyDecode, sampleNextToken, } from "./inference/index.ts";