Kv cache

core/src/main/scala/deepwit/attention/Attention.scala

@@ -31,26 +31,41 @@ abstract class Attention[Source: Λ, SourceEmbedding: Λ, Target: Λ, TargetEmbe
     val queries = target.vmap(Axis[Target])(projectQuery)
     val keys = source.vmap(Axis[Source])(projectKey)
     val values = source.vmap(Axis[Source])(projectValue)
-    val attentionWeights = calculateAttentionWeights(queries, keys)
-    val attentionValues = attentionWeights.dot(Axis[Source])(values)
-    (attentionValues, (queries = queries, keys = keys, values = values))
-
-  private final def calculateAttentionWeights(queries: Tensor2[Target, Query, V], keys: Tensor2[Source, Key, V]) =
-    val attentionScores = attentionScore(queries, keys)
-    val attentionMask = createAttentionMask(attentionScores.shape)
-    val maskedScores = where_!(attentionMask, attentionScores, Float.NegativeInfinity)
-    maskedScores.vapply(Axis[Source])(softmax)
+    val targetPositions = Tensor1(Axis[Target]).fromRange(0 until target.shape(Axis[Target]))
+    val sourcePositions = Tensor1(Axis[Source]).fromRange(0 until source.shape(Axis[Source]))
+    val attended = Attention.attend(attentionScore, createAttentionMask)(queries, keys, values, targetPositions, sourcePositions)
+    (attended, (queries = queries, keys = keys, values = values))
 
   /** Creates a boolean mask indicating which source positions each target position may attend to.
     *
-    * @param scoreShape The shape of the attention scores (same shape as attention weights).
     * @return A boolean tensor where `mask(row, col) == true` means the target position
     *         at `row` is allowed to attend to the source position at `col`.
     */
-  protected def createAttentionMask(scoreShape: Shape2[Target, Source]): Tensor2[Target, Source, Bool]
+  protected def createAttentionMask(targetPositions: Tensor1[Target, Int32], sourcePositions: Tensor1[Source, Int32]): Tensor2[Target, Source, Bool]
 
 object Attention:
 
+  /** Attends from the `queries` to the `keys` and `values`: scores them with `attentionScore`, hides what
+    * `attentionMask` does not allow, and averages the values weighted by the softmax of the scores.
+    *
+    * @param targetPositions The positions of the queries, e.g. a single one when decoding with a [[KVCache]].
+    * @param sourcePositions The positions of the keys and values.
+    */
+  def attend[Source: Λ, Target: Λ, Query, Key, Value: Λ, V: IsFloating](
+      attentionScore: AttentionScore[Target, Source, Query, Key, V],
+      attentionMask: AttentionMask[Target, Source]
+  )(
+      queries: Tensor2[Target, Query, V],
+      keys: Tensor2[Source, Key, V],
+      values: Tensor2[Source, Value, V],
+      targetPositions: Tensor1[Target, Int32],
+      sourcePositions: Tensor1[Source, Int32]
+  ): Tensor2[Target, Value, V] =
+    val attentionScores = attentionScore(queries, keys)
+    val maskedScores = where_!(attentionMask(targetPositions, sourcePositions), attentionScores, Float.NegativeInfinity)
+    val attentionWeights = maskedScores.vapply(Axis[Source])(softmax)
+    attentionWeights.dot(Axis[Source])(values)
+
   /** Holds the intermediate tensors computed during attention. */
   type Intermediates[Source, Target, Query, Key, Value, V] = (
       queries: Tensor2[Target, Query, V],
@@ -90,7 +105,8 @@ class FullAttention[Source: Λ, SourceEmbedding: Λ, Target: Λ, TargetEmbedding
     attentionScore: AttentionScore[Target, Source, Query, Key, V]
 ) extends Attention[Source, SourceEmbedding, Target, TargetEmbedding, Query, Key, Value, V](params, attentionScore):
 
-  override protected def createAttentionMask(scoreShape: Shape2[Target, Source]): Tensor2[Target, Source, Bool] = fullMask(scoreShape)
+  override protected def createAttentionMask(targetPositions: Tensor1[Target, Int32], sourcePositions: Tensor1[Source, Int32]): Tensor2[Target, Source, Bool] =
+    fullMask(targetPositions, sourcePositions)
 
 /** Attention where a target position may only attend to source positions up to its own index. */
 class CausalAttention[Source: Λ, SourceEmbedding: Λ, Target: Λ, TargetEmbedding: Λ, Query: Λ, Key: Λ, Value: Λ, V: IsFloating](
@@ -100,16 +116,23 @@ class CausalAttention[Source: Λ, SourceEmbedding: Λ, Target: Λ, TargetEmbeddi
     attentionScore: AttentionScore[Target, Source, Query, Key, V]
 ) extends Attention[Source, SourceEmbedding, Target, TargetEmbedding, Query, Key, Value, V](params, attentionScore):
 
-  override protected def createAttentionMask(scoreShape: Shape2[Target, Source]): Tensor2[Target, Source, Bool] = causalMask(scoreShape)
+  override protected def createAttentionMask(targetPositions: Tensor1[Target, Int32], sourcePositions: Tensor1[Source, Int32]): Tensor2[Target, Source, Bool] =
+    causalMask(targetPositions, sourcePositions)
 
-/** @param mask Creates a boolean mask indicating which source positions each target position may attend to. */
+/** @param attentionMask Creates a boolean mask indicating which source positions each target position may attend to. */
 class CustomAttention[Source: Λ, SourceEmbedding: Λ, Target: Λ, TargetEmbedding: Λ, Query: Λ, Key: Λ, Value: Λ, V: IsFloating](
     params: Attention.Params[SourceEmbedding, TargetEmbedding, Query, Key, Value, V],
-    attentionMask: Shape2[Target, Source] => Tensor2[Target, Source, Bool],
+    attentionMask: AttentionMask[Target, Source],
     attentionScore: AttentionScore[Target, Source, Query, Key, V]
 ) extends Attention[Source, SourceEmbedding, Target, TargetEmbedding, Query, Key, Value, V](params, attentionScore):
 
-  override protected def createAttentionMask(scoreShape: Shape2[Target, Source]): Tensor2[Target, Source, Bool] = attentionMask(scoreShape)
+  override protected def createAttentionMask(targetPositions: Tensor1[Target, Int32], sourcePositions: Tensor1[Source, Int32]): Tensor2[Target, Source, Bool] =
+    attentionMask(targetPositions, sourcePositions)
+
+/** A target position may attend to the source positions up to and including its own. */
+private def causalMask[Target: Λ, Source: Λ](targetPositions: Tensor1[Target, Int32], sourcePositions: Tensor1[Source, Int32]): Tensor2[Target, Source, Bool] =
+  targetPositions.vmap(Axis[Target])(targetPosition => sourcePositions <=! targetPosition)
 
-private def causalMask[Context: Λ, CrossContext: Λ](scoreShape: Shape2[Context, CrossContext]): Tensor[(Context, CrossContext), Bool] = tril(fullMask(scoreShape))
-private def fullMask[Context: Λ, CrossContext: Λ](scoreShape: Shape2[Context, CrossContext]): Tensor[(Context, CrossContext), Bool] = Tensor(scoreShape).fill(true)
+/** A target position may attend to every source position. */
+private def fullMask[Target: Λ, Source: Λ](targetPositions: Tensor1[Target, Int32], sourcePositions: Tensor1[Source, Int32]): Tensor2[Target, Source, Bool] =
+  Tensor(Shape(targetPositions.shape.extent(Axis[Target]), sourcePositions.shape.extent(Axis[Source]))).fill(true)

core/src/main/scala/deepwit/attention/KVCache.scala

@@ -0,0 +1,28 @@
+package deepwit.attention
+
+import dimwit.*
+
+/** The keys and values of every position of a context, so that decoding a position can attend to the
+  * positions before it without projecting them again.
+  *
+  * The cache spans the whole context from the start, so its shape stays the same while decoding, as `jit`
+  * and the loops of DimWit require. Positions not decoded yet hold arbitrary values, which are never attended to.
+  *
+  * @see [[MultiHeadCausalSelfAttention.prefill]] and [[MultiHeadCausalSelfAttention.decode]]
+  */
+case class KVCache[Context, V](
+    keys: Tensor3[Context, Head, HeadKey, V],
+    values: Tensor3[Context, Head, HeadValue, V]
+)
+
+object KVCache:
+
+  /** A cache spanning `contextExtent` for the attention with `params`, before any position is decoded. */
+  def empty[Context: Label, Embedding, V: IsFloating](contextExtent: AxisExtent[Context], params: MultiHeadSelfAttention.Params[Embedding, V]): KVCache[Context, V] =
+    val keyWeights = params.multiHeadAttention.keyWeights
+    val valueWeights = params.multiHeadAttention.valueWeights
+    val headExtent = keyWeights.shape.extent(Axis[Head])
+    KVCache(
+      keys = Tensor(Shape(contextExtent, headExtent, keyWeights.shape.extent(Axis[HeadKey])), VType[V]).fill(0f),
+      values = Tensor(Shape(contextExtent, headExtent, valueWeights.shape.extent(Axis[HeadValue])), VType[V]).fill(0f)
+    )

core/src/main/scala/deepwit/attention/MultiHeadAttention.scala

@@ -22,14 +22,17 @@ abstract class MultiHeadAttention[Source: Λ, SourceEmbedding: Λ, Target: Λ, T
       Tensor2[Target, TargetEmbedding, V],
       MultiHeadAttention.Intermediates[Source, Target, V]
   ) =
-    val (heads, intermediates) = attend(source, target)
-    val res = heads.vmap(Axis[Target])(heads => projectHeadValues(heads.flatten))
-    (res, intermediates)
+    val queries = target.dot(Axis[TargetEmbedding])(params.queryWeights)
+    val keys = source.dot(Axis[SourceEmbedding])(params.keyWeights)
+    val values = source.dot(Axis[SourceEmbedding])(params.valueWeights)
+    (projectHeads(attend(queries, keys, values)), (queries = queries, keys = keys, values = values))
 
-  protected def attend(source: Tensor2[Source, SourceEmbedding, V], target: Tensor2[Target, TargetEmbedding, V]): (
-      Tensor3[Target, Head, HeadValue, V],
-      MultiHeadAttention.Intermediates[Source, Target, V]
-  )
+  /** Projects the values attended to by all heads back into the target embedding space. */
+  final def projectHeads(headValues: Tensor3[Target, Head, HeadValue, V]): Tensor2[Target, TargetEmbedding, V] =
+    headValues.vmap(Axis[Target])(headValues => projectHeadValues(headValues.flatten))
+
+  /** Attends from the `queries` of every head to its `keys` and `values`. */
+  protected def attend(queries: Tensor3[Target, Head, HeadQuery, V], keys: Tensor3[Source, Head, HeadKey, V], values: Tensor3[Source, Head, HeadValue, V]): Tensor3[Target, Head, HeadValue, V]
 
 object MultiHeadAttention:
 

core/src/main/scala/deepwit/attention/MultiHeadFusedAttention.scala

@@ -11,10 +11,6 @@ import scala.util.Try
 /** Marker trait for multi-head attention implementations that use [[FusedMultiHeadAttentionKernel]]. Restricted to BFloat16 due to kernel constraints. */
 trait MultiHeadFusedAttention[Source: Λ, SourceEmbedding: Λ, Target: Λ, TargetEmbedding: Λ]:
   this: MultiHeadAttention[Source, SourceEmbedding, Target, TargetEmbedding, BFloat16] =>
-  protected final def headAttention(
-      headParams: Attention.Params[SourceEmbedding, TargetEmbedding, HeadQuery, HeadKey, HeadValue, BFloat16]
-  ): Attention[Source, SourceEmbedding, Target, TargetEmbedding, HeadQuery, HeadKey, HeadValue, BFloat16] =
-    throw new RuntimeException("Fused attention does not run head-by-head, so this method should never be called.")
 
 // --- Fused multi-head attention implementations ---
 
@@ -25,10 +21,11 @@ class MultiHeadFusedFullAttention[Source: Λ, SourceEmbedding: Λ, Target: Λ, T
 ) extends MultiHeadFullAttention[Source, SourceEmbedding, Target, TargetEmbedding, BFloat16](sourceAxis, targetAxis, params) with MultiHeadFusedAttention[Source, SourceEmbedding, Target, TargetEmbedding]:
 
   override protected def attend(
-      source: Tensor2[Source, SourceEmbedding, BFloat16],
-      target: Tensor2[Target, TargetEmbedding, BFloat16]
-  ): (Tensor3[Target, Head, HeadValue, BFloat16], MultiHeadAttention.Intermediates[Source, Target, BFloat16]) =
-    FusedMultiHeadAttentionKernel.attend(params, source, target, isCausal = false)
+      queries: Tensor3[Target, Head, HeadQuery, BFloat16],
+      keys: Tensor3[Source, Head, HeadKey, BFloat16],
+      values: Tensor3[Source, Head, HeadValue, BFloat16]
+  ): Tensor3[Target, Head, HeadValue, BFloat16] =
+    FusedMultiHeadAttentionKernel.attend(queries, keys, values, isCausal = false)
 
 object MultiHeadFusedFullAttention:
   export FusedMultiHeadAttentionKernel.canRun as isAvailable
@@ -40,10 +37,11 @@ class MultiHeadFusedCausalAttention[Source: Λ, SourceEmbedding: Λ, Target: Λ,
 ) extends MultiHeadCausalAttention[Source, SourceEmbedding, Target, TargetEmbedding, BFloat16](sourceAxis, targetAxis, params) with MultiHeadFusedAttention[Source, SourceEmbedding, Target, TargetEmbedding]:
 
   override protected def attend(
-      source: Tensor2[Source, SourceEmbedding, BFloat16],
-      target: Tensor2[Target, TargetEmbedding, BFloat16]
-  ): (Tensor3[Target, Head, HeadValue, BFloat16], MultiHeadAttention.Intermediates[Source, Target, BFloat16]) =
-    FusedMultiHeadAttentionKernel.attend(params, source, target, isCausal = true)
+      queries: Tensor3[Target, Head, HeadQuery, BFloat16],
+      keys: Tensor3[Source, Head, HeadKey, BFloat16],
+      values: Tensor3[Source, Head, HeadValue, BFloat16]
+  ): Tensor3[Target, Head, HeadValue, BFloat16] =
+    FusedMultiHeadAttentionKernel.attend(queries, keys, values, isCausal = true)
 
 object MultiHeadFusedCausalAttention:
   export FusedMultiHeadAttentionKernel.canRun as isAvailable
@@ -60,23 +58,17 @@ private[attention] object FusedMultiHeadAttentionKernel:
     dtype == DType.BFloat16 && hasCudaDevice
 
   /** Kernel implementation. */
-  def attend[Source: Λ, SourceEmbedding: Λ, Target: Λ, TargetEmbedding: Λ](
-      params: MultiHeadAttention.Params[SourceEmbedding, TargetEmbedding, BFloat16],
-      source: Tensor2[Source, SourceEmbedding, BFloat16],
-      target: Tensor2[Target, TargetEmbedding, BFloat16],
+  def attend[Source: Λ, Target: Λ](
+      queries: Tensor3[Target, Head, HeadQuery, BFloat16],
+      keys: Tensor3[Source, Head, HeadKey, BFloat16],
+      values: Tensor3[Source, Head, HeadValue, BFloat16],
       isCausal: Boolean
-  ): (Tensor3[Target, Head, HeadValue, BFloat16], MultiHeadAttention.Intermediates[Source, Target, BFloat16]) =
-    val queries = target.dot(Axis[TargetEmbedding])(params.queryWeights)
-    val keys = source.dot(Axis[SourceEmbedding])(params.keyWeights)
-    val values = source.dot(Axis[SourceEmbedding])(params.valueWeights)
-
-    val res = PyBridge.liftPyTensor[(Target, Head, HeadValue), BFloat16]:
+  ): Tensor3[Target, Head, HeadValue, BFloat16] =
+    PyBridge.liftPyTensor[(Target, Head, HeadValue), BFloat16]:
       Jax.jnn.dot_product_attention(
         PyBridge.toPyTensor(queries),
         PyBridge.toPyTensor(keys),
         PyBridge.toPyTensor(values),
         is_causal = isCausal,
         implementation = "cudnn"
       )
-
-    (res, (queries = queries, keys = keys, values = values))

core/src/main/scala/deepwit/attention/MultiHeadSelfAttention.scala

@@ -35,21 +35,50 @@ class MultiHeadFullSelfAttention[Context: Λ, Embedding: Λ, V: IsFloating](
     params: MultiHeadSelfAttention.Params[Embedding, V]
 ) extends MultiHeadSelfAttention[Context, Embedding, V](MultiHeadFullAttention(contextAxis, contextAxis, params.multiHeadAttention))
 
-/** Multi-head self-attention with [[MultiHeadCausalAttention]]. */
+/** Multi-head self-attention with [[MultiHeadCausalAttention]].
+  *
+  * Besides attending over a whole context at once, it can decode one position at a time with a [[KVCache]]:
+  * [[prefill]] attends over a prompt and keeps its keys and values, and [[decode]] continues from there.
+  */
 class MultiHeadCausalSelfAttention[Context: Λ, Embedding: Λ, V: IsFloating](
     contextAxis: Axis[Context],
     params: MultiHeadSelfAttention.Params[Embedding, V]
-) extends MultiHeadSelfAttention[Context, Embedding, V](MultiHeadCausalAttention(contextAxis, contextAxis, params.multiHeadAttention))
+) extends MultiHeadSelfAttention[Context, Embedding, V](MultiHeadCausalAttention(contextAxis, contextAxis, params.multiHeadAttention)):
+
+  private val weights = params.multiHeadAttention
+
+  // Decoding attends from one position to the cache, which the fused kernel cannot, so it always attends head by head.
+  private val decoding = MultiHeadUnfusedCausalAttention(contextAxis, contextAxis, weights, AttentionScore.scaledDotProduct)
+
+  /** Attends over `context` like [[apply]], and keeps its keys and values to [[decode]] the positions after it. */
+  def prefill(context: Tensor2[Context, Embedding, V]): (Tensor2[Context, Embedding, V], KVCache[Context, V]) =
+    val (attended, intermediates) = applyWithIntermediates(context)
+    (attended, KVCache(intermediates.keys, intermediates.values))
+
+  /** Attends from the `embedding` at `position` to itself and every position before it, as [[apply]] would.
+    *
+    * @param cache Holds the keys and values of the positions before `position`.
+    * @param position May be traced, e.g. under `jit` or in a loop.
+    * @return The attended embedding, and `cache` with the key and value of `position` written into it.
+    */
+  def decode(cache: KVCache[Context, V], position: Tensor0[Int32], embedding: Tensor1[Embedding, V]): (Tensor1[Embedding, V], KVCache[Context, V]) =
+    val query = embedding.dot(Axis[Embedding])(weights.queryWeights)
+    val key = embedding.dot(Axis[Embedding])(weights.keyWeights)
+    val value = embedding.dot(Axis[Embedding])(weights.valueWeights)
+    val keys = cache.keys.updated(contextAxis.at(position))(key)
+    val values = cache.values.updated(contextAxis.at(position))(value)
+    val headValues = decoding.attendAt(
+      queries = stack(Seq(query), contextAxis),
+      keys = keys,
+      values = values,
+      targetPositions = stack(Seq(position), contextAxis),
+      sourcePositions = Tensor1(contextAxis).fromRange(0 until keys.shape(contextAxis))
+    )
+    (decoding.projectHeads(headValues).select(contextAxis.at(0)), KVCache(keys, values))
 
 /** Multi-head self-attention with [[MultiHeadCustomAttention]]. */
 class MultiHeadCustomSelfAttention[Context: Λ, Embedding: Λ, V: IsFloating](
     params: MultiHeadSelfAttention.Params[Embedding, V],
-    mask: AxisExtent[Context] => Tensor2[Context, Context, Bool],
+    mask: AttentionMask[Context, Context],
     attentionScore: AttentionScore[Context, Context, HeadQuery, HeadKey, V]
-) extends MultiHeadSelfAttention[Context, Embedding, V](
-      MultiHeadCustomAttention(
-        params.multiHeadAttention,
-        shape2 => mask(shape2.extent(Axis[Context])),
-        attentionScore
-      )
-    )
+) extends MultiHeadSelfAttention[Context, Embedding, V](MultiHeadCustomAttention(params.multiHeadAttention, mask, attentionScore))