webinfer 0.0.2 → 0.0.4

package/README.md

@@ -2,53 +2,58 @@
 
 High-performance LLM inference kernels for WebGPU.
 
+WebGPU implementation of FlashInfer APIs for browser-based LLM inference.
+
+## Features
+
+- **Flash Attention for WebGPU**: Online softmax and paged KV cache for memory-efficient attention
+- **Subgroup Optimization**: Uses WebGPU subgroup operations for 2-4x faster reductions when available
+- **Fused Kernels**: Reduces memory bandwidth with operations like `silu_and_mul`, `fused_add_rmsnorm`
+- **JIT Compilation**: Runtime kernel generation and pipeline caching
+- **f16 Safe Accumulation**: Uses f32 accumulators for f16 computation to prevent precision loss
+- **FlashInfer Compatible API**: Drop-in replacement for FlashInfer with identical naming
+
 ## Install
 
 ```bash
 npm install webinfer
 ```
 
-## Usage
-
-```typescript
-import { WebInferDevice, Tensor, matmul, flashAttention } from 'webinfer';
+## API Status
 
-// 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);
-
-// Read result back to CPU
-const result = await c.toArray();
-```
+**Phase 1: Core + Simple Ops**
+- `webinfer.WebInferContext.create()` - Initialize WebGPU context
+- `webinfer.norm.rmsnorm()` - Root mean square normalization
+- `webinfer.activation.silu_and_mul()` - SiLU activation with gating
+- `webinfer.activation.gelu_and_mul()` - GELU activation with gating
 
-## Operations
+**Phase 2: Single Attention**
+- `webinfer.decode.single_decode_with_kv_cache()` - Flash attention for decode phase
+- `webinfer.prefill.single_prefill_with_kv_cache()` - Flash attention for prefill phase with causal masking
 
-| 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 |
+**Phase 3: Paged KV + Batched Operations**
+- `webinfer.page.append_paged_kv_cache()` - Append new KV pairs to paged cache
+- `webinfer.decode.BatchDecodeWithPagedKVCacheWrapper` - Batched decode with paged KV cache
+- `webinfer.prefill.BatchPrefillWithPagedKVCacheWrapper` - Batched prefill with paged KV cache
 
-## Requirements
+**Phase 4: Additional Operations**
+- `webinfer.rope.apply_rope_inplace()` - Apply RoPE position encoding
+- `webinfer.rope.apply_llama31_rope_inplace()` - Llama 3.1 RoPE variant for long context
+- `webinfer.sampling.sampling_from_probs()` - Categorical sampling
+- `webinfer.sampling.top_k_sampling_from_probs()` - Top-k sampling
+- `webinfer.sampling.top_p_sampling_from_probs()` - Top-p (nucleus) sampling
+- `webinfer.sampling.min_p_sampling_from_probs()` - Min-p sampling
+- `webinfer.sampling.top_k_top_p_sampling_from_probs()` - Combined top-k and top-p
+- `webinfer.gemm.bmm_fp16()` - Batched matrix multiplication (fp16)
+- `webinfer.gemm.bmm_fp32()` - Batched matrix multiplication (fp32)
 
-- Browser with WebGPU support (Chrome 113+, Edge 113+)
-- Or Node.js with `@aspect-build/aspect-cli` for server-side WebGPU
+See [dev/design.md](dev/design.md) for full design and implementation plan.
 
-## 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/activation/index.d.ts

@@ -0,0 +1,30 @@
+/**
+ * Activation functions
+ */
+import type { WebInferContext } from '../core/context.ts';
+import type { Tensor } from '../core/tensor.ts';
+/**
+ * SiLU (Swish) activation with gating
+ *
+ * Input tensor has shape [..., 2 * hidden_size]
+ * Splits into gate[..., :hidden_size] and up[..., hidden_size:]
+ * Computes: output = silu(gate) * up
+ * Where silu(x) = x * sigmoid(x) = x / (1 + exp(-x))
+ *
+ * @param ctx WebInfer context
+ * @param input Input tensor of shape [..., 2 * hidden_size]
+ * @returns Activated tensor of shape [..., hidden_size]
+ */
+export declare function silu_and_mul(ctx: WebInferContext, input: Tensor): Promise<Tensor>;
+/**
+ * GELU activation with gating
+ *
+ * Input tensor has shape [..., 2 * hidden_size]
+ * Computes: output = gelu(gate) * up
+ * Where gelu(x) ≈ 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x^3)))
+ *
+ * @param ctx WebInfer context
+ * @param input Input tensor of shape [..., 2 * hidden_size]
+ * @returns Activated tensor of shape [..., hidden_size]
+ */
+export declare function gelu_and_mul(ctx: WebInferContext, input: Tensor): Promise<Tensor>;

