Class LlamaAttention

constructor

• new LlamaAttention(
      weights: Weights,
      config: ModelConfig,
      prefix: string,
      reads?: AttentionReadStrategy,
  ): LlamaAttention

  reads is how the block reads its cache: the cache's own read, fixed when the block is built.

  Parameters

  • weights: Weights
  • config: ModelConfig
  • prefix: string
  • reads: AttentionReadStrategy = cacheRead

  Returns LlamaAttention

ReadonlyheadDim

ReadonlykProj

ReadonlynHeads

ReadonlynKvHeads

ReadonlyoProj

ReadonlyqProj

Readonlyreads

ReadonlyropeBase

Readonlyscale

ReadonlyvProj

forward

• forward(x: MlxArray, cache: Cache, phase: ReadPhase): MlxArray

  x [B, L, hidden] (borrowed) to an owned [B, L, hidden], appending to cache and attending through the read phase names.

  Parameters

  • x: MlxArray
  • cache: Cache
  • phase: ReadPhase

  Returns MlxArray

forwardTraining

• forwardTraining(x: MlxArray, cache: TrainingAttentionCache): MlxArray

  The training forward over a training cache (TrainingCache, PrefixSharedCache), roped at the positions the cache owns (block-wise for a prefix-shared concat), read through its window read, which owns the mask (causal, padded or block-sparse) and the kernel the run composed (the SDPA, or the validated flash kernel).

  Parameters

  • x: MlxArray
  • cache: TrainingAttentionCache

  Returns MlxArray