The forward entry with one level of blocks expanded. Every edge carries the tensor type the compiler inferred at that point, in the model's own generics.

Entries

Entry	Signature
forward	forward<B: Dim, S: Dim>(tokens: Tensor[B, S; i32]) -> Tensor[B, S, Vocab; T]
next_token	next_token<B: Dim, S: Dim>(tokens: Tensor[B, S; i32]) -> Tensor[B, Vocab; T]
decode	decode(token: Tensor[Batch, 1; i32], pos: i32) -> Tensor[Batch, Vocab; T]
prefill	prefill<S: Dim>(tokens: Tensor[Batch, S; i32], pos: i32) -> Tensor[Batch, Vocab; T]
decode_rows	decode_rows(tokens: Tensor[Batch, 1; i32], positions: Tensor[Batch; i32]) -> Tensor[Batch, Vocab; T]
prefill_slots	prefill_slots<M: Dim, S: Dim>(tokens: Tensor[M, S; i32], slots: Tensor[M; i32], lengths: Tensor[M; i32]) -> Tensor[M, Vocab; T]
loss_packed	loss_packed<P: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], targets: Tensor[P; i64], weights: Tensor[P; f32]) -> f32
log_probs_packed	log_probs_packed<P: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], targets: Tensor[P; i64]) -> Tensor[P; f32]
hidden_packed	hidden_packed<P: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32]) -> Tensor[P, H; T]
prefill_packed	prefill_packed<P: Dim>(tokens: Tensor[P; i32], rows: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], last: Tensor[Batch; i32]) -> Tensor[Batch, Vocab; T]
step_packed	step_packed<P: Dim>(tokens: Tensor[P; i32], rows: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], last: Tensor[Batch; i32], step_tokens: Tensor[Batch, 1; i32], step_positions: Tensor[Batch; i32]) -> Tensor[2 * Batch, Vocab; T]
decode_paged	decode_paged<Rows: Dim, Pages: Dim>(tokens: Tensor[Rows, 1; i32], positions: Tensor[Rows; i32], table: Tensor[Rows, Pages; i32]) -> Tensor[Rows, Vocab; T]
prefill_paged	prefill_paged<P: Dim, Rows: Dim, Pages: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], rows: Tensor[P; i32], slots: Tensor[P; i32], last: Tensor[Rows; i32], table: Tensor[Rows, Pages; i32]) -> Tensor[Rows, Vocab; T]
step_paged	step_paged<P: Dim, Rows: Dim, Pages: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], rows: Tensor[P; i32], slots: Tensor[P; i32], last: Tensor[Rows; i32], step_tokens: Tensor[Rows, 1; i32], step_positions: Tensor[Rows; i32], table: Tensor[Rows, Pages; i32]) -> Tensor[2 * Rows, Vocab; T]
generate	generate<Steps: Dim>(token: Tensor[Batch, 1; i32], pos: i32) -> Tensor[Batch, Steps; i32]
sample	sample<Steps: Dim>(token: Tensor[Batch, 1; i32], pos: i32, key: Tensor[2; i64], temperature: f32) -> Tensor[Batch, Steps; i32]
generate_until	generate_until<MaxNew: Dim>(token: Tensor[Batch, 1; i32], pos: i32, eos: i32) -> (Tensor[Batch, MaxNew; i32], i32)

Generics

The root block's generics as this checkpoint binds them.

Vocab	128256
H	4096
Heads	32
KvHeads	8
Inner	14336
Layers	32
Batch	1
MaxSeq	8192
T	bf16

Blocks

Every block of the program with its members and functions, as linnet inspect prints them.

DecoderLayer

llama::DecoderLayer<H: Dim, Heads: Dim, KvHeads: Dim, Inner: Dim, Batch: Dim, MaxSeq: Dim, T: Float, Shards: Dim = 1, PageSize: Dim = 64>
  sub attention_norm: RmsNorm<H, T>
  sub attention: GroupedQueryAttention<H, Heads, KvHeads, Batch, MaxSeq, T, Shards, PageSize>
  sub mlp_norm: RmsNorm<H, T>
  sub mlp: SwiGlu<H, Inner / Shards, T>
  pub fn forward<B: Dim, S: Dim>(x: Tensor[B, S, H; T]) -> Tensor[B, S, H; T]
  pub fn decode(x: Tensor[Batch, 1, H; T], pos: i32) -> Tensor[Batch, 1, H; T]
  pub fn prefill<S: Dim>(x: Tensor[Batch, S, H; T], pos: i32) -> Tensor[Batch, S, H; T]
  pub fn decode_rows(x: Tensor[Batch, 1, H; T], positions: Tensor[Batch; i32]) -> Tensor[Batch, 1, H; T]
  pub fn prefill_slots<M: Dim, S: Dim>(x: Tensor[M, S, H; T], slots: Tensor[M; i32]) -> Tensor[M, S, H; T]
  pub fn prefill_packed<P: Dim>(x: Tensor[1, P, H; T], rows: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32]) -> Tensor[1, P, H; T]
  pub fn packed<P: Dim>(x: Tensor[1, P, H; T], positions: Tensor[P; i32], segments: Tensor[P; i32]) -> Tensor[1, P, H; T]
  pub fn step_packed<P: Dim>(x: Tensor[1, Batch + P, H; T], rows: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], step_positions: Tensor[Batch; i32]) -> Tensor[1, Batch + P, H; T]
  pub fn decode_paged<Rows: Dim, Pages: Dim>(x: Tensor[Rows, 1, H; T], positions: Tensor[Rows; i32], table: Tensor[Rows, Pages; i32]) -> Tensor[Rows, 1, H; T]
  pub fn prefill_paged<P: Dim, Rows: Dim, Pages: Dim>(x: Tensor[1, P, H; T], positions: Tensor[P; i32], rows: Tensor[P; i32], slots: Tensor[P; i32], table: Tensor[Rows, Pages; i32]) -> Tensor[1, P, H; T]
  pub fn step_paged<P: Dim, Rows: Dim, Pages: Dim>(x: Tensor[1, P + Rows, H; T], positions: Tensor[P; i32], rows: Tensor[P; i32], slots: Tensor[P; i32], step_positions: Tensor[Rows; i32], table: Tensor[Rows, Pages; i32]) -> Tensor[1, P + Rows, H; T]

