/**
 * 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>;
}