Documentation
¶
Index ¶
- type Cache
- type KVCache
- func (c *KVCache) CanTrim() bool
- func (c *KVCache) Clone() Cache
- func (c *KVCache) Free()
- func (c *KVCache) Len() int
- func (c *KVCache) Offset() int
- func (c *KVCache) State() (*mlx.Array, *mlx.Array)
- func (c *KVCache) Trim(n int) int
- func (c *KVCache) Update(keys, values *mlx.Array) (*mlx.Array, *mlx.Array)
- type RecurrentCache
- func (c *RecurrentCache) Advance(n int)
- func (c *RecurrentCache) CanTrim() bool
- func (c *RecurrentCache) Clone() Cache
- func (c *RecurrentCache) ConvState(batch int, dtype mlx.DType) *mlx.Array
- func (c *RecurrentCache) DeltaState(batch int, dtype mlx.DType) *mlx.Array
- func (c *RecurrentCache) Free()
- func (c *RecurrentCache) Len() int
- func (c *RecurrentCache) Offset() int
- func (c *RecurrentCache) SetConvState(v *mlx.Array)
- func (c *RecurrentCache) SetDeltaState(v *mlx.Array)
- func (c *RecurrentCache) State() (*mlx.Array, *mlx.Array)
- func (c *RecurrentCache) Trim(n int) int
- func (c *RecurrentCache) Update(keys, values *mlx.Array) (*mlx.Array, *mlx.Array)
- type RotatingKVCache
- func (c *RotatingKVCache) CanTrim() bool
- func (c *RotatingKVCache) Clone() Cache
- func (c *RotatingKVCache) Len() int
- func (c *RotatingKVCache) State() (*mlx.Array, *mlx.Array)
- func (c *RotatingKVCache) Trim(n int) int
- func (c *RotatingKVCache) Update(keys, values *mlx.Array) (*mlx.Array, *mlx.Array)
Constants ¶
This section is empty.
Variables ¶
This section is empty.
Functions ¶
This section is empty.
Types ¶
type KVCache ¶
type KVCache struct {
// contains filtered or unexported fields
}
func NewKVCache ¶
func NewKVCache() *KVCache
type RecurrentCache ¶
type RecurrentCache struct {
// contains filtered or unexported fields
}
RecurrentCache stores state for linear-recurrent layers.
Conv state shape: [B, convTail, convDim] Delta state shape: [B, numVHeads, headVDim, headKDim]
func NewRecurrentCache ¶
func NewRecurrentCache(convTail, convDim, numVHeads, headVDim, headKDim int32) *RecurrentCache
func (*RecurrentCache) Advance ¶
func (c *RecurrentCache) Advance(n int)
func (*RecurrentCache) CanTrim ¶
func (c *RecurrentCache) CanTrim() bool
func (*RecurrentCache) Clone ¶
func (c *RecurrentCache) Clone() Cache
func (*RecurrentCache) DeltaState ¶
func (*RecurrentCache) Free ¶
func (c *RecurrentCache) Free()
func (*RecurrentCache) Len ¶
func (c *RecurrentCache) Len() int
func (*RecurrentCache) Offset ¶
func (c *RecurrentCache) Offset() int
func (*RecurrentCache) SetConvState ¶
func (c *RecurrentCache) SetConvState(v *mlx.Array)
func (*RecurrentCache) SetDeltaState ¶
func (c *RecurrentCache) SetDeltaState(v *mlx.Array)
func (*RecurrentCache) Trim ¶
func (c *RecurrentCache) Trim(n int) int
type RotatingKVCache ¶
type RotatingKVCache struct {
*KVCache
// contains filtered or unexported fields
}
RotatingKVCache implements sliding window attention with bounded memory
func NewRotatingKVCache ¶
func NewRotatingKVCache(maxSize int) *RotatingKVCache
func (*RotatingKVCache) CanTrim ¶
func (c *RotatingKVCache) CanTrim() bool
func (*RotatingKVCache) Clone ¶
func (c *RotatingKVCache) Clone() Cache
func (*RotatingKVCache) Len ¶
func (c *RotatingKVCache) Len() int
func (*RotatingKVCache) Trim ¶
func (c *RotatingKVCache) Trim(n int) int
Click to show internal directories.
Click to hide internal directories.