package/dist/core/context.d.ts

@@ -0,0 +1,60 @@
+/**
+ * WebInferContext manages WebGPU device and provides core functionality
+ */
+import type { DType } from './types.ts';
+import { Tensor } from './tensor.ts';
+/**
+ * Device vendor detection
+ */
+export type Vendor = 'apple' | 'nvidia' | 'amd' | 'intel' | 'unknown';
+/**
+ * Device features detected at runtime
+ */
+export interface DeviceFeatures {
+    readonly vendor: Vendor;
+    readonly hasSubgroups: boolean;
+    readonly hasF16: boolean;
+    readonly maxWorkgroupSize: number;
+    readonly maxStorageBufferSize: number;
+    readonly maxBufferSize: number;
+    readonly maxWorkgroupsPerDimension: number;
+}
+/**
+ * WebInferContext - main context for all WebInfer operations
+ */
+export declare class WebInferContext {
+    readonly device: GPUDevice;
+    readonly features: DeviceFeatures;
+    private pipelineCache;
+    private disposed;
+    private constructor();
+    /**
+     * Create a new WebInferContext
+     */
+    static create(): Promise<WebInferContext>;
+    /**
+     * Create a GPU buffer
+     */
+    createBuffer(size: number, usage?: GPUBufferUsageFlags): GPUBuffer;
+    /**
+     * Create a tensor with uninitialized data
+     */
+    createTensor(shape: number[], dtype: DType): Tensor;
+    /**
+     * Create a tensor from TypedArray data
+     */
+    createTensorFromData(data: Float32Array | Uint16Array, shape: number[], dtype: DType): Tensor;
+    /**
+     * Get or create a cached compute pipeline
+     */
+    getOrCreatePipeline(key: string, createFn: () => Promise<GPUComputePipeline> | GPUComputePipeline): Promise<GPUComputePipeline>;
+    /**
+     * Clear pipeline cache
+     */
+    clearPipelineCache(): void;
+    /**
+     * Dispose the context and cleanup resources
+     */
+    dispose(): void;
+    private checkDisposed;
+}

package/dist/core/paged-kv-cache.d.ts

@@ -0,0 +1,33 @@
+/**
+ * Paged KV Cache implementation
+ */
+import type { Tensor } from './tensor.ts';
+/**
+ * Paged KV Cache for efficient memory management
+ *
+ * Layout: [num_pages, 2, page_size, num_kv_heads, head_dim]
+ * where dim 1 has [0] = keys, [1] = values
+ */
+export interface PagedKvCache {
+    /**
+     * The backing tensor for the paged KV cache
+     * Shape: [num_pages, 2, page_size, num_kv_heads, head_dim]
+     */
+    data: Tensor;
+    /**
+     * Number of tokens per page
+     */
+    page_size: number;
+    /**
+     * Number of KV heads (for GQA/MQA)
+     */
+    num_kv_heads: number;
+    /**
+     * Dimension of each head
+     */
+    head_dim: number;
+}
+/**
+ * Create a paged KV cache
+ */
+export declare function createPagedKvCache(data: Tensor, page_size: number, num_kv_heads: number, head_dim: number): PagedKvCache;

package/dist/core/tensor.d.ts

@@ -1,25 +1,44 @@
 /**
- * Tensor - GPU-backed multidimensional array
+ * Tensor wrapper for WebGPU buffers
+ */
+import type { DType } from './types.ts';
+/**
+ * Compute strides from shape (row-major order)
+ */
+export declare function computeStrides(shape: number[]): number[];
+/**
+ * Compute total number of elements from shape
+ */
+export declare function computeSize(shape: number[]): number;
+/**
+ * Tensor class wrapping a GPUBuffer with shape and dtype information
  */
