/**
 * JIT compilation module for WebInfer
 *
 * This module provides functionality to compile kernel specifications into
 * executable WGSL shaders at runtime. It follows a similar pattern to
 * FlashInfer's JIT compilation but targets WebGPU instead of CUDA.
 */
/**
 * Kernel specification for JIT compilation.
 * This is the format used by TVM to pass kernel configurations to WebInfer.
 */
export interface KernelSpec {
    /** Type of kernel to compile */
    kernel_type: 'batch_prefill_paged' | 'batch_decode_paged' | 'single_prefill' | 'single_decode' | 'rmsnorm' | 'silu_and_mul' | 'gelu_and_mul' | 'rope' | 'sampling';
    /** Data type for computation */
    dtype: 'float16' | 'float32';
    /** Number of query/output heads */
    num_qo_heads?: number;
    /** Number of key/value heads (for GQA/MQA) */
    num_kv_heads?: number;
    /** Head dimension for query and key */
    qk_head_dim?: number;
    /** Head dimension for value (may differ from qk_head_dim) */
    v_head_dim?: number;
    /** Page size for paged KV cache */
    page_size?: number;
    /** Whether to use causal masking */
    causal?: boolean;
    /** Whether to enable inline RoPE */
    enable_inline_rope?: boolean;
    /** Hidden dimension (for normalization kernels) */
    hidden_dim?: number;
    /** Epsilon for numerical stability (for normalization) */
    eps?: number;
    /** RoPE theta base */
    rope_theta?: number;
    /** RoPE scaling factor */
    rope_scale?: number;
}
/**
 * Information about a binding in the compiled shader
 */
export interface BindingInfo {
    /** Binding index */
    binding: number;
    /** Name of the binding (for documentation) */
    name: string;
    /** Type of binding */
    type: 'storage' | 'storage_read' | 'uniform';
    /** Data type */
    dtype: string;
}
/**
 * Result of JIT compilation
 */
export interface CompiledKernel {
    /** Generated WGSL shader code */
    wgsl: string;
    /** Workgroup size [x, y, z] */
    workgroupSize: [number, number, number];
    /** Binding layout information */
    bindings: BindingInfo[];
    /** Entry point function name */
    entryPoint: string;
    /**
     * Calculate dispatch size based on input parameters
     * @param params - Parameters like batch_size, seq_len, etc.
     * @returns Dispatch size [x, y, z]
     */
    dispatchSize: (params: Record<string, number>) => [number, number, number];
    /** Original kernel specification */
    spec: KernelSpec;
}
/**
 * Plan information returned by plan phase
 */
export interface PlanInfo {
    /** Unique key for this configuration */
    key: string;
    /** Required workspace size in bytes */
    workspaceSize: number;
    /** Compiled kernel reference */
    kernel: CompiledKernel;
    /** Additional configuration for execution */
    config: Record<string, number | boolean>;
}
/**
 * Compile a kernel from specification.
 *
 * This is the main entry point for JIT compilation. It takes a kernel
 * specification and returns a compiled kernel that can be executed.
 *
 * @param spec - Kernel specification
 * @returns Compiled kernel
 */
export declare function compileKernel(spec: KernelSpec): CompiledKernel;
/**
 * Generate a unique key for a kernel specification.
 * Used for caching compiled pipelines.
 */
export declare function getSpecKey(spec: KernelSpec): string;
/**
 * Registry entry for a compiled kernel
 */
export interface CompiledKernelEntry {
    /** The compiled kernel information */
    kernel: CompiledKernel;
    /** The GPU compute pipeline */
    pipeline: GPUComputePipeline;
    /** The bind group layout for this kernel */
    bindGroupLayout: GPUBindGroupLayout;
}
/**
 * Registry of compiled kernels, keyed by spec key
 */
export interface CompiledKernelRegistry {
    [specKey: string]: CompiledKernelEntry;
}
/**
 * Initialize multiple kernels from specifications.
 *
 * This function takes a list of kernel specs and compiles them all,
 * returning a registry that can be used to quickly lookup pipelines
 * at execution time.
 *
 * @param device - The GPU device to create pipelines on
 * @param specs - Array of kernel specifications to compile
 * @returns Promise resolving to a registry of compiled kernels
 */
export declare function initFromSpecs(device: GPUDevice, specs: KernelSpec[]): Promise<CompiledKernelRegistry>;
/**
 * Get a compiled kernel entry from the registry.
 *
 * @param registry - The kernel registry
 * @param spec - The kernel specification to look up
 * @returns The compiled kernel entry, or undefined if not found
 */
export declare function getCompiledKernel(registry: CompiledKernelRegistry, spec: KernelSpec): CompiledKernelEntry | undefined;