/**
 * Prefill 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_prefill_plan()
 */
export interface BatchPrefillPlanOptions {
    /** Number of sequences in the batch */
    batchSize: number;
    /** Total number of query tokens across all sequences */
    totalQoLen: 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 for query and key */
    headDim: number;
    /** Whether to apply causal masking (default: true) */
    causal?: boolean;
}
/**
 * Plan information for batch prefill with paged KV cache.
 * This is returned by batch_prefill_plan() and passed to batch_prefill_run().
 */
export interface BatchPrefillPlanInfo {
    /** 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;
    /** Whether causal masking is enabled */
    causal: boolean;
    /** Batch size */
    batch_size: number;
    /** Total query/output length */
    total_qo_len: number;
    /** Required workspace size in bytes (currently 0) */
    workspaceSize: number;
}
/**
 * Plan batch prefill with paged KV cache.
 *
 * @param options - Configuration options
 * @returns Plan information for execution
 *
 * @example
 * const plan = batch_prefill_plan({
 *   batchSize: 4,
 *   totalQoLen: 2048,
 *   pageSize: 16,
 *   numQoHeads: 32,
 *   numKvHeads: 8,
 *   headDim: 128,
 *   causal: true
 * });
 */
export declare function batch_prefill_plan(options: BatchPrefillPlanOptions): BatchPrefillPlanInfo;
/**
 * Execute batch prefill with paged KV cache.
 *
 * This function executes the attention computation using the plan
 * prepared by batch_prefill_plan().
 *
 * @param ctx - WebInfer context
 * @param planInfo - Plan information from batch_prefill_plan()
 * @param q - Query tensor [total_qo_len, num_qo_heads, head_dim]
 * @param pagedKvCache - Paged KV cache
 * @param qoIndptr - Query indirection pointer [batch_size + 1]
 * @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 [total_qo_len, num_qo_heads, head_dim]
 * @param lse - Log-sum-exp output [total_qo_len, num_qo_heads] (optional)
 */
export declare function batch_prefill_run(ctx: WebInferContext, planInfo: BatchPrefillPlanInfo, q: Tensor, pagedKvCache: PagedKvCache, qoIndptr: Tensor, pageIndptr: Tensor, pageIndices: Tensor, lastPageLen: Tensor, output: Tensor, lse?: Tensor): Promise<void>;
/**
 * Single prefill with KV cache
 *
 * Performs attention for multiple query tokens (prefill phase) with KV cache.
 * Uses flash attention with online softmax and supports causal masking.
 *
 * @param ctx WebInfer context
 * @param q Query tensor of shape [qo_len, num_qo_heads, head_dim]
 * @param k Key tensor of shape [kv_len, num_kv_heads, head_dim]
 * @param v Value tensor of shape [kv_len, num_kv_heads, head_dim]
 * @param output Optional pre-allocated output tensor [qo_len, 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 causal Whether to apply causal masking (default: true)
 * @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 [qo_len, num_qo_heads, head_dim]
 */
export declare function single_prefill_with_kv_cache(ctx: WebInferContext, q: Tensor, k: Tensor, v: Tensor, output?: Tensor, causal?: boolean, pos_encoding_mode?: PosEncodingMode, sm_scale?: number, rope_scale?: number, rope_theta?: number): Promise<Tensor>;
/**
 * Batched prefill with paged KV cache wrapper
 *
 * FlashInfer-compatible wrapper for batched prefill attention with paged KV cache.
 */
export declare class BatchPrefillWithPagedKVCacheWrapper {
    private ctx;
    private planned;
    private num_qo_heads;
    private num_kv_heads;
    private head_dim;
    private page_size;
    private sm_scale;
    private causal;
    private pos_encoding_mode;
    constructor(ctx: WebInferContext);
    /**
     * Plan the batched prefill operation
     *
     * @param qo_indptr Query indirection pointer [batch_size + 1]
     * @param paged_kv_indptr Paged KV indirection pointer [batch_size + 1]
     * @param paged_kv_indices Paged KV indices [nnz_pages]
     * @param paged_kv_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 causal Whether to apply causal masking (default: true)
     * @param pos_encoding_mode Position encoding mode (default: NONE)
     */
    plan(qo_indptr: Tensor, paged_kv_indptr: Tensor, paged_kv_indices: Tensor, paged_kv_last_page_len: Tensor, num_qo_heads: number, num_kv_heads: number, head_dim: number, page_size: number, causal?: boolean, pos_encoding_mode?: PosEncodingMode): void;
    /**
     * Run the batched prefill operation
     *
     * @param q Query tensor [total_qo_len, num_qo_heads, head_dim]
     * @param paged_kv_cache Paged KV cache
     * @param qo_indptr Query indirection pointer [batch_size + 1]
     * @param paged_kv_indptr Paged KV indirection pointer [batch_size + 1]
     * @param paged_kv_indices Paged KV indices [nnz_pages]
     * @param paged_kv_last_page_len Last page lengths [batch_size]
     * @param output Optional pre-allocated output tensor [total_qo_len, num_qo_heads, head_dim].
     *               If provided, results are written to this tensor (zero-copy for TVM integration).
     * @returns Output tensor [total_qo_len, num_qo_heads, head_dim]
     */
    run(q: Tensor, paged_kv_cache: PagedKvCache, qo_indptr: Tensor, paged_kv_indptr: Tensor, paged_kv_indices: Tensor, paged_kv_last_page_len: Tensor, output?: Tensor): Promise<Tensor>;
}