webinfer 0.0.3 → 0.0.5

package/README.md

@@ -1,44 +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)
 
-## Install
+High-performance LLM inference kernels for WebGPU. A browser-native implementation of [FlashInfer](https://github.com/flashinfer-ai/flashinfer) APIs.
+
+## Why WebInfer?
+
+Running LLMs in the browser requires efficient GPU kernels. WebInfer brings FlashInfer's battle-tested attention mechanisms to WebGPU:
+
+- **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
 ```
 
-## Usage
+## Quick Start
 
 ```typescript
-import { WebInferDevice, attention } from 'webinfer';
+import * as webinfer from 'webinfer';
 
-const device = await WebInferDevice.create();
+// Initialize context
+const ctx = await webinfer.WebInferContext.create();
 
-// 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);
+// Run RMSNorm
+const output = ctx.norm.rmsnorm(input, weight, epsilon);
 
-const output = await attention(device, { q, k, v });
+// 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]
+);
 ```
 
-## API
+## Browser Compatibility
 
-| 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` |
+WebInfer requires WebGPU support:
+
+| 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/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,70 @@
+/**
+ * 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 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>;
+    /**
+     * 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,57 @@
 /**
- * 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[]);
+    /**
+     * 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
+     */
+    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,140 @@
+/**
+ * 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';
+/**
+ * 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
+ *
+ * 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 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, output?: 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]
+     * @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, output?: Tensor): Promise<Tensor>;
+}

package/dist/gemm/index.d.ts

@@ -0,0 +1,27 @@
+/**
+ * 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]
+ * @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, 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, output?: Tensor): Promise<Tensor>;

package/dist/index.d.ts

@@ -1,24 +1,32 @@
 /**
  * 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 { 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, 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";
+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';
+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;
+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;