-import type { WebInferDevice } from "./device.ts";
-export type DType = "f32" | "f16" | "i32" | "u32";
 export declare class Tensor {
-    private _device;
-    private _shape;
-    private _dtype;
-    private _buffer;
-    private _disposed;
-    constructor(device: WebInferDevice, shape: number[], dtype?: DType, data?: Float32Array | Uint32Array | Int32Array);
-    static fromArray(device: WebInferDevice, shape: number[], data: Float32Array, dtype?: DType): Promise<Tensor>;
-    static zeros(device: WebInferDevice, shape: number[], dtype?: DType): Tensor;
-    static rand(device: WebInferDevice, shape: number[], dtype?: DType): Tensor;
-    get shape(): readonly number[];
-    get dtype(): DType;
-    get numel(): number;
-    get byteSize(): number;
-    get buffer(): GPUBuffer;
-    get device(): WebInferDevice;
-    toArray(): Promise<Float32Array>;
+    readonly buffer: GPUBuffer;
+    readonly shape: readonly number[];
+    readonly dtype: DType;
+    readonly strides: readonly number[];
+    readonly size: number;
+    readonly byteSize: number;
+    constructor(buffer: GPUBuffer, shape: number[], dtype: DType, strides?: number[]);
+    /**
+     * Get the number of dimensions
+     */
+    get ndim(): number;
+    /**
+     * Check if tensor is contiguous (row-major)
+     */
+    get isContiguous(): boolean;
+    /**
+     * Get a view of the tensor with new shape (must have same size)
+     */
     reshape(newShape: number[]): Tensor;
+    /**
+     * Destroy the underlying GPU buffer
+     */
     dispose(): void;
+    /**
+     * Get string representation
+     */
+    toString(): string;
 }

package/dist/core/types.d.ts

@@ -0,0 +1,27 @@
+/**
+ * Core type definitions for WebInfer
+ */
+/**
+ * Data type for tensors
+ */
+export type DType = 'float32' | 'float16';
+/**
+ * Position encoding modes for attention
+ */
+export declare enum PosEncodingMode {
+    NONE = 0,
+    ROPE_LLAMA = 1,
+    ALIBI = 2
+}
+/**
+ * Get the byte size for a given dtype
+ */
+export declare function dtypeByteSize(dtype: DType): number;
+/**
+ * Get the GPUTextureFormat for a given dtype
+ */
+export declare function dtypeToGPUFormat(dtype: DType): string;
+/**
+ * Get WGSL type string for a given dtype
+ */
+export declare function dtypeToWGSL(dtype: DType): 'f32' | 'f16';

package/dist/decode/index.d.ts

@@ -0,0 +1,65 @@
+/**
+ * Decode attention operations
+ */
+import type { WebInferContext } from '../core/context.ts';
+import type { Tensor } from '../core/tensor.ts';
+import type { PagedKvCache } from '../core/paged-kv-cache.ts';
+import { PosEncodingMode } from '../core/types.ts';
+/**
+ * Single decode with KV cache
+ *
+ * Performs attention for a single token (decode phase) with KV cache.
+ * Uses flash attention with online softmax to avoid materializing full attention matrix.
+ *
+ * @param ctx WebInfer context
+ * @param q Query tensor of shape [num_qo_heads, head_dim]
+ * @param k_cache Key cache tensor of shape [seq_len, num_kv_heads, head_dim]
+ * @param v_cache Value cache tensor of shape [seq_len, num_kv_heads, head_dim]
+ * @param pos_encoding_mode Position encoding mode (default: NONE)
+ * @param sm_scale Softmax scale (default: 1/sqrt(head_dim))
+ * @param rope_scale RoPE scale (default: 1.0) - not yet implemented
+ * @param rope_theta RoPE theta (default: 10000.0) - not yet implemented
+ * @returns Output tensor of shape [num_qo_heads, head_dim]
+ */
+export declare function single_decode_with_kv_cache(ctx: WebInferContext, q: Tensor, k_cache: Tensor, v_cache: Tensor, pos_encoding_mode?: PosEncodingMode, sm_scale?: number, rope_scale?: number, rope_theta?: number): Promise<Tensor>;
+/**
+ * Batched decode with paged KV cache wrapper
+ *
+ * FlashInfer-compatible wrapper for batched decode attention with paged KV cache.
+ */
+export declare class BatchDecodeWithPagedKVCacheWrapper {
+    private ctx;
+    private planned;
+    private num_qo_heads;
+    private num_kv_heads;
+    private head_dim;
+    private page_size;
+    private batch_size;
+    private sm_scale;
+    private pos_encoding_mode;
+    constructor(ctx: WebInferContext);
+    /**
+     * Plan the batched decode operation
+     *
+     * @param indptr Indirection pointer [batch_size + 1]
+     * @param indices Page indices [nnz_pages]
+     * @param last_page_len Last page lengths [batch_size]
+     * @param num_qo_heads Number of query heads
+     * @param num_kv_heads Number of KV heads
+     * @param head_dim Head dimension
+     * @param page_size Page size
+     * @param pos_encoding_mode Position encoding mode (default: NONE)
+     */
+    plan(indptr: Tensor, indices: Tensor, last_page_len: Tensor, num_qo_heads: number, num_kv_heads: number, head_dim: number, page_size: number, pos_encoding_mode?: PosEncodingMode): void;
+    /**
+     * Run the batched decode operation
+     *
+     * @param q Query tensor [batch_size, num_qo_heads, head_dim]
+     * @param paged_kv_cache Paged KV cache
+     * @param indptr Indirection pointer [batch_size + 1]
+     * @param indices Page indices [nnz_pages]
+     * @param last_page_len Last page lengths [batch_size]
+     * @returns Output tensor [batch_size, num_qo_heads, head_dim]
+     */
+    run(q: Tensor, paged_kv_cache: PagedKvCache, indptr: Tensor, indices: Tensor, last_page_len: Tensor): Promise<Tensor>;
+}