Embedding

std.nn.embedding::Embedding<Vocab: Dim, H: Dim, T: Float = bf16>
  param weight: Tensor[Vocab, H; T]
  pub fn forward<*S: Shape>(ids: Tensor[*S; i32]) -> Tensor[*S, H; T]

GroupedQueryAttention

llama.attention::GroupedQueryAttention<H: Dim, Heads: Dim, KvHeads: Dim, Batch: Dim, MaxSeq: Dim, T: Float, Shards: Dim = 1, PageSize: Dim = 64>
  sub q_proj: Linear<H, (Heads / Shards) * (H / Heads), T>
  sub k_proj: Linear<H, (KvHeads / Shards) * (H / Heads), T>
  sub v_proj: Linear<H, (KvHeads / Shards) * (H / Heads), T>
  sub o_proj: Linear<(Heads / Shards) * (H / Heads), H, T>
  state cache_k: Tensor[Batch, KvHeads / Shards, MaxSeq, H / Heads; T]
  state cache_v: Tensor[Batch, KvHeads / Shards, MaxSeq, H / Heads; T]
  pub fn forward<B: Dim, S: Dim>(x: Tensor[B, S, H; T]) -> Tensor[B, S, H; T]
  pub fn decode(x: Tensor[Batch, 1, H; T], pos: i32) -> Tensor[Batch, 1, H; T]
  pub fn prefill<S: Dim>(x: Tensor[Batch, S, H; T], pos: i32) -> Tensor[Batch, S, H; T]
  pub fn decode_rows(x: Tensor[Batch, 1, H; T], positions: Tensor[Batch; i32]) -> Tensor[Batch, 1, H; T]
  pub fn prefill_slots<M: Dim, S: Dim>(x: Tensor[M, S, H; T], slots: Tensor[M; i32]) -> Tensor[M, S, H; T]
  pub fn prefill_packed<P: Dim>(x: Tensor[1, P, H; T], rows: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32]) -> Tensor[1, P, H; T]
  pub fn packed<P: Dim>(x: Tensor[1, P, H; T], positions: Tensor[P; i32], segments: Tensor[P; i32]) -> Tensor[1, P, H; T]
  pub fn step_packed<P: Dim>(x: Tensor[1, Batch + P, H; T], rows: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], step_positions: Tensor[Batch; i32]) -> Tensor[1, Batch + P, H; T]
  pub fn decode_paged<Rows: Dim, Pages: Dim>(x: Tensor[Rows, 1, H; T], positions: Tensor[Rows; i32], table: Tensor[Rows, Pages; i32]) -> Tensor[Rows, 1, H; T]
  pub fn prefill_paged<P: Dim, Rows: Dim, Pages: Dim>(x: Tensor[1, P, H; T], positions: Tensor[P; i32], rows: Tensor[P; i32], slots: Tensor[P; i32], table: Tensor[Rows, Pages; i32]) -> Tensor[1, P, H; T]
  pub fn step_paged<P: Dim, Rows: Dim, Pages: Dim>(x: Tensor[1, P + Rows, H; T], positions: Tensor[P; i32], rows: Tensor[P; i32], slots: Tensor[P; i32], step_positions: Tensor[Rows; i32], table: Tensor[Rows, Pages; i32]) -> Tensor[1, P + Rows, H; T]

Linear

std.nn.linear::Linear<In: Dim, Out: Dim, T: Float = bf16>
  param weight: Tensor[Out, In; T]
  param bias: Tensor[Out; T]?
  pub fn forward<*S: Shape>(x: Tensor[*S, In; T]) -> Tensor[*S, Out; T]

Model

