Class MiniMaxIndexCache

java.lang.Object

smile.llm.model.minimax.MiniMaxIndexCache

All Implemented Interfaces:

AutoCloseable

public class MiniMaxIndexCache
extends Object
implements AutoCloseable

Key-only cache of MSA index keys, one [numKeyHeads, indexDim] vector per token and sparse layer. It is addressed by the physical slots of the main KvCachePool, so allocation, radix prefix reuse and eviction need no separate bookkeeping.

Constructor Summary

Constructors

Constructor	Description
MiniMaxIndexCache(int numSparseLayers, int numSlots, int numKeyHeads, int indexDim, Device device, ScalarType dtype)	Constructor.

Method Summary

Modifier and Type	Method	Description
void	close()
Tensor	get(int ordinal, long[] slots)	Reads index keys.
Tensor	get(int ordinal, Tensor idx)	get(int, long[]) with the slots already on the device.
Tensor	layer(int ordinal)	One layer of the cache [numSlots, numKeyHeads, indexDim] as a view; the caller closes it.
int	numKeyHeads()	Returns the number of index-key heads per token.
void	put(int ordinal, long[] slots, Tensor keys)	Writes index keys.
void	put(int ordinal, Tensor idx, Tensor keys)	put(int, long[], Tensor) with the slots already on the device.

Methods inherited from class Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Constructor Details

MiniMaxIndexCache

public MiniMaxIndexCache(int numSparseLayers,
 int numSlots,
 int numKeyHeads,
 int indexDim,
 Device device,
 ScalarType dtype)

Constructor.

Parameters:

numSparseLayers - number of MSA layers.

numSlots - token slots, equal to the KV pool's.

numKeyHeads - index-key heads per token.

indexDim - index head dimension.

device - device of the KV pool.

dtype - storage dtype.

Method Details

put

public void put(int ordinal,
 long[] slots,
 Tensor keys)

Writes index keys.

Parameters:

ordinal - sparse-layer ordinal.

slots - destination slots, length N.

keys - index keys, any shape holding N * numKeyHeads * indexDim elements.

put

public void put(int ordinal,
 Tensor idx,
 Tensor keys)

put(int, long[], Tensor) with the slots already on the device.

Parameters:

ordinal - sparse-layer ordinal.

idx - device tensor of destination slots, length N.

keys - index keys, any shape holding N * numKeyHeads * indexDim elements.

get

public Tensor get(int ordinal,
 Tensor idx)

get(int, long[]) with the slots already on the device.

Parameters:

ordinal - sparse-layer ordinal.

idx - device tensor of source slots, length N.

Returns:

keys [N, numKeyHeads, indexDim]; caller owns the tensor.

layer

public Tensor layer(int ordinal)

One layer of the cache [numSlots, numKeyHeads, indexDim] as a view; the caller closes it.

Parameters:

ordinal - sparse-layer ordinal.

Returns:

the layer view tensor.

numKeyHeads

public int numKeyHeads()

Returns the number of index-key heads per token.

Returns:

index-key heads per token.

get

public Tensor get(int ordinal,
 long[] slots)

Reads index keys.

Parameters:

ordinal - sparse-layer ordinal.

slots - source slots, length N.

Returns:

keys [N, numKeyHeads, indexDim]; caller owns the tensor.

close

public void close()

Specified by:

close in interface AutoCloseable