package/dist/gemm/index.d.ts

@@ -0,0 +1,25 @@
+/**
+ * GEMM (General Matrix Multiplication) operations
+ */
+import type { WebInferContext } from '../core/context.ts';
+import type { Tensor } from '../core/tensor.ts';
+/**
+ * Batched matrix multiplication for fp16
+ *
+ * Computes C = A @ B for batched matrices
+ *
+ * @param ctx WebInfer context
+ * @param a Input tensor A [B, M, K]
+ * @param b Input tensor B [B, K, N]
+ * @returns Output tensor C [B, M, N]
+ */
+export declare function bmm_fp16(ctx: WebInferContext, a: Tensor, b: Tensor): Promise<Tensor>;
+/**
+ * Batched matrix multiplication for fp32
+ *
+ * @param ctx WebInfer context
+ * @param a Input tensor A [B, M, K]
+ * @param b Input tensor B [B, K, N]
+ * @returns Output tensor C [B, M, N]
+ */
+export declare function bmm_fp32(ctx: WebInferContext, a: Tensor, b: Tensor): Promise<Tensor>;

package/dist/index.d.ts

@@ -1,22 +1,29 @@
 /**
  * WebInfer - High-performance LLM inference kernels for WebGPU
- * "The cuDNN/FlashInfer of WebGPU"
+ *
+ * WebGPU implementation of FlashInfer APIs for browser-based LLM inference.
  */
-export { WebInferDevice, type DeviceInfo } from "./core/device.ts";
-export { Tensor, type DType } from "./core/tensor.ts";
-export { BufferPool } from "./core/buffer-pool.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 { 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";
+export { WebInferContext } from './core/context.ts';
+export { Tensor } from './core/tensor.ts';
+export type { DType } from './core/types.ts';
+export { PosEncodingMode } from './core/types.ts';
+export type { PagedKvCache } from './core/paged-kv-cache.ts';
+export { createPagedKvCache } from './core/paged-kv-cache.ts';
+export { WebInferContext as create_context } from './core/context.ts';
+import * as norm from './norm/index.ts';
+import * as activation from './activation/index.ts';
+import * as decode from './decode/index.ts';
+import * as prefill from './prefill/index.ts';
+import * as page from './page/index.ts';
+import * as rope from './rope/index.ts';
+import * as sampling from './sampling/index.ts';
+import * as gemm from './gemm/index.ts';
+export { norm, activation, decode, prefill, page, rope, sampling, gemm };
+export declare const rmsnorm: typeof norm.rmsnorm, fused_add_rmsnorm: typeof norm.fused_add_rmsnorm, gemma_rmsnorm: typeof norm.gemma_rmsnorm;
+export declare const silu_and_mul: typeof activation.silu_and_mul, gelu_and_mul: typeof activation.gelu_and_mul;
+export declare const single_decode_with_kv_cache: typeof decode.single_decode_with_kv_cache, BatchDecodeWithPagedKVCacheWrapper: typeof decode.BatchDecodeWithPagedKVCacheWrapper;
+export declare const single_prefill_with_kv_cache: typeof prefill.single_prefill_with_kv_cache, BatchPrefillWithPagedKVCacheWrapper: typeof prefill.BatchPrefillWithPagedKVCacheWrapper;
+export declare const append_paged_kv_cache: typeof page.append_paged_kv_cache;
+export declare const apply_rope_inplace: typeof rope.apply_rope_inplace, apply_llama31_rope_inplace: typeof rope.apply_llama31_rope_inplace;
+export declare const sampling_from_probs: typeof sampling.sampling_from_probs, top_k_sampling_from_probs: typeof sampling.top_k_sampling_from_probs, top_p_sampling_from_probs: typeof sampling.top_p_sampling_from_probs, min_p_sampling_from_probs: typeof sampling.min_p_sampling_from_probs, top_k_top_p_sampling_from_probs: typeof sampling.top_k_top_p_sampling_from_probs;
+export declare const bmm_fp16: typeof gemm.bmm_fp16, bmm_fp32: typeof gemm.bmm_fp32;