Class LlamaModel

java.lang.Object

smile.deep.layer.LayerBlock

smile.llm.model.llama.LlamaModel

All Implemented Interfaces:

Function<Tensor,Tensor>, Layer

public class LlamaModel
extends LayerBlock

Llama dense decoder: token embeddings, stacked LlamaBlocks, final RMSNorm, and the output projection.

Construct on CPU, load weights, then call to(Device) to place parameters and the RoPE frequency table. Install the KV pool afterward via setKvCachePool(KvCachePool).

Field Summary

Fields inherited from class LayerBlock

device, dtype, module

Constructor Summary

Constructors

Constructor	Description
LlamaModel(int dim, int numLayers, int numHeads, int numKvHeads, int vocabSize, Integer intermediateSize, int multipleOf, Double ffnDimMultiplier, double normEps, double ropeTheta, boolean scaledRope, int maxSeqLen)	Constructs the module graph on CPU with a host-side RoPE table.

Method Summary

Modifier and Type	Method	Description
Tensor	forward(Tensor tokens)	Forward propagation (or forward pass) through the layer.
Tensor	forward(Tensor tokens, int startPos)	Forward pass through the model.
Tensor	forward(Tensor tokens, int[] positions)	Decode forward with per-row absolute positions (seqLen must be 1).
KvCachePool	kvCachePool()	Returns the shared KV cache pool.
List<LlamaBlock>	layers()	Returns the decoder blocks.
int	numLayers()	Returns the number of transformer blocks.
void	setKvCachePool(KvCachePool pool)	Replaces the shared KV cache pool on every attention layer.
void	setKvCachePool(KvCachePool pool, boolean closePrevious)	Replaces the shared KV cache pool on every attention layer.
LlamaModel	to(Device device)	Moves parameters and the RoPE frequency table to device.
LlamaModel	to(Device device, ScalarType dtype)	Moves parameters and the RoPE frequency table to device / dtype.

Methods inherited from class LayerBlock

add, add, device, dtype, eval, isTraining, load, loadStateDict, module, save, setRequiresGrad, toString, train

Methods inherited from class Object

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

Methods inherited from interface Function

andThen, compose

Methods inherited from interface Layer

apply, name

Constructor Details

LlamaModel

public LlamaModel(int dim,
 int numLayers,
 int numHeads,
 int numKvHeads,
 int vocabSize,
 Integer intermediateSize,
 int multipleOf,
 Double ffnDimMultiplier,
 double normEps,
 double ropeTheta,
 boolean scaledRope,
 int maxSeqLen)

Constructs the module graph on CPU with a host-side RoPE table. Call to(Device) after weight load; then setKvCachePool(KvCachePool).

Parameters:

dim - token embedding dimension.

numLayers - number of transformer blocks.

numHeads - number of query heads.

numKvHeads - number of key/value heads.

vocabSize - vocabulary size.

intermediateSize - explicit FFN hidden size, or null.

multipleOf - FFN rounding multiple when deriving size.

ffnDimMultiplier - optional FFN dim multiplier.

normEps - RMSNorm epsilon.

ropeTheta - RoPE theta.

scaledRope - whether to use Llama-3 scaled RoPE.

maxSeqLen - maximum sequence length (RoPE table uses 2 * maxSeqLen).

Method Details

to

public LlamaModel to(Device device)

Moves parameters and the RoPE frequency table to device.

Specified by:

to in interface Layer

Overrides:

to in class LayerBlock

Parameters:

device - the compute device.

Returns:

this layer.

to

public LlamaModel to(Device device,
 ScalarType dtype)

Moves parameters and the RoPE frequency table to device / dtype. RoPE cis stays complex64 (device move only); casting it to a real dtype would drop the imaginary part and break rotary embeddings.

Specified by:

to in interface Layer

Overrides:

to in class LayerBlock

Parameters:

device - the compute device.

dtype - the data type.

Returns:

this layer.

kvCachePool

public KvCachePool kvCachePool()

Returns the shared KV cache pool.

Returns:

the KV cache pool.

numLayers

public int numLayers()

Returns the number of transformer blocks.

Returns:

number of transformer blocks.

layers

public List<LlamaBlock> layers()

Returns the decoder blocks.

Returns:

decoder blocks (mutable list owned by this model).

setKvCachePool

public void setKvCachePool(KvCachePool pool)

Replaces the shared KV cache pool on every attention layer. Closes the previous pool. Intended to be called once after weight loading so the pool can be sized from residual device memory.

Parameters:

pool - the new pool (must not be null).

setKvCachePool

public void setKvCachePool(KvCachePool pool,
 boolean closePrevious)

Replaces the shared KV cache pool on every attention layer.

Parameters:

pool - the new pool (must not be null).

closePrevious - when true, closes the previously installed pool.

forward

public Tensor forward(Tensor tokens,
 int startPos)

Forward pass through the model.

Parameters:

tokens - the input token indices.

startPos - the starting position for attention caching.

Returns:

the output tensor.

forward

public Tensor forward(Tensor tokens,
 int[] positions)

Decode forward with per-row absolute positions (seqLen must be 1).

Parameters:

tokens - token ids [B, 1].

positions - write position per batch row.

Returns:

logits [B, 1, V].

forward

public Tensor forward(Tensor tokens)

Description copied from interface: Layer

Forward propagation (or forward pass) through the layer.

Parameters:

tokens - the input tensor.

Returns:

the output tensor.