llama::Model<Vocab: Dim, H: Dim, Heads: Dim, KvHeads: Dim, Inner: Dim, Layers: Dim, Batch: Dim, MaxSeq: Dim, T: Float = bf16, Shards: Dim = 1, PageSize: Dim = 64>
  sub embedding: Embedding<Vocab, H, T>
  sub layers: [DecoderLayer<H, Heads, KvHeads, Inner, Batch, MaxSeq, T, Shards, PageSize>; Layers]
  sub norm: RmsNorm<H, T>
  sub lm_head: Linear<H, Vocab / Shards, T>
  pub entry forward<B: Dim, S: Dim>(tokens: Tensor[B, S; i32]) -> Tensor[B, S, Vocab; T]
  pub entry next_token<B: Dim, S: Dim>(tokens: Tensor[B, S; i32]) -> Tensor[B, Vocab; T]
  pub entry decode(token: Tensor[Batch, 1; i32], pos: i32) -> Tensor[Batch, Vocab; T]
  pub entry prefill<S: Dim>(tokens: Tensor[Batch, S; i32], pos: i32) -> Tensor[Batch, Vocab; T]
  pub entry decode_rows(tokens: Tensor[Batch, 1; i32], positions: Tensor[Batch; i32]) -> Tensor[Batch, Vocab; T]
  pub entry prefill_slots<M: Dim, S: Dim>(tokens: Tensor[M, S; i32], slots: Tensor[M; i32], lengths: Tensor[M; i32]) -> Tensor[M, Vocab; T]
  pub entry loss_packed<P: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], targets: Tensor[P; i64], weights: Tensor[P; f32]) -> f32
  pub entry log_probs_packed<P: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], targets: Tensor[P; i64]) -> Tensor[P; f32]
  pub entry hidden_packed<P: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32]) -> Tensor[P, H; T]
  fn packed_states<P: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32]) -> Tensor[P, H; T]
  pub entry prefill_packed<P: Dim>(tokens: Tensor[P; i32], rows: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], last: Tensor[Batch; i32]) -> Tensor[Batch, Vocab; T]
  pub entry step_packed<P: Dim>(tokens: Tensor[P; i32], rows: Tensor[P; i32], positions: Tensor[P; i32], segments: Tensor[P; i32], last: Tensor[Batch; i32], step_tokens: Tensor[Batch, 1; i32], step_positions: Tensor[Batch; i32]) -> Tensor[2 * Batch, Vocab; T]
  pub entry decode_paged<Rows: Dim, Pages: Dim>(tokens: Tensor[Rows, 1; i32], positions: Tensor[Rows; i32], table: Tensor[Rows, Pages; i32]) -> Tensor[Rows, Vocab; T]
  pub entry prefill_paged<P: Dim, Rows: Dim, Pages: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], rows: Tensor[P; i32], slots: Tensor[P; i32], last: Tensor[Rows; i32], table: Tensor[Rows, Pages; i32]) -> Tensor[Rows, Vocab; T]
  pub entry step_paged<P: Dim, Rows: Dim, Pages: Dim>(tokens: Tensor[P; i32], positions: Tensor[P; i32], rows: Tensor[P; i32], slots: Tensor[P; i32], last: Tensor[Rows; i32], step_tokens: Tensor[Rows, 1; i32], step_positions: Tensor[Rows; i32], table: Tensor[Rows, Pages; i32]) -> Tensor[2 * Rows, Vocab; T]
  pub entry generate<Steps: Dim>(token: Tensor[Batch, 1; i32], pos: i32) -> Tensor[Batch, Steps; i32]
  pub entry sample<Steps: Dim>(token: Tensor[Batch, 1; i32], pos: i32, key: Tensor[2; i64], temperature: f32) -> Tensor[Batch, Steps; i32]
  pub entry generate_until<MaxNew: Dim>(token: Tensor[Batch, 1; i32], pos: i32, eos: i32) -> (Tensor[Batch, MaxNew; i32], i32)
  fn step(token: Tensor[Batch, 1; i32], pos: i32) -> Tensor[Batch, Vocab; T]
  fn logits<R: Dim>(x: Tensor[R, H; T]) -> Tensor[R, Vocab; T]
  fn hidden<B: Dim, S: Dim>(tokens: Tensor[B, S; i32]) -> Tensor[B, S, H; T]
  fn hidden_states<B: Dim, S: Dim>(tokens: Tensor[B, S; i32]) -> Tensor[B, S, H; T]

RmsNorm

std.nn.norm::RmsNorm<H: Dim, T: Float = bf16>
  param weight: Tensor[H; T]
  pub fn forward<*S: Shape>(x: Tensor[*S, H; T]) -> Tensor[*S, H; T]

SwiGlu

std.nn.mlp::SwiGlu<H: Dim, Inner: Dim, T: Float = bf16>
  sub gate: Linear<H, Inner, T>
  sub up: Linear<H, Inner, T>
  sub down: Linear<Inner, H, T>
  pub fn forward<*S: Shape>(x: Tensor[*S, H; T]) -> Tensor[*S, H; T]