import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
import { Tensor } from '@tensorflow/tfjs-core';
import { GPTConfig } from '../models/config';
export interface KVCache {
    k?: Tensor;
    v?: Tensor;
    length: number;
    cumulativeLength: number;
}
export interface AttentionScores {
    meanOfHeads?: boolean;
    attentionOut?: Tensor[];
}
interface AttentionForwardAttributes extends ForwardAttributes {
    attentionScores?: AttentionScores;
    pastKV?: KVCache;
    seed?: number;
}
export interface CausalSelfAttentionConfig {
    useQKNorm?: boolean;
}
export default class CausalSelfAttention extends BaseLayer<AttentionForwardAttributes> {
    private readonly attentionConfig;
    private divisor;
    private index;
    private units;
    private projUnits;
    private ATTN;
    private PROJ;
    constructor(index: number, config: GPTConfig, attentionConfig: CausalSelfAttentionConfig, parent?: BaseLayer);
    protected build(): void;
    private getAttentionScores;
    private getQKV;
    private getOutputProjection;
    private updateCache;
    forward(attr: AttentionForwardAttributes, x: Tensor): Tensor;
}
export {};
