webinfer 0.0.4 → 0.0.5

package/README.md

@@ -1,59 +1,88 @@
 # WebInfer
 
-High-performance LLM inference kernels for WebGPU.
+[![npm version](https://img.shields.io/npm/v/webinfer.svg)](https://www.npmjs.com/package/webinfer)
+[![License](https://img.shields.io/badge/license-Apache--2.0-blue.svg)](LICENSE)
 
-WebGPU implementation of FlashInfer APIs for browser-based LLM inference.
+High-performance LLM inference kernels for WebGPU. A browser-native implementation of [FlashInfer](https://github.com/flashinfer-ai/flashinfer) APIs.
 
-## Features
+## Why WebInfer?
 
-- **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
+Running LLMs in the browser requires efficient GPU kernels. WebInfer brings FlashInfer's battle-tested attention mechanisms to WebGPU:
 
-## Install
+- **Flash Attention**: Online softmax with O(1) memory for long sequences
+- **Subgroup Optimizations**: 2-4x faster reductions on supported hardware
+- **JIT Compilation**: Runtime kernel generation with pipeline caching
+- **Dynamic Tile Sizes**: Auto-tuned tile dimensions based on workload and hardware
+
+## Installation
 
 ```bash
 npm install webinfer
 ```
 
-## API Status
+## Quick Start
+
+```typescript
+import * as webinfer from 'webinfer';
 
-**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
+// Initialize context
+const ctx = await webinfer.WebInferContext.create();
 
-**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
+// Run RMSNorm
+const output = ctx.norm.rmsnorm(input, weight, epsilon);
+
+// Flash attention decode
+const attnOutput = ctx.decode.single_decode_with_kv_cache(
+  query,      // [num_heads, head_dim]
+  kv_cache,   // [seq_len, 2, num_kv_heads, head_dim]
+);
+```
 
-**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
+## Browser Compatibility
 
-**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)
+WebInfer requires WebGPU support:
 
-See [dev/design.md](dev/design.md) for full design and implementation plan.
+| Browser | Status |
+|---------|--------|
+| Chrome 113+ | Supported |
+| Edge 113+ | Supported |
+| Firefox Nightly | Behind flag |
+| Safari 18+ | Supported |
+
+Check for WebGPU support:
+
+```typescript
+if (!navigator.gpu) {
+  console.error('WebGPU not supported');
+}
+```
 
-## Release
+## Development
 
 ```bash
+# Run tests
+bun test
+
+# Build
 bun run build
-npm publish
+```
+
+## Acknowledgments
+
+- [FlashInfer](https://github.com/flashinfer-ai/flashinfer) - The original CUDA implementation
+- [WebLLM](https://github.com/mlc-ai/web-llm) - Browser LLM runtime
+
+## Citation
+
+If you use WebInfer in your research, please cite:
+
+```bibtex
+@software{webinfer2025,
+  author       = {Guan-Ming, Chiu},
+  title        = {WebInfer: High-Performance LLM Inference Kernels for WebGPU},
+  year         = {2026},
+  url          = {https://github.com/guan404ming/webinfer}
+}
 ```
 
 ## License

package/dist/core/context.d.ts

@@ -29,7 +29,17 @@ export declare class WebInferContext {
     private disposed;
     private constructor();
     /**
-     * Create a new WebInferContext
+     * Create a WebInferContext from an existing GPUDevice.
+     * Use this when integrating with TVM or other WebGPU frameworks
+     * that already have a GPUDevice to enable zero-copy buffer sharing.
+     *
+     * @param device - The existing GPUDevice to use
+     * @param featureOverrides - Optional overrides for detected features
+     * @returns A new WebInferContext using the provided device
+     */
+    static fromDevice(device: GPUDevice, featureOverrides?: Partial<DeviceFeatures>): WebInferContext;
+    /**
+     * Create a new WebInferContext with its own GPUDevice
      */
     static create(): Promise<WebInferContext>;
     /**

package/dist/core/tensor.d.ts

@@ -21,6 +21,19 @@ export declare class Tensor {
     readonly size: number;
     readonly byteSize: number;
     constructor(buffer: GPUBuffer, shape: number[], dtype: DType, strides?: number[]);
+    /**
+     * Create a Tensor from an existing GPUBuffer (zero-copy).
+     *
+     * This is useful for TVM integration where the buffer is already allocated
+     * by TVM's WebGPU runtime and we want to wrap it without copying data.
+     *
+     * @param buffer - The existing GPUBuffer to wrap
+     * @param shape - The shape of the tensor
+     * @param dtype - The data type of the tensor
+     * @param strides - Optional custom strides (default: row-major)
+     * @returns A new Tensor wrapping the provided buffer
+     */
+    static fromBuffer(buffer: GPUBuffer, shape: number[], dtype: DType, strides?: number[]): Tensor;
     /**
      * Get the number of dimensions
      */

package/dist/decode/index.d.ts

@@ -5,6 +5,76 @@ 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';
+/**
+ * Options for batch_decode_plan()
+ */
+export interface BatchDecodePlanOptions {
+    /** Number of sequences in the batch */
+    batchSize: number;
+    /** Size of each page in the paged KV cache */
+    pageSize: number;
+    /** Number of query/output heads */
+    numQoHeads: number;
+    /** Number of key/value heads */
+    numKvHeads: number;
+    /** Head dimension */
+    headDim: number;
+}
+/**
+ * Plan information for batch decode with paged KV cache.
+ * This is returned by batch_decode_plan() and passed to batch_decode_run().
+ */
+export interface BatchDecodePlanInfo {
+    /** Unique key for this configuration */
+    key: string;
+    /** Number of query/output heads */
+    num_qo_heads: number;
+    /** Number of key/value heads */
+    num_kv_heads: number;
+    /** Head dimension */
+    head_dim: number;
+    /** Page size */
+    page_size: number;
+    /** Softmax scale */
+    sm_scale: number;
+    /** Batch size */
+    batch_size: number;
+    /** Required workspace size in bytes (currently 0) */
+    workspaceSize: number;
+}
+/**
+ * Plan batch decode with paged KV cache.
+ *
+ * @param options - Configuration options
+ * @returns Plan information for execution
+ *
+ * @example
+ * const plan = batch_decode_plan({
+ *   batchSize: 4,
+ *   pageSize: 16,
+ *   numQoHeads: 32,
+ *   numKvHeads: 8,
+ *   headDim: 128
+ * });
+ */
+export declare function batch_decode_plan(options: BatchDecodePlanOptions): BatchDecodePlanInfo;
+/**
+ * Execute batch decode with paged KV cache.
+ *
+ * This function executes the attention computation using the plan
+ * prepared by batch_decode_plan().
+ *
+ * @param ctx - WebInfer context
+ * @param planInfo - Plan information from batch_decode_plan()
+ * @param q - Query tensor [batch_size, num_qo_heads, head_dim]
+ * @param pagedKvCache - Paged KV cache
+ * @param pageIndptr - Page indirection pointer [batch_size + 1]
+ * @param pageIndices - Page indices [nnz_pages]
+ * @param lastPageLen - Last page lengths [batch_size]
+ * @param output - Output tensor [batch_size, num_qo_heads, head_dim]
+ * @param lse - Log-sum-exp output [batch_size, num_qo_heads] (optional)
+ */
+export declare function batch_decode_run(ctx: WebInferContext, planInfo: BatchDecodePlanInfo, q: Tensor, pagedKvCache: PagedKvCache, pageIndptr: Tensor, pageIndices: Tensor, lastPageLen: Tensor, output: Tensor, lse?: Tensor): Promise<void>;
 /**
  * Single decode with KV cache
  *
@@ -15,13 +85,16 @@ import { PosEncodingMode } from '../core/types.ts';
  * @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 output Optional pre-allocated output tensor [num_qo_heads, head_dim].
+ *               If provided, results are written to this tensor (zero-copy for TVM integration).
+ *               If not provided, a new tensor is created and returned.
  * @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>;
+export declare function single_decode_with_kv_cache(ctx: WebInferContext, q: Tensor, k_cache: Tensor, v_cache: Tensor, output?: Tensor, pos_encoding_mode?: PosEncodingMode, sm_scale?: number, rope_scale?: number, rope_theta?: number): Promise<Tensor>;
 /**
  * Batched decode with paged KV cache wrapper
  *
@@ -59,7 +132,9 @@ export declare class BatchDecodeWithPagedKVCacheWrapper {
      * @param indptr Indirection pointer [batch_size + 1]
      * @param indices Page indices [nnz_pages]
      * @param last_page_len Last page lengths [batch_size]
+     * @param output Optional pre-allocated output tensor [batch_size, num_qo_heads, head_dim].
+     *               If provided, results are written to this tensor (zero-copy for TVM integration).
      * @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>;
+    run(q: Tensor, paged_kv_cache: PagedKvCache, indptr: Tensor, indices: Tensor, last_page_len: Tensor, output?: Tensor): Promise<Tensor>;
 }

package/dist/gemm/index.d.ts

@@ -11,15 +11,17 @@ import type { Tensor } from '../core/tensor.ts';
  * @param ctx WebInfer context
  * @param a Input tensor A [B, M, K]
  * @param b Input tensor B [B, K, N]
+ * @param output Optional output tensor C [B, M, N] for in-place operation
  * @returns Output tensor C [B, M, N]
  */
-export declare function bmm_fp16(ctx: WebInferContext, a: Tensor, b: Tensor): Promise<Tensor>;
+export declare function bmm_fp16(ctx: WebInferContext, a: Tensor, b: Tensor, output?: 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]
+ * @param output Optional output tensor C [B, M, N] for in-place operation
  * @returns Output tensor C [B, M, N]
  */
-export declare function bmm_fp32(ctx: WebInferContext, a: Tensor, b: Tensor): Promise<Tensor>;
+export declare function bmm_fp32(ctx: WebInferContext, a: Tensor, b: Tensor, output?: Tensor): Promise<Tensor>;

package/dist/index.d.ts

@@ -18,7 +18,10 @@ 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 };
+import * as jit from './jit/index.ts';
+export { norm, activation, decode, prefill, page, rope, sampling, gemm, jit };
+export type { KernelSpec, CompiledKernel, BindingInfo, PlanInfo, CompiledKernelEntry, CompiledKernelRegistry, } from './jit/index.ts';
+export { compileKernel, getSpecKey, initFromSpecs, getCompiledKernel, } from './jit/index.ts';
 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;