/**
 * mamba_model.ts – HybridMambaModel: Mamba-1/2/3 and Attention layer scheduling.
 *
 * Replaces the fixed MambaBlock[] array with a SequenceLayer[] built from a
 * per-layer type schedule.  MambaModel is kept as a backward-compatible alias
 * (all-mamba1 schedule).
 *
 * MBJS binary format:
 *   Version 1 (legacy): [magic][v=1][nParams][numel[]][ f32 data ]
 *   Version 2 (new):    [magic][v=2][nLayers][layerType[]][padding][nParams][numel[]][ f32 data ]
 *     layerType: 0=mamba1, 1=mamba2, 2=mamba3, 3=attention
 */
import type { SequenceLayer, LayerParam, LayerType } from './sequence_layer.js';
import type { Mamba1BlockConfig } from './mamba1_block.js';
import type { Mamba2BlockConfig } from './mamba2_block.js';
import type { Mamba3BlockConfig } from './mamba3_block.js';
import type { AttentionBlockConfig } from './attention_block.js';
export interface LayerSpec {
    type: LayerType;
    config?: Partial<Mamba1BlockConfig | Mamba2BlockConfig | Mamba3BlockConfig | AttentionBlockConfig>;
}
export interface HybridMambaModelConfig {
    vocabSize: number;
    dModel: number;
    numLayers: number;
    /**
     * Per-layer type schedule.  Length must equal numLayers.
     * Defaults to all 'mamba1' (backward-compatible).
     */
    layers?: LayerSpec[];
    defaultMamba1?: Partial<Mamba1BlockConfig>;
    defaultMamba2?: Partial<Mamba2BlockConfig>;
    defaultMamba3?: Partial<Mamba3BlockConfig>;
    defaultAttention?: Partial<AttentionBlockConfig>;
    dState?: number;
    dConv?: number;
    expand?: number;
    nHeads?: number;
    nGroups?: number;
    chunkLen?: number;
    mimoGroup?: number;
    eosId?: number;
    /**
     * Optional deterministic seed for weight initialisation. When set, the
     * embedding table and all block weights are initialised reproducibly — the
     * same seed yields byte-identical initial weights on any machine. When
     * omitted, weights use `Math.random` (non-reproducible) as before.
     */
    seed?: number;
}
/** Legacy Mamba-1-only config (fully backward-compatible). */
export interface MambaModelConfig {
    vocabSize: number;
    dModel: number;
    numLayers: number;
    dState?: number;
    dConv?: number;
    expand?: number;
    eosId?: number;
}
export interface ModelForwardResult {
    logits: Float32Array;
    gpuLogits: GPUBuffer;
    caches: unknown[];
}
export interface SamplingOptions {
    temperature?: number;
    topK?: number;
    topP?: number;
}
export declare class HybridMambaModel {
    device: GPUDevice;
    config: Required<HybridMambaModelConfig>;
    gpuEmbedding: GPUBuffer;
    layers: SequenceLayer[];
    layerSpecs: LayerSpec[];
    gpuFinalNorm: GPUBuffer;
    tiedEmbedding: boolean;
    gpuLMHeadBias: GPUBuffer;
    private _lmHeadPipeline;
    private _rmsnormPipeline;
    private _embedPipeline;
    private _wslaMode;
    constructor(device: GPUDevice, config: HybridMambaModelConfig);
    private _buildLayer;
    embedTokens(tokenIds: number[] | Uint32Array, batch: number, seqLen: number): GPUBuffer;
    forward(tokenIds: number[] | Uint32Array, batch: number, seqLen: number): Promise<ModelForwardResult>;
    /**
     * Produces a single fixed-length embedding vector for a token sequence.
     *
     * Runs the full layer stack plus the final RMSNorm — i.e. the same hidden
     * state the LM head consumes — then mean-pools across sequence positions and
     * L2-normalises the result. The returned vector has length `dModel` and is
     * suitable for cosine-similarity semantic search.
     *
     * Unlike `forward()`, this skips the (expensive) LM-head projection: it only
     * needs the `dModel`-wide hidden state, not `vocabSize` logits.
     *
     * The embedding reflects whatever the model currently knows — an untrained
     * model behaves like a random projection of the token embeddings (still
     * lexically discriminative), and the representation sharpens automatically as
     * the model is adapted/distilled.
     */
    embed(tokenIds: number[] | Uint32Array): Promise<Float32Array>;
    generate(promptIds: number[], maxNewTokens?: number, samplingOpts?: SamplingOptions): Promise<number[]>;
    parameters(): LayerParam[];
    /**
     * The parameters a `learn()` / `adapt()` step is allowed to update.
     *
     * Under WSLA (write-through / narrow adaptation — the Evermind "update ==
     * replace" path) this is ONLY the per-layer narrow subset (the input
     * projections that produce Δ/B/C). The backbone — token embedding, final
     * norm, and every block's `A_log` / conv / output projection — stays FROZEN.
     * Freezing `A_log` is the key stability guarantee: the state-transition
     * decay can never drift into a degenerate regime (state death / no-decay)
     * across repeated adapts, which is what made the model "die after several
     * executions". Outside WSLA every parameter trains (full fine-tune / distill).
     *
     * Param names are namespaced `layer{i}.{name}` to match {@link parameters},
     * so a name-keyed optimizer (see MambaTrainer) reuses the same Adam moments
     * whether it's stepping the full set or just the WSLA subset.
     */
    getTrainableParams(): LayerParam[];
    setWSLAMode(enabled: boolean): void;
    /**
     * Export all parameters to an ArrayBuffer.
     *
     * MBJS v2/v3 format (identical header; only the data encoding differs):
     *   [0..3]   magic    : uint32 = 0x4D424A53
     *   [4..7]   version  : uint32 = 2 (fp32 data) | 3 (fp16 data)
     *   [8..11]  nLayers  : uint32
     *   [12 .. 12+nLayers-1]  layerType[i]: uint8 (0=m1, 1=m2, 2=m3, 3=attn)
     *   aligned to 4 bytes: padding
     *   [next 4] nParams  : uint32
     *   [next 4*nParams]  numel[i]: uint32
     *   [data]  float32 values (v2)  |  float16 values (v3, half the size)
     *
     * Pass `{ fp16: true }` to emit a v3 checkpoint — roughly half the bytes,
     * with a small precision loss that is negligible for SSM weights.
     */
    exportWeights(opts?: {
        fp16?: boolean;
    }): Promise<ArrayBuffer>;
    /**
     * Load parameters from an MBJS v1, v2, or v3 ArrayBuffer.
     *
     * v1: assumes all layers are mamba1 (backward compatible).
     * v2: reads layer type array and validates per-layer parameter counts (fp32 data).
     * v3: identical layout to v2 but the data section is fp16 (dequantised on load).
     */
    loadWeights(buffer: ArrayBuffer): Promise<void>;
    destroy(): void;
}
export declare class MambaModel extends HybridMambaModel {
    constructor(device: GPUDevice, config: MambaModelConfig);
}
//# sourceMappingURL=mamba_model.d.ts.map