diff --git a/core/src/main/scala/deepwit/attention/Attention.scala b/core/src/main/scala/deepwit/attention/Attention.scala
index eee0a07..916780a 100644
--- a/core/src/main/scala/deepwit/attention/Attention.scala
+++ b/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)
diff --git a/core/src/main/scala/deepwit/attention/KVCache.scala b/core/src/main/scala/deepwit/attention/KVCache.scala
new file mode 100644
index 0000000..8046cb8
--- /dev/null
+++ b/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)
+    )
diff --git a/core/src/main/scala/deepwit/attention/MultiHeadAttention.scala b/core/src/main/scala/deepwit/attention/MultiHeadAttention.scala
index 6429cb8..2e87f0b 100644
--- a/core/src/main/scala/deepwit/attention/MultiHeadAttention.scala
+++ b/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:
 
diff --git a/core/src/main/scala/deepwit/attention/MultiHeadFusedAttention.scala b/core/src/main/scala/deepwit/attention/MultiHeadFusedAttention.scala
index 972c172..c951271 100644
--- a/core/src/main/scala/deepwit/attention/MultiHeadFusedAttention.scala
+++ b/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,17 +58,13 @@ 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),
@@ -78,5 +72,3 @@ private[attention] object FusedMultiHeadAttentionKernel:
         is_causal = isCausal,
         implementation = "cudnn"
       )
-
-    (res, (queries = queries, keys = keys, values = values))
diff --git a/core/src/main/scala/deepwit/attention/MultiHeadSelfAttention.scala b/core/src/main/scala/deepwit/attention/MultiHeadSelfAttention.scala
index 2c701ab..7479b24 100644
--- a/core/src/main/scala/deepwit/attention/MultiHeadSelfAttention.scala
+++ b/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))
diff --git a/core/src/main/scala/deepwit/attention/MultiHeadUnfusedAttention.scala b/core/src/main/scala/deepwit/attention/MultiHeadUnfusedAttention.scala
index 7df48b8..201f917 100644
--- a/core/src/main/scala/deepwit/attention/MultiHeadUnfusedAttention.scala
+++ b/core/src/main/scala/deepwit/attention/MultiHeadUnfusedAttention.scala
@@ -7,22 +7,33 @@ import dimwit.Label as Λ
 trait MultiHeadUnfusedAttention[Source: Λ, SourceEmbedding: Λ, Target: Λ, TargetEmbedding: Λ, V: IsFloating]:
   this: MultiHeadAttention[Source, SourceEmbedding, Target, TargetEmbedding, V] =>
 
-  /** Implements the multi-head attention by attending head-by-head and concatenating the results. */
-  protected final def attend(source: Tensor2[Source, SourceEmbedding, V], target: Tensor2[Target, TargetEmbedding, V]): (
-      Tensor3[Target, Head, HeadValue, V],
-      MultiHeadAttention.Intermediates[Source, Target, V]
-  ) =
-    val (headValues, queries, keys, values) =
-      zipvmap(Axis[Head])(params.queryWeights, params.keyWeights, params.valueWeights):
-        case (q, k, v) =>
-          val (headValue, perHead) = headAttention(Attention.Params(q, k, v)).applyWithIntermediates(source, target)
-          (headValue, perHead.queries, perHead.keys, perHead.values)
-    (headValues.swap(Axis[Target], Axis[Head]), (queries = queries.swap(Axis[Target], Axis[Head]), keys = keys.swap(Axis[Source], Axis[Head]), values = values.swap(Axis[Source], Axis[Head])))
-
-  /** Per-head attention implementation. */
-  protected def headAttention(
-      headParams: Attention.Params[SourceEmbedding, TargetEmbedding, HeadQuery, HeadKey, HeadValue, V]
-  ): Attention[Source, SourceEmbedding, Target, TargetEmbedding, HeadQuery, HeadKey, HeadValue, V]
+  /** Implements the multi-head attention by attending head-by-head, over every position of the target and the source. */
+  protected final 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] =
+    val targetPositions = Tensor1(Axis[Target]).fromRange(0 until queries.shape(Axis[Target]))
+    val sourcePositions = Tensor1(Axis[Source]).fromRange(0 until keys.shape(Axis[Source]))
+    attendAt(queries, keys, values, targetPositions, sourcePositions)
+
+  /** Attends head by head from the `queries` at `targetPositions` to the `keys` and `values` at `sourcePositions`.
+    *
+    * Decoding with a [[KVCache]] attends from a single target position to every position of the cache.
+    */
+  final def attendAt(
+      queries: Tensor3[Target, Head, HeadQuery, V],
+      keys: Tensor3[Source, Head, HeadKey, V],
+      values: Tensor3[Source, Head, HeadValue, V],
+      targetPositions: Tensor1[Target, Int32],
+      sourcePositions: Tensor1[Source, Int32]
+  ): Tensor3[Target, Head, HeadValue, V] =
+    val attended = zipvmap(Axis[Head])(queries, keys, values):
+      case (headQueries, headKeys, headValues) =>
+        Attention.attend(attentionScore, attentionMask)(headQueries, headKeys, headValues, targetPositions, sourcePositions)
+    attended.swap(Axis[Head], Axis[Target])
+
+  /** How every head scores its queries against its keys. */
+  protected def attentionScore: AttentionScore[Target, Source, HeadQuery, HeadKey, V]
+
+  /** Which source positions every head lets each target position attend to. */
+  protected def attentionMask: AttentionMask[Target, Source]
 
 // --- Unfused multi-head attention implementations ---
 
@@ -30,33 +41,24 @@ class MultiHeadUnfusedFullAttention[Source: Λ, SourceEmbedding: Λ, Target: Λ,
     sourceAxis: Axis[Source],
     targetAxis: Axis[Target],
     params: MultiHeadAttention.Params[SourceEmbedding, TargetEmbedding, V],
-    attentionScore: AttentionScore[Target, Source, HeadQuery, HeadKey, V]
+    protected val attentionScore: AttentionScore[Target, Source, HeadQuery, HeadKey, V]
 ) extends MultiHeadFullAttention[Source, SourceEmbedding, Target, TargetEmbedding, V](sourceAxis, targetAxis, params) with MultiHeadUnfusedAttention[Source, SourceEmbedding, Target, TargetEmbedding, V]:
 
-  override protected def headAttention(
-      headParams: Attention.Params[SourceEmbedding, TargetEmbedding, HeadQuery, HeadKey, HeadValue, V]
-  ): Attention[Source, SourceEmbedding, Target, TargetEmbedding, HeadQuery, HeadKey, HeadValue, V] =
-    FullAttention(sourceAxis, targetAxis, headParams, attentionScore)
+  override protected def attentionMask: AttentionMask[Target, Source] = fullMask
 
 class MultiHeadUnfusedCausalAttention[Source: Λ, SourceEmbedding: Λ, Target: Λ, TargetEmbedding: Λ, V: IsFloating](
     sourceAxis: Axis[Source],
     targetAxis: Axis[Target],
     params: MultiHeadAttention.Params[SourceEmbedding, TargetEmbedding, V],
-    attentionScore: AttentionScore[Target, Source, HeadQuery, HeadKey, V]
+    protected val attentionScore: AttentionScore[Target, Source, HeadQuery, HeadKey, V]
 ) extends MultiHeadCausalAttention[Source, SourceEmbedding, Target, TargetEmbedding, V](sourceAxis, targetAxis, params) with MultiHeadUnfusedAttention[Source, SourceEmbedding, Target, TargetEmbedding, V]:
 
-  override protected def headAttention(
-      headParams: Attention.Params[SourceEmbedding, TargetEmbedding, HeadQuery, HeadKey, HeadValue, V]
-  ): Attention[Source, SourceEmbedding, Target, TargetEmbedding, HeadQuery, HeadKey, HeadValue, V] =
-    CausalAttention(sourceAxis, targetAxis, headParams, attentionScore)
+  override protected def attentionMask: AttentionMask[Target, Source] = causalMask
 
 class MultiHeadCustomAttention[Source: Λ, SourceEmbedding: Λ, Target: Λ, TargetEmbedding: Λ, V: IsFloating](
     params: MultiHeadAttention.Params[SourceEmbedding, TargetEmbedding, V],
-    mask: Shape2[Target, Source] => Tensor2[Target, Source, Bool],
-    attentionScore: AttentionScore[Target, Source, HeadQuery, HeadKey, V]
+    mask: AttentionMask[Target, Source],
+    protected val attentionScore: AttentionScore[Target, Source, HeadQuery, HeadKey, V]
 ) extends MultiHeadAttention[Source, SourceEmbedding, Target, TargetEmbedding, V](params) with MultiHeadUnfusedAttention[Source, SourceEmbedding, Target, TargetEmbedding, V]:
 
-  override protected def headAttention(
-      headParams: Attention.Params[SourceEmbedding, TargetEmbedding, HeadQuery, HeadKey, HeadValue, V]
-  ): Attention[Source, SourceEmbedding, Target, TargetEmbedding, HeadQuery, HeadKey, HeadValue, V] =
-    CustomAttention(headParams, mask, attentionScore)
+  override protected def attentionMask: AttentionMask[Target, Source] = mask
diff --git a/core/src/main/scala/deepwit/attention/ReferenceMultiHeadAttention.scala b/core/src/main/scala/deepwit/attention/ReferenceMultiHeadAttention.scala
index ae72e19..e59f7c9 100644
--- a/core/src/main/scala/deepwit/attention/ReferenceMultiHeadAttention.scala
+++ b/core/src/main/scala/deepwit/attention/ReferenceMultiHeadAttention.scala
@@ -13,7 +13,7 @@ import dimwit.Label as Λ
   */
 class ReferenceMultiHeadAttention[Source: Λ, SourceEmbedding: Λ, Target: Λ, TargetEmbedding: Λ, V: IsFloating](
     params: ReferenceMultiHeadAttention.Params[SourceEmbedding, TargetEmbedding, V],
-    createAttentionMask: Shape2[Target, Source] => Tensor2[Target, Source, Bool],
+    createAttentionMask: AttentionMask[Target, Source],
     attentionScore: AttentionScore[Target, Source, HeadQuery, HeadKey, V]
 ) extends ((Tensor2[Source, SourceEmbedding, V], Tensor2[Target, TargetEmbedding, V]) => Tensor2[Target, TargetEmbedding, V]):
 
diff --git a/core/src/main/scala/deepwit/attention/package.scala b/core/src/main/scala/deepwit/attention/package.scala
index 20e503c..87a4698 100644
--- a/core/src/main/scala/deepwit/attention/package.scala
+++ b/core/src/main/scala/deepwit/attention/package.scala
@@ -1,6 +1,6 @@
 package deepwit.attention
 
-import dimwit.Label
+import dimwit.*
 
 /** The multi-head attention heads. */
 trait Head derives Label
@@ -13,3 +13,8 @@ trait HeadKey derives Label
 
 /** The space a head projects values into. */
 trait HeadValue derives Label
+
+/** Which source positions each target position may attend to, given the positions of both:
+  * `mask(targetPositions, sourcePositions)(row, col) == true` allows the target at `row` to attend to the source at `col`.
+  */
+type AttentionMask[Target, Source] = (Tensor1[Target, Int32], Tensor1[Source, Int32]) => Tensor2[Target, Source, Bool]
diff --git a/core/src/main/scala/deepwit/embedder/LearnedAbsolutePositionalInjector.scala b/core/src/main/scala/deepwit/embedder/LearnedAbsolutePositionalInjector.scala
index 7da2015..ec5abe0 100644
--- a/core/src/main/scala/deepwit/embedder/LearnedAbsolutePositionalInjector.scala
+++ b/core/src/main/scala/deepwit/embedder/LearnedAbsolutePositionalInjector.scala
@@ -21,7 +21,11 @@ class LearnedAbsolutePositionalInjector[Context: Λ, Embedding: Λ, V: IsFloatin
     val length = context.shape(Axis[Context])
     val capacity = params.positionalEmbeddings.shape(Axis[Context])
     require(length <= capacity, s"A context of $length positions was given, but only $capacity were learned.")
-    context + params.positionalEmbeddings.slice(Axis[Context].at(0 until length))
+    context + params.positionalEmbeddings(Axis[Context].at(0 until length))
+
+  /** Adds the positional embedding of `position` to the `embedding` there, e.g. when decoding one position at a time. */
+  def injectAt(position: Tensor0[Int32])(embedding: Tensor1[Embedding, V]): Tensor1[Embedding, V] =
+    embedding + params.positionalEmbeddings(Axis[Context].at(position))
 
 object LearnedAbsolutePositionalInjector:
 
diff --git a/core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala b/core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala
index a425a1a..e2c2013 100644
--- a/core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala
+++ b/core/src/main/scala/deepwit/embedder/VocabularyEmbedder.scala
@@ -9,7 +9,7 @@ import dimwit.Label as Λ
 class VocabularyEmbedder[Vocab: Λ, Embedding: Λ, V: IsFloating](params: VocabularyEmbedder.Params[Vocab, Embedding, V]) extends (Tensor0[Int32] => Tensor1[Embedding, V]):
 
   override def apply(token: Tensor0[Int32]): Tensor1[Embedding, V] =
-    params.vocabularyEmbeddings.slice(Axis[Vocab].at(token))
+    params.vocabularyEmbeddings(Axis[Vocab].at(token))
 
   /** Scores every token by projecting onto its embedding, through the matrix the lookup uses —
     * weight tying, as described in [Using the Output Embedding to Improve Language Models](https://arxiv.org/abs/1608.05859).
diff --git a/core/src/main/scala/deepwit/loss/CrossEntropy.scala b/core/src/main/scala/deepwit/loss/CrossEntropy.scala
index 3c16798..c2d2840 100644
--- a/core/src/main/scala/deepwit/loss/CrossEntropy.scala
+++ b/core/src/main/scala/deepwit/loss/CrossEntropy.scala
@@ -29,7 +29,7 @@ object CategoricalCrossEntropy:
     * @param logits Raw model outputs in (-inf, +inf).
     */
   def fromLogits[L: Λ, V: IsFloating](target: Tensor0[Int32], logits: Tensor1[L, V]): Tensor0[V] =
-    val targetLogit = logits.slice(Axis[L].at(target))
+    val targetLogit = logits(Axis[L].at(target))
     val logNormalizer = logsumexp(logits)
     logNormalizer - targetLogit
 
@@ -38,7 +38,7 @@ object CategoricalCrossEntropy:
     */
   def stable[L: Λ, V: IsFloating](ε: Tensor0[V])(target: Tensor0[Int32], prediction: Tensor1[L, V]): Tensor0[V] =
     val p = prediction.clip(ε, 1f - ε)
-    -p.slice(Axis[L].at(target)).log
+    -p(Axis[L].at(target)).log
 
 /** * Bernoulli Cross Entropy Losses for discrete binary labels.
   * * Specifically for classification tasks where the outcome is Boolean.
diff --git a/core/src/main/scala/deepwit/optimizer/LossScale.scala b/core/src/main/scala/deepwit/optimizer/LossScale.scala
index 87df54a..5c5dc76 100644
--- a/core/src/main/scala/deepwit/optimizer/LossScale.scala
+++ b/core/src/main/scala/deepwit/optimizer/LossScale.scala
@@ -18,8 +18,8 @@ import dimwit.TreeOf.ops.**!
   * val (params, optState) = optimizer.update(grads, state.params, state.optState)
   * val finite = allFinite(grads)
   * TrainState(
-  *   select(finite, params, state.params),
-  *   select(finite, optState, state.optState),
+  *   ifElse(finite, params, state.params),
+  *   ifElse(finite, optState, state.optState),
   *   state.lossScale.next(finite)
   * )
   * }}}
@@ -47,8 +47,8 @@ case class LossScale(scale: Tensor0[Float32], finiteSteps: Tensor0[Int32]):
     val counted = LossScale(scale, finiteSteps + 1)
     val grown = LossScale(scale * growthFactor, 0)
     val backedOff = LossScale(maximum(scale * backoffFactor, minScale), 0)
-    val afterFiniteStep = select(counted.finiteSteps >= growthInterval, grown, counted)
-    select(gradsFinite, afterFiniteStep, backedOff)
+    val afterFiniteStep = ifElse(counted.finiteSteps >= growthInterval, grown, counted)
+    ifElse(gradsFinite, afterFiniteStep, backedOff)
 
 object LossScale:
   def initial(scale: Float = 65536f): LossScale = LossScale(Tensor0(scale), Tensor0(0))
@@ -57,5 +57,5 @@ def allFinite[H: TensorTree, V](grads: Grad[H])(using TreeOf[H, V])(using IsFloa
   grads.value.mapLeaves([T <: Tuple] => (labels: Labels[T]) ?=> (grad: Tensor[T, V]) => grad.isfinite.all).reduce(_ and _)
 
 /** `ifTrue` if `condition` holds, else `ifFalse`, leaf by leaf. Unlike an `if`, this works inside `jit`. */
-def select[P: TensorTree](condition: Tensor0[Bool], ifTrue: P, ifFalse: P): P =
+def ifElse[P: TensorTree](condition: Tensor0[Bool], ifTrue: P, ifFalse: P): P =
   TensorTree[P].zipMap(ifTrue, ifFalse, [T <: Tuple, V] => (labels: Labels[T]) ?=> (a: Tensor[T, V], b: Tensor[T, V]) => where_!(condition, a, b))
diff --git a/core/src/test/scala/deepwit/attention/AttentionSuite.scala b/core/src/test/scala/deepwit/attention/AttentionSuite.scala
index 408dc94..70abd77 100644
--- a/core/src/test/scala/deepwit/attention/AttentionSuite.scala
+++ b/core/src/test/scala/deepwit/attention/AttentionSuite.scala
@@ -81,7 +81,7 @@ class AttentionSuite extends AnyFunSpec with Matchers:
 
     it("gives the first target position only the first source value under a causal mask"):
       val attention = CausalAttention(Axis[Src], Axis[Tgt], paramsWithKeyWeights(identity2(Axis[SrcEmb], Axis[Ky])), AttentionScore.scaledDotProduct)
-      val firstRow = attention(source, target).slice(Axis[Tgt].at(0))
+      val firstRow = attention(source, target).select(Axis[Tgt].at(0))
       firstRow should approxEqual(Tensor(Shape1(Axis[Vl] -> 2)).fromArray(Array(1f, 2f)), 1e-5f)
 
     it("honours a custom attention score function"):
diff --git a/core/src/test/scala/deepwit/attention/KVCacheSuite.scala b/core/src/test/scala/deepwit/attention/KVCacheSuite.scala
new file mode 100644
index 0000000..8faba39
--- /dev/null
+++ b/core/src/test/scala/deepwit/attention/KVCacheSuite.scala
@@ -0,0 +1,53 @@
+package deepwit.attention
+
+import deepwit.*
+import dimwit.*
+import dimwit.stats.Normal
+import org.scalatest.matchers.should.Matchers
+import org.scalatest.funspec.AnyFunSpec
+
+class KVCacheSuite extends AnyFunSpec with Matchers:
+
+  trait Ctx derives Label
+  trait Emb derives Label
+
+  private val ctxExtent = Axis[Ctx] -> 6
+  private val params = MultiHeadSelfAttention.Params.init(numTransformerLayers = 2, numHeads = 2, embeddingExtent = Axis[Emb] -> 4, key = Random.Key(42))
+  private val attention = MultiHeadCausalSelfAttention(Axis[Ctx], params)
+  private val context = Normal.standardNormal(Shape(ctxExtent, Axis[Emb] -> 4)).sample(Random.Key(7))
+  private val positions = Tensor(Shape1(ctxExtent)).fromArray(Array.range(0, ctxExtent.size))
+
+  /** Decodes every position of `context`, one after another, starting from `cache`. */
+  private def decodeFrom(cache: KVCache[Ctx, Float32]): Tensor2[Ctx, Emb, Float32] =
+    val (_, attended) = scan(Axis[Ctx])(cache, (positions, context)):
+      case (cache, (position, embedding)) =>
+        val (attended, nextCache) = attention.decode(cache, position, embedding)
+        (nextCache, attended)
+    attended
+
+  describe("decode"):
+
+    it("reproduces causal attention over the whole context, one position at a time"):
+      decodeFrom(KVCache.empty(ctxExtent, params)) should approxEqual(attention(context), 1e-5f)
+
+    it("works under jit"):
+      val decode = jit(attention.decode)
+      val (attended, _) = decode(KVCache.empty(ctxExtent, params), Tensor0(0), context(Axis[Ctx].at(0)))
+      attended should approxEqual(attention(context).select(Axis[Ctx].at(0)), 1e-5f)
+
+  describe("prefill"):
+
+    it("attends like apply and caches the keys and values of every position"):
+      val (attended, cache) = attention.prefill(context)
+      attended should approxEqual(attention(context), 1e-6f)
+      decodeFrom(cache) should approxEqual(attention(context), 1e-5f)
+
+    it("lets decoding continue after a prompt, whatever follows the prompt in the context"):
+      val promptLength = 3
+      val noise = Normal.standardNormal(context.shape).sample(Random.Key(8))
+      val promptOnly = where_!(Tensor0(promptLength) >! positions, context, noise)
+      val (_, cache) = attention.prefill(promptOnly)
+      (promptLength until ctxExtent.size).foldLeft(cache): (cache, position) =>
+        val (attended, nextCache) = attention.decode(cache, Tensor0(position), context(Axis[Ctx].at(position)))
+        attended should approxEqual(attention(context).select(Axis[Ctx].at(position)), 1e-5f)
+        nextCache
diff --git a/core/src/test/scala/deepwit/attention/MaskingSuite.scala b/core/src/test/scala/deepwit/attention/MaskingSuite.scala
index 087a367..0b7843a 100644
--- a/core/src/test/scala/deepwit/attention/MaskingSuite.scala
+++ b/core/src/test/scala/deepwit/attention/MaskingSuite.scala
@@ -7,10 +7,13 @@ import org.scalatest.funspec.AnyFunSpec
 
 class MaskingSuite extends AnyFunSpec with Matchers:
 
+  private def positionsA(range: Range) = Tensor1(Axis[A]).fromRange(range)
+  private def positionsB(range: Range) = Tensor1(Axis[B]).fromRange(range)
+
   describe("causalMask"):
 
     it("is lower triangular including the diagonal"):
-      val mask = causalMask(Shape2(Axis[A] -> 3, Axis[B] -> 3))
+      val mask = causalMask(positionsA(0 until 3), positionsB(0 until 3))
       // format: off
       val expected = Tensor(Shape(Axis[A] -> 3, Axis[B] -> 3)).fromArray(Array[Float](
         1f, 0f, 0f,
@@ -21,19 +24,23 @@ class MaskingSuite extends AnyFunSpec with Matchers:
       mask.asFloat32 should approxEqual(expected, 1e-6f)
 
     it("lets row i attend to exactly i + 1 positions"):
-      val mask = causalMask(Shape2(Axis[A] -> 4, Axis[B] -> 4))
+      val mask = causalMask(positionsA(0 until 4), positionsB(0 until 4))
       val visible = mask.asFloat32.sum(Axis[B])
       visible should approxEqual(Tensor(Shape1(Axis[A] -> 4)).fromArray(Array(1f, 2f, 3f, 4f)), 1e-6f)
 
     it("works on a non-square score shape"):
-      val mask = causalMask(Shape2(Axis[A] -> 2, Axis[B] -> 4))
+      val mask = causalMask(positionsA(0 until 2), positionsB(0 until 4))
       mask.shape(Axis[A]) shouldBe 2
       mask.shape(Axis[B]) shouldBe 4
       val visible = mask.asFloat32.sum(Axis[B])
       visible should approxEqual(Tensor(Shape1(Axis[A] -> 2)).fromArray(Array(1f, 2f)), 1e-6f)
 
+    it("lets a single position attend to itself and every position before it, as when decoding"):
+      val mask = causalMask(positionsA(2 to 2), positionsB(0 until 4))
+      mask.asFloat32 should approxEqual(Tensor(Shape(Axis[A] -> 1, Axis[B] -> 4)).fromArray(Array(1f, 1f, 1f, 0f)), 1e-6f)
+
   describe("fullMask"):
 
     it("lets every position attend everywhere"):
-      val mask = fullMask(Shape2(Axis[A] -> 2, Axis[B] -> 3))
+      val mask = fullMask(positionsA(0 until 2), positionsB(0 until 3))
       mask.asFloat32 should approxEqual(Tensor(Shape(Axis[A] -> 2, Axis[B] -> 3)).fill(1f), 1e-6f)
diff --git a/core/src/test/scala/deepwit/attention/MultiHeadAttentionSuite.scala b/core/src/test/scala/deepwit/attention/MultiHeadAttentionSuite.scala
index 1aa2dc3..5178b8b 100644
--- a/core/src/test/scala/deepwit/attention/MultiHeadAttentionSuite.scala
+++ b/core/src/test/scala/deepwit/attention/MultiHeadAttentionSuite.scala
@@ -52,8 +52,8 @@ class MultiHeadAttentionSuite extends AnyFunSpec with Matchers:
       val attention = MultiHeadFullAttention(Axis[Src], Axis[Tgt], zeroed)
       val expectedRow = Tensor(Shape1(embExtent)).fromArray(Array(1f, 2f, 3f, 4f))
       val result = attention(source(), target)
-      result.slice(Axis[Tgt].at(0)) should approxEqual(expectedRow, 1e-5f)
-      result.slice(Axis[Tgt].at(2)) should approxEqual(expectedRow, 1e-5f)
+      result(Axis[Tgt].at(0)) should approxEqual(expectedRow, 1e-5f)
+      result(Axis[Tgt].at(2)) should approxEqual(expectedRow, 1e-5f)
 
     it("attends across a source whose embedding space differs from the target's"):
       // Head query and key extents both follow the target embedding, so the dot product lines up.
@@ -86,8 +86,8 @@ class MultiHeadAttentionSuite extends AnyFunSpec with Matchers:
       val attention = MultiHeadFullAttention(Axis[Src], Axis[Tgt], params())
       val (_, intermediates) = attention.applyWithIntermediates(source(), target)
       // The queries are the target read through each head's own query weights.
-      val firstHeadQueries = target.dot(Axis[TgtEmb])(params().queryWeights.slice(Axis[Head].at(0)))
-      intermediates.queries.slice(Axis[Head].at(0)) should approxEqual(firstHeadQueries, 1e-5f)
+      val firstHeadQueries = target.dot(Axis[TgtEmb])(params().queryWeights(Axis[Head].at(0)))
+      intermediates.queries(Axis[Head].at(0)) should approxEqual(firstHeadQueries, 1e-5f)
 
   describe("MultiHeadAttention.Params.xavierUniformDepthScaled"):
 
diff --git a/core/src/test/scala/deepwit/attention/MultiHeadSelfAttentionSuite.scala b/core/src/test/scala/deepwit/attention/MultiHeadSelfAttentionSuite.scala
index 2145844..71a63e8 100644
--- a/core/src/test/scala/deepwit/attention/MultiHeadSelfAttentionSuite.scala
+++ b/core/src/test/scala/deepwit/attention/MultiHeadSelfAttentionSuite.scala
@@ -54,15 +54,20 @@ class MultiHeadSelfAttentionSuite extends AnyFunSpec with Matchers:
       val a = attention(context(13f))
       val b = attention(context(-99f))
       // Positions 0 to 2 cannot attend to the perturbed final position.
-      a.slice(Axis[Ctx].at(0 until 3)) should approxEqual(b.slice(Axis[Ctx].at(0 until 3)), 1e-5f)
+      a(Axis[Ctx].at(0 until 3)) should approxEqual(b(Axis[Ctx].at(0 until 3)), 1e-5f)
       // The final position does see the change.
-      (a.slice(Axis[Ctx].at(3)) - b.slice(Axis[Ctx].at(3))).abs.max.item should be > 1e-3f
+      (a(Axis[Ctx].at(3)) - b(Axis[Ctx].at(3))).abs.max.item should be > 1e-3f
 
     it("lets a full mask expose later positions to earlier ones"):
       val attention = MultiHeadFullSelfAttention(Axis[Ctx], params)
       val a = attention(context(13f))
       val b = attention(context(-99f))
-      (a.slice(Axis[Ctx].at(0)) - b.slice(Axis[Ctx].at(0))).abs.max.item should be > 1e-3f
+      (a(Axis[Ctx].at(0)) - b(Axis[Ctx].at(0))).abs.max.item should be > 1e-3f
+
+    it("is causal self-attention with a custom causal mask"):
+      val custom = MultiHeadCustomSelfAttention(params, causalMask[Ctx, Ctx], AttentionScore.scaledDotProduct)
+      val causal = MultiHeadCausalSelfAttention(Axis[Ctx], params)
+      custom(context(13f)) should approxEqual(causal(context(13f)), 1e-6f)
 
   describe("MultiHeadSelfAttention.Params.xavierUniformDepthScaled"):
 
diff --git a/core/src/test/scala/deepwit/attention/ReferenceMultiHeadAttentionSuite.scala b/core/src/test/scala/deepwit/attention/ReferenceMultiHeadAttentionSuite.scala
index e5910ec..78884e5 100644
--- a/core/src/test/scala/deepwit/attention/ReferenceMultiHeadAttentionSuite.scala
+++ b/core/src/test/scala/deepwit/attention/ReferenceMultiHeadAttentionSuite.scala
@@ -32,9 +32,9 @@ class ReferenceMultiHeadAttentionSuite extends AnyFunSpec with Matchers:
   private val perHead = ReferenceMultiHeadAttention.Params(
     heads = (0 until numHeads).toList.map(head =>
       Attention.Params(
-        batched.queryWeights.slice(Axis[Head].at(head)),
-        batched.keyWeights.slice(Axis[Head].at(head)),
-        batched.valueWeights.slice(Axis[Head].at(head))
+        batched.queryWeights(Axis[Head].at(head)),
+        batched.keyWeights(Axis[Head].at(head)),
+        batched.valueWeights(Axis[Head].at(head))
       )
     ),
     outputProjection = batched.outputProjection
diff --git a/core/src/test/scala/deepwit/embedder/LearnedAbsolutePositionalInjectorSuite.scala b/core/src/test/scala/deepwit/embedder/LearnedAbsolutePositionalInjectorSuite.scala
index 9ad1244..f0323a3 100644
--- a/core/src/test/scala/deepwit/embedder/LearnedAbsolutePositionalInjectorSuite.scala
+++ b/core/src/test/scala/deepwit/embedder/LearnedAbsolutePositionalInjectorSuite.scala
@@ -35,8 +35,8 @@ class LearnedAbsolutePositionalInjectorSuite extends AnyFunSpec with Matchers:
       val injector = LearnedAbsolutePositionalInjector(params)
       val context = Tensor(Shape(contextExtent, embeddingExtent)).fill(0f)
       val injected = injector(context)
-      val first = injected.slice(Axis[A].at(0))
-      val second = injected.slice(Axis[A].at(1))
+      val first = injected(Axis[A].at(0))
+      val second = injected(Axis[A].at(1))
       (first - second).abs.max.item should be > 0f
 
     it("embeds a prefix with the positions it starts at"):
diff --git a/core/src/test/scala/deepwit/embedder/PositionalEncodingSuite.scala b/core/src/test/scala/deepwit/embedder/PositionalEncodingSuite.scala
index 3be3f67..877d695 100644
--- a/core/src/test/scala/deepwit/embedder/PositionalEncodingSuite.scala
+++ b/core/src/test/scala/deepwit/embedder/PositionalEncodingSuite.scala
@@ -26,16 +26,16 @@ class PositionalEncodingSuite extends AnyFunSpec with Matchers:
 
     it("gives distinct encodings to distinct positions"):
       val encoding = PositionalEncoding.sinusoidal2D[A, B, C, Float32](shape)
-      val at00 = encoding.slice(Axis[A].at(0)).slice(Axis[B].at(0))
-      val at10 = encoding.slice(Axis[A].at(1)).slice(Axis[B].at(0))
-      val at01 = encoding.slice(Axis[A].at(0)).slice(Axis[B].at(1))
+      val at00 = encoding.select(Axis[A].at(0)).select(Axis[B].at(0))
+      val at10 = encoding.select(Axis[A].at(1)).select(Axis[B].at(0))
+      val at01 = encoding.select(Axis[A].at(0)).select(Axis[B].at(1))
       (at00 - at10).abs.max.item should be > 1e-3f
       (at00 - at01).abs.max.item should be > 1e-3f
 
     it("encodes the origin as sines of zero and cosines of one"):
       // Per axis the layout is [sin(scales), cos(scales)], and both axes are concatenated.
       val encoding = PositionalEncoding.sinusoidal2D[A, B, C, Float32](shape)
-      val origin = encoding.slice(Axis[A].at(0)).slice(Axis[B].at(0))
+      val origin = encoding.select(Axis[A].at(0)).select(Axis[B].at(0))
       val expected = Tensor(Shape1(Axis[C] -> 8)).fromArray(Array(0f, 0f, 1f, 1f, 0f, 0f, 1f, 1f))
       origin should approxEqual(expected, 1e-6f)
 
@@ -53,18 +53,18 @@ class PositionalEncodingSuite extends AnyFunSpec with Matchers:
     it("encodes a position the same way however densely it is sampled"):
       val coarse = PositionalEncoding.sinusoidal(positions(0f, 1f, 2f, 3f), Axis[C] -> 8)
       val fine = PositionalEncoding.sinusoidal(positions(0f, 0.5f, 1f, 1.5f, 2f, 2.5f, 3f, 3.5f), Axis[C] -> 8)
-      coarse.slice(Axis[A].at(1)) should approxEqual(fine.slice(Axis[A].at(2)), 1e-6f)
-      coarse.slice(Axis[A].at(3)) should approxEqual(fine.slice(Axis[A].at(6)), 1e-6f)
+      coarse(Axis[A].at(1)) should approxEqual(fine(Axis[A].at(2)), 1e-6f)
+      coarse(Axis[A].at(3)) should approxEqual(fine(Axis[A].at(6)), 1e-6f)
 
     it("accepts a position that lies between two grid indices"):
-      val between = PositionalEncoding.sinusoidal(positions(0.5f), Axis[C] -> 8).slice(Axis[A].at(0))
+      val between = PositionalEncoding.sinusoidal(positions(0.5f), Axis[C] -> 8).select(Axis[A].at(0))
       between.abs.max.item should be <= 1f
 
     it("spreads its slowest scale further as the frequency range grows"):
       val narrow = PositionalEncoding.sinusoidal(positions(0f, 8f), Axis[C] -> 8, frequencyRange = 10f)
       val wide = PositionalEncoding.sinusoidal(positions(0f, 8f), Axis[C] -> 8, frequencyRange = 10000f)
-      val narrowTravel = (narrow.slice(Axis[A].at(0)) - narrow.slice(Axis[A].at(1))).abs.max.item
-      val wideTravel = (wide.slice(Axis[A].at(0)) - wide.slice(Axis[A].at(1))).abs.max.item
+      val narrowTravel = (narrow(Axis[A].at(0)) - narrow(Axis[A].at(1))).abs.max.item
+      val wideTravel = (wide(Axis[A].at(0)) - wide(Axis[A].at(1))).abs.max.item
       narrowTravel should be > wideTravel
 
     it("requires an even embedding dimension"):
diff --git a/core/src/test/scala/deepwit/optimizer/LossScaleSuite.scala b/core/src/test/scala/deepwit/optimizer/LossScaleSuite.scala
index ff3af5f..bb63375 100644
--- a/core/src/test/scala/deepwit/optimizer/LossScaleSuite.scala
+++ b/core/src/test/scala/deepwit/optimizer/LossScaleSuite.scala
@@ -57,7 +57,7 @@ class LossScaleSuite extends AnyFunSpec with Matchers:
     it("does not back off below minScale"):
       LossScale(Tensor0(1f), Tensor0(0)).next(Tensor0(false)).scale.item shouldBe 1f
 
-  describe("select"):
+  describe("ifElse"):
 
     it("keeps or replaces every leaf, whatever its shape, inside jit"):
       val optimizer = Adam(Tensor0(0.1f))
@@ -65,7 +65,7 @@ class LossScaleSuite extends AnyFunSpec with Matchers:
       val (newParams, newOptState) = optimizer.update(Autodiff.grad(cost)(params), params, optState)
 
       val step = jit: (finite: Tensor0[Bool]) =>
-        (select(finite, newParams, params), select(finite, newOptState, optState))
+        (ifElse(finite, newParams, params), ifElse(finite, newOptState, optState))
 
       val (skippedParams, skippedOptState) = step(Tensor0(false))
       skippedParams.weight should approxEqual(params.weight, 0f)
diff --git a/core/src/test/scala/deepwit/regularization/PerturbationSuite.scala b/core/src/test/scala/deepwit/regularization/PerturbationSuite.scala
index 19955a1..051b6f7 100644
--- a/core/src/test/scala/deepwit/regularization/PerturbationSuite.scala
+++ b/core/src/test/scala/deepwit/regularization/PerturbationSuite.scala
@@ -23,26 +23,26 @@ class PerturbationSuite extends AnyFunSpec with Matchers:
 
   /** Which directions a diagonal lets through, as a mask of ones and zeros. */
   private def survivors(diagonal: Tensor1[Hidden, Float32]): Seq[Float] =
-    (0 until hiddenExtent.size).map(i => if diagonal.slice(Axis[Hidden].at(i)).item == 0f then 0f else 1f)
+    (0 until hiddenExtent.size).map(i => if diagonal(Axis[Hidden].at(i)).item == 0f then 0f else 1f)
 
   private def readScaleOf(thinned: AffineLayer.Params[Hidden, Out, Float32], i: Int): Seq[Float] =
     (0 until outExtent.size).map: o =>
-      val before = layer.weight.slice(Axis[Hidden].at(i)).slice(Axis[Out].at(o)).item
-      val after = thinned.weight.slice(Axis[Hidden].at(i)).slice(Axis[Out].at(o)).item
+      val before = layer.weight.select(Axis[Hidden].at(i)).select(Axis[Out].at(o)).item
+      val after = thinned.weight.select(Axis[Hidden].at(i)).select(Axis[Out].at(o)).item
       if before == 0f then 0f else after / before
 
   describe("Perturbation.thinningDiagonal"):
 
     it("deletes directions and rescales the survivors by one over the keep probability"):
       val diagonal = Perturbation.thinningDiagonal(hiddenExtent, probability, Random.Key(42))
-      val values = (0 until hiddenExtent.size).map(i => diagonal.slice(Axis[Hidden].at(i)).item)
+      val values = (0 until hiddenExtent.size).map(i => diagonal(Axis[Hidden].at(i)).item)
       values should contain(0f) // this draw deleted something, so there is a deletion to look at
       values.filter(_ != 0f).distinct shouldBe Seq(2f) // survivors scaled by 1 / (1 - 0.5)
 
     it("is exactly the identity at a probability of zero, by arithmetic rather than a branch"):
       val diagonal = Perturbation.thinningDiagonal(hiddenExtent, Tensor0(0f), Random.Key(7))
       (0 until hiddenExtent.size).foreach: i =>
-        diagonal.slice(Axis[Hidden].at(i)).item shouldBe 1f
+        diagonal(Axis[Hidden].at(i)).item shouldBe 1f
 
   describe("Perturbation.thin"):
 
@@ -51,7 +51,7 @@ class PerturbationSuite extends AnyFunSpec with Matchers:
       val thinned = Perturbation.thin(layer, probability, key)
       val diagonal = Perturbation.thinningDiagonal(hiddenExtent, probability, key)
       (0 until hiddenExtent.size).foreach: i =>
-        val expected = diagonal.slice(Axis[Hidden].at(i)).item
+        val expected = diagonal(Axis[Hidden].at(i)).item
         readScaleOf(thinned, i).foreach(_ shouldBe expected +- 1e-5f)
       thinned.bias should approxEqual(layer.bias, 1e-6f)
 
@@ -82,5 +82,5 @@ class PerturbationSuite extends AnyFunSpec with Matchers:
       val gradient = Autodiff.grad(cost)(layer).value
       val diagonal = Perturbation.thinningDiagonal(hiddenExtent, probability, key)
       survivors(diagonal).zipWithIndex.foreach: (survived, i) =>
-        val g = gradient.weight.slice(Axis[Hidden].at(i)).slice(Axis[Out].at(0)).item
+        val g = gradient.weight.select(Axis[Hidden].at(i)).select(Axis[Out].at(0)).item
         if survived == 0f then g shouldBe 0f else g shouldBe 2f +- 1e-5f
diff --git a/examples/src/main/scala/deepwit/examples/autoencoder/AutoEncoderEval.scala b/examples/src/main/scala/deepwit/examples/autoencoder/AutoEncoderEval.scala
index 0ee1b68..d3cdb03 100644
--- a/examples/src/main/scala/deepwit/examples/autoencoder/AutoEncoderEval.scala
+++ b/examples/src/main/scala/deepwit/examples/autoencoder/AutoEncoderEval.scala
@@ -18,8 +18,7 @@ def eval(): Unit =
   require(iterations.nonEmpty, s"No checkpoint to render in ${checkpointer.rootPath}.")
   println(s"Loading checkpoints from: ${checkpointer.rootPath}, steps ${iterations.mkString(", ")}")
 
-  val originals = MNISTLoader.createTestDataset().get.images
-    .slice(Axis[TestSample].at(0 until 64))
+  val originals = MNISTLoader.createTestDataset().get.images(Axis[TestSample].at(0 until 64))
 
   val originalsImgPanel = toImgPanel(originals)
   display(vconcat(
@@ -38,11 +37,9 @@ def eval(): Unit =
 /** Lays 64 samples out as an 8x8 tiling of the digits, so a whole batch reads as one image. */
 private def toImgPanel(tensor: Tensor3[TestSample, Height, Width, Float32]): Tensor2[Prime[Width] |*| Width, Prime[Height] |*| Height, UInt8] =
   tensor
-    .relabel(Axis[TestSample].as(Axis[Prime[Height] |*| Prime[Width]]))
-    .rearrange(
-      (Axis[Prime[Height] |*| Height], Axis[Prime[Width] |*| Width]),
-      (Axis[Prime[Height]] -> 8, Axis[Prime[Width]] -> 8, Axis[Height] -> 28, Axis[Width] -> 28)
-    )
+    .unflatten(Axis[TestSample], Shape(Axis[Prime[Height]] -> 8, Axis[Prime[Width]] -> 8))
+    .flatten((Axis[Prime[Height]], Axis[Height]))
+    .flatten((Axis[Prime[Width]], Axis[Width]))
     .clip(0f, 1f)
     .scale(255f)
     .asInt(VType[UInt8])
diff --git a/examples/src/main/scala/deepwit/examples/dataset/MNISTLoader.scala b/examples/src/main/scala/deepwit/examples/dataset/MNISTLoader.scala
index efea95a..fb0563b 100644
--- a/examples/src/main/scala/deepwit/examples/dataset/MNISTLoader.scala
+++ b/examples/src/main/scala/deepwit/examples/dataset/MNISTLoader.scala
@@ -19,8 +19,8 @@ case class MNISTDataset[S: Label](images: Tensor3[S, Height, Width, Float32], la
     val batchSize = batchExtent.size
     Iterator.iterate(0)(_ + batchSize).map: offset =>
       val batchIds = (0 until batchSize).map(i => (offset + i) % totalSamples)
-      val batchImages = images.slice(Axis[S].at(batchIds)).relabel(Axis[S], Axis[Batch])
-      val batchLabels = labels.slice(Axis[S].at(batchIds)).relabel(Axis[S], Axis[Batch])
+      val batchImages = images(Axis[S].at(batchIds)).relabel(Axis[S], Axis[Batch])
+      val batchLabels = labels(Axis[S].at(batchIds)).relabel(Axis[S], Axis[Batch])
       MNISTBatchSample(batchImages, batchLabels)
 
 object MNISTLoader:
diff --git a/examples/src/main/scala/deepwit/examples/dataset/TwoMoons.scala b/examples/src/main/scala/deepwit/examples/dataset/TwoMoons.scala
index 1d6f68b..f94178e 100644
--- a/examples/src/main/scala/deepwit/examples/dataset/TwoMoons.scala
+++ b/examples/src/main/scala/deepwit/examples/dataset/TwoMoons.scala
@@ -30,8 +30,8 @@ object TwoMoons:
       val batchSize = batchExtent.size
       Iterator.iterate(0)(_ + batchSize).map: offset =>
         val batchIds = (0 until batchSize).map(i => (offset + i) % totalSamples)
-        val batchFeatures = features.slice(Axis[Sample].at(batchIds)).relabel(Axis[Sample], Axis[Batch])
-        val batchLabels = labels.slice(Axis[Sample].at(batchIds)).relabel(Axis[Sample], Axis[Batch])
+        val batchFeatures = features(Axis[Sample].at(batchIds)).relabel(Axis[Sample], Axis[Batch])
+        val batchLabels = labels(Axis[Sample].at(batchIds)).relabel(Axis[Sample], Axis[Batch])
         BatchSample(batchFeatures, batchLabels)
 
   def sampleFix(count: Int, noiseScale: Float): Dataset = sample(count, noiseScale, Key(42))
diff --git a/examples/src/main/scala/deepwit/examples/gpt/FineWebDataset.scala b/examples/src/main/scala/deepwit/examples/gpt/FineWebDataset.scala
index f68cc4b..0ecf666 100644
--- a/examples/src/main/scala/deepwit/examples/gpt/FineWebDataset.scala
+++ b/examples/src/main/scala/deepwit/examples/gpt/FineWebDataset.scala
@@ -32,10 +32,10 @@ object FineWebDataset:
     ).sample(key)
     val shiftedIndices = randomIndices +! 1
     val inputs = randomIndices.vmap(Axis[Sample])(startIndex =>
-      data.dynamicSlice(startIndex, contextLength).relabelTo(Axis[Context])
+      data(Axis[Sample].at(startIndex, contextLength)).relabelTo(Axis[Context])
     )
     val targets = shiftedIndices.vmap(Axis[Sample])(startIndex =>
-      data.dynamicSlice(startIndex, contextLength).relabelTo(Axis[Context])
+      data(Axis[Sample].at(startIndex, contextLength)).relabelTo(Axis[Context])
     )
     val gpu = GPU.devices.head
     BatchSample(targets.asInt32.toDevice(gpu), inputs.asInt32.toDevice(gpu))
diff --git a/examples/src/main/scala/deepwit/examples/gpt/GPT.scala b/examples/src/main/scala/deepwit/examples/gpt/GPT.scala
index 1b130fa..84a8499 100644
--- a/examples/src/main/scala/deepwit/examples/gpt/GPT.scala
+++ b/examples/src/main/scala/deepwit/examples/gpt/GPT.scala
@@ -2,7 +2,7 @@ package deepwit.examples.gpt
 
 import dimwit.*
 import deepwit.activation.softmax
-import deepwit.attention.{Head, HeadKey, HeadQuery, HeadValue}
+import deepwit.attention.{Head, HeadKey, HeadQuery, HeadValue, KVCache}
 import dimwit.stats.Categorical
 import deepwit.embedder.VocabularyEmbedder
 import deepwit.embedder.LearnedAbsolutePositionalInjector
@@ -30,35 +30,37 @@ class GPT[V: IsFloating](params: GPT.Params[V]):
     logits(tokenContext)
       .argmax(Axis[Vocab])
 
-  def generate(
-      prompt: Seq[Int],
-      contextSize: Int = 1024,
-      temperature: Float = 1.0f
-  )(using key: Key): LazyList[Int] =
-    val initialContext = if prompt.size > contextSize then prompt.takeRight(contextSize) else prompt
-    val fastLogits = jit(logits)
-
-    LazyList.unfold((initialContext, key)): (currentContext, currentKey) =>
-
-      val padLength = contextSize - currentContext.size
-      val padded = currentContext ++ Seq.fill(padLength)(0)
-
-      val contextTensor = Tensor1(Axis[Context], VType[Int32]).fromArray(padded.toArray)
-      val stepLogits = fastLogits(contextTensor)
-      val lastLogits = stepLogits.slice(Axis[Context].at(currentContext.size - 1))
-      val scaledLogits = lastLogits /! temperature
-      val probs = Prob(softmax(scaledLogits.asFloat32))
-
-      val (nextKey, sampleKey) = currentKey.split2()
-      val nextTokenTensor = Categorical(probs).sample(sampleKey)
-      val nextToken = nextTokenTensor.item
-
-      val nextContext = if currentContext.size >= contextSize then
-        currentContext.tail :+ nextToken
-      else
-        currentContext :+ nextToken
-
-      Some((nextToken, (nextContext, nextKey)))
+  /** The [[KVCache]]s of a prompt, which fills `tokenContext` from its start, to [[decode]] the positions after it. */
+  def prefill(tokenContext: Tensor1[Context, Int32]): List[KVCache[Context, V]] =
+    val embeddingContext = tokenContext.vmap(Axis[Context])(embedder)
+    val (_, caches) = causalTransformer.prefill(positionalInjector(embeddingContext))
+    caches
+
+  /** The logits for the token after `token` at `position`, as [[logits]] would compute them, given the caches of the positions before. */
+  def decode(caches: List[KVCache[Context, V]], position: Tensor0[Int32], token: Tensor0[Int32]): (Tensor1[Vocab, V], List[KVCache[Context, V]]) =
+    val embedding = positionalInjector.injectAt(position)(embedder(token))
+    val (mixed, nextCaches) = causalTransformer.decode(caches, position, embedding)
+    (outputProjection(mixed), nextCaches)
+
+  /** Samples `numTokens` tokens following `prompt`, decoding one position at a time with [[KVCache]]s.
+    *
+    * The whole loop runs as one `scan` on the device. The last token of the prompt is decoded once more, so that
+    * every step samples from the logits it decoded itself.
+    */
+  def generate(prompt: Seq[Int], numTokens: Int, temperature: Float = 1.0f)(using key: Key): Seq[Int] =
+    val contextSize = params.positionalInjectorParams.positionalEmbeddings.shape(Axis[Context])
+    require(prompt.nonEmpty, "Generating needs a prompt of at least one token.")
+    require(prompt.size + numTokens <= contextSize, s"A prompt of ${prompt.size} tokens leaves room for ${contextSize - prompt.size} more, not $numTokens.")
+
+    val caches = prefill(Tensor1(Axis[Context], VType[Int32]).fromArray(prompt.padTo(contextSize, 0).toArray))
+    val positions = Tensor(Shape1(Axis[Context] -> numTokens)).fromArray(Array.range(prompt.size - 1, prompt.size - 1 + numTokens))
+    val (_, tokens) = scan(Axis[Context])((caches, Tensor0(prompt.last), key), positions):
+      case ((caches, token, key), position) =>
+        val (logits, nextCaches) = decode(caches, position, token)
+        val (nextKey, sampleKey) = key.split2()
+        val nextToken = Categorical(Prob(softmax(logits.asFloat32 /! temperature))).sample(sampleKey)
+        ((nextCaches, nextToken, nextKey), nextToken)
+    tokens.toArray.toSeq
 
 object GPT:
 
diff --git a/examples/src/main/scala/deepwit/examples/gpt/GPT2Transformer.scala b/examples/src/main/scala/deepwit/examples/gpt/GPT2Transformer.scala
index aa6e9ad..8dee5b3 100644
--- a/examples/src/main/scala/deepwit/examples/gpt/GPT2Transformer.scala
+++ b/examples/src/main/scala/deepwit/examples/gpt/GPT2Transformer.scala
@@ -3,7 +3,7 @@ package deepwit.examples.gpt
 import dimwit.*
 import dimwit.Label as Λ
 
-import deepwit.attention.{MultiHeadCausalSelfAttention, MultiHeadSelfAttention}
+import deepwit.attention.{KVCache, MultiHeadCausalSelfAttention, MultiHeadSelfAttention}
 import deepwit.normalization.LayerNorm
 import deepwit.transformer.TransformerBlock
 
@@ -32,6 +32,22 @@ class GPT2Transformer[Context: Λ, Embedding: Λ, V: IsFloating](
       case (context_i, block) => block(context_i)
     res.vmap(Axis[Context])(finalNorm)
 
+  /** Transforms `context` like [[apply]], and keeps one [[KVCache]] per block to [[decode]] the positions after it. */
+  def prefill(context: Tensor2[Context, Embedding, V]): (Tensor2[Context, Embedding, V], List[KVCache[Context, V]]) =
+    val (res, caches) = blocks.foldLeft((context, List.empty[KVCache[Context, V]])):
+      case ((context_i, caches), block) =>
+        val (next, cache) = block.prefill(context_i)
+        (next, caches :+ cache)
+    (res.vmap(Axis[Context])(finalNorm), caches)
+
+  /** Transforms the `embedding` at `position`, as [[apply]] would, given one [[KVCache]] per block. */
+  def decode(caches: List[KVCache[Context, V]], position: Tensor0[Int32], embedding: Tensor1[Embedding, V]): (Tensor1[Embedding, V], List[KVCache[Context, V]]) =
+    val (res, nextCaches) = blocks.zip(caches).foldLeft((embedding, List.empty[KVCache[Context, V]])):
+      case ((embedding_i, nextCaches), (block, cache)) =>
+        val (next, nextCache) = block.decode(cache, position, embedding_i)
+        (next, nextCaches :+ nextCache)
+    (finalNorm(res), nextCaches)
+
 object GPT2Transformer:
 
   case class Params[Embedding, V](
@@ -80,6 +96,18 @@ class GPT2TransformerBlock[Context: Λ, Embedding: Λ, V: IsFloating](
   override protected def contextMixer(context: Tensor2[Context, Embedding, V]): Tensor2[Context, Embedding, V] =
     selfAttention(context.vmap(Axis[Context])(selfAttentionPreNorm))
 
+  /** Applies the block like [[apply]], and keeps the [[KVCache]] of its attention. */
+  def prefill(context: Tensor2[Context, Embedding, V]): (Tensor2[Context, Embedding, V], KVCache[Context, V]) =
+    val (attended, cache) = selfAttention.prefill(context.vmap(Axis[Context])(selfAttentionPreNorm))
+    val contextMixed = context + attended
+    (contextMixed + contextMixed.vmap(Axis[Context])(embeddingMixer), cache)
+
+  /** Applies the block to the `embedding` at `position`, as [[apply]] would, attending to the positions in `cache`. */
+  def decode(cache: KVCache[Context, V], position: Tensor0[Int32], embedding: Tensor1[Embedding, V]): (Tensor1[Embedding, V], KVCache[Context, V]) =
+    val (attended, nextCache) = selfAttention.decode(cache, position, selfAttentionPreNorm(embedding))
+    val contextMixed = embedding + attended
+    (contextMixed + embeddingMixer(contextMixed), nextCache)
+
 object GPT2TransformerBlock:
 
   case class Params[Embedding, V](
diff --git a/examples/src/main/scala/deepwit/examples/gpt/GPTGen.scala b/examples/src/main/scala/deepwit/examples/gpt/GPTGen.scala
index 0e55205..92659b2 100644
--- a/examples/src/main/scala/deepwit/examples/gpt/GPTGen.scala
+++ b/examples/src/main/scala/deepwit/examples/gpt/GPTGen.scala
@@ -29,16 +29,10 @@ object BPETokenizer:
   print("Response: ")
 
   val promptTokens = BPETokenizer.encode(promptText)
-  val tokenStream = model.generate(promptTokens, contextSize = contextExtent.size, temperature = 0.8f)
   val maxTokens = 200
   val eosTokenId = 50256
 
-  tokenStream
-    .takeWhile(_ != eosTokenId)
-    .take(maxTokens)
-    .foreach: tokenId =>
-      val textChunk = BPETokenizer.decode(Seq(tokenId))
-      print(textChunk)
-      System.out.flush()
+  val response = model.generate(promptTokens, numTokens = maxTokens, temperature = 0.8f).takeWhile(_ != eosTokenId)
+  print(BPETokenizer.decode(response))
 
   println("\n\n[Generation Complete]")
diff --git a/examples/src/main/scala/deepwit/examples/mixedPrecision/MixedPrecision.scala b/examples/src/main/scala/deepwit/examples/mixedPrecision/MixedPrecision.scala
index da290a4..c775f5c 100644
--- a/examples/src/main/scala/deepwit/examples/mixedPrecision/MixedPrecision.scala
+++ b/examples/src/main/scala/deepwit/examples/mixedPrecision/MixedPrecision.scala
@@ -9,7 +9,7 @@ import deepwit.training.after
 import deepwit.activation.gelu
 import deepwit.base.{AffineFormLayer, AffineLayer}
 import deepwit.loss.SquaredError
-import deepwit.optimizer.{LossScale, allFinite, select}
+import deepwit.optimizer.{LossScale, allFinite, ifElse}
 
 trait Sample derives Label
 trait Feature derives Label
@@ -66,7 +66,7 @@ def mixedPrecision(): Unit =
   val (dataKey, initKey) = Key(42).split2()
   val shape = Shape(Axis[Sample] -> numSamples, Axis[Feature] -> 1)
   val xs = Uniform(Tensor(shape).fill(-1f), Tensor(shape).fill(1f)).sample(dataKey)
-  val ys = (xs.slice(Axis[Feature].at(0)) *! (2f * math.Pi.toFloat)).sin *! amplitude
+  val ys = (xs(Axis[Feature].at(0)) *! (2f * math.Pi.toFloat)).sin *! amplitude
 
   /** Runs the model in `vtype` and the loss in Float32. The master weights stay Float32, and so do their gradients. */
   def costIn[V: IsFloating](vtype: VType[V])(params: MLP.Params[Float32]): Tensor0[Float32] =
@@ -89,8 +89,8 @@ def mixedPrecision(): Unit =
     val (params, optimizerState) = optimizer.update(grads, state.params, state.optimizerState)
     val finite = allFinite(grads)
     TrainState(
-      select(finite, params, state.params),
-      select(finite, optimizerState, state.optimizerState),
+      ifElse(finite, params, state.params),
+      ifElse(finite, optimizerState, state.optimizerState),
       lossScale.next(finite, growthInterval = 200),
       where(finite, state.skippedSteps, state.skippedSteps + 1)
     )
diff --git a/examples/src/main/scala/deepwit/examples/neuralImage/NeuralImageEval.scala b/examples/src/main/scala/deepwit/examples/neuralImage/NeuralImageEval.scala
index 46c0979..9054ee8 100644
--- a/examples/src/main/scala/deepwit/examples/neuralImage/NeuralImageEval.scala
+++ b/examples/src/main/scala/deepwit/examples/neuralImage/NeuralImageEval.scala
@@ -44,8 +44,8 @@ def eval(): Unit =
   val zoom = 4
 
   val trainedCrop = fourierImage
-    .slice(Axis[Height].at((cropTop until cropTop + cropSize).toIndexedSeq))
-    .slice(Axis[Width].at((cropLeft until cropLeft + cropSize).toIndexedSeq))
+    .select(Axis[Height].at((cropTop until cropTop + cropSize).toIndexedSeq))
+    .select(Axis[Width].at((cropLeft until cropLeft + cropSize).toIndexedSeq))
 
   // The network never saw a position between two pixels; the encoding is defined at them all the same.
   val zoomedSize = cropSize * zoom
@@ -81,5 +81,5 @@ private def render(
 private def nearest(image: Tensor3[Height, Width, Channel, Float32], zoom: Int): Tensor3[Height, Width, Channel, Float32] =
   def repeated(extent: Int) = (0 until extent * zoom).map(_ / zoom).toIndexedSeq
   image
-    .slice(Axis[Height].at(repeated(image.shape(Axis[Height]))))
-    .slice(Axis[Width].at(repeated(image.shape(Axis[Width]))))
+    .select(Axis[Height].at(repeated(image.shape(Axis[Height]))))
+    .select(Axis[Width].at(repeated(image.shape(Axis[Width]))))
diff --git a/examples/src/main/scala/deepwit/examples/neuralImage/NeuralImageTrain.scala b/examples/src/main/scala/deepwit/examples/neuralImage/NeuralImageTrain.scala
index 4c246c5..f0281b0 100644
--- a/examples/src/main/scala/deepwit/examples/neuralImage/NeuralImageTrain.scala
+++ b/examples/src/main/scala/deepwit/examples/neuralImage/NeuralImageTrain.scala
@@ -40,13 +40,13 @@ def batchStream(
 ): Iterator[TrainBatch] =
   // permute across pixels to avoid biasing the network to any particular region of the image
   val permutation = Random.permutation(coordinates.shape.extent(Axis[Pixel]))(key)
-  val coordinatesShuffled = coordinates.take(Axis[Pixel])(permutation)
-  val targetsShuffled = targets.take(Axis[Pixel])(permutation)
+  val coordinatesShuffled = coordinates(Axis[Pixel].at(permutation))
+  val targetsShuffled = targets(Axis[Pixel].at(permutation))
   val totalPixels = coordinates.shape(Axis[Pixel])
   Iterator.iterate(0)(_ + batchSize).map: offset =>
     val batchIds = (0 until batchSize).map(i => (offset + i) % totalPixels)
-    val batchPixels = coordinatesShuffled.slice(Axis[Pixel].at(batchIds)).relabel(Axis[Pixel], Axis[Batch])
-    val batchLabels = targetsShuffled.slice(Axis[Pixel].at(batchIds)).relabel(Axis[Pixel], Axis[Batch])
+    val batchPixels = coordinatesShuffled(Axis[Pixel].at(batchIds)).relabel(Axis[Pixel], Axis[Batch])
+    val batchLabels = targetsShuffled(Axis[Pixel].at(batchIds)).relabel(Axis[Pixel], Axis[Batch])
     TrainBatch(batchPixels, batchLabels)
 
 @main
diff --git a/examples/src/main/scala/deepwit/examples/regression/Regression.scala b/examples/src/main/scala/deepwit/examples/regression/Regression.scala
index fa577bc..092a8bc 100644
--- a/examples/src/main/scala/deepwit/examples/regression/Regression.scala
+++ b/examples/src/main/scala/deepwit/examples/regression/Regression.scala
@@ -136,14 +136,14 @@ def eval(): Unit =
   val grid = NoisyCurve.grid(gridSize)
   display(overlay(
     plots.scatterPlot(
-      data.xs.slice(Axis[Feature].at(0)),
+      data.xs(Axis[Feature].at(0)),
       data.ys,
       _.encoding.x.title := "x",
       _.encoding.y.title := "y",
       _.encoding.size := Json.obj("value" -> Json.fromInt(18))
     ),
     plots.linePlot(
-      grid.slice(Axis[Feature].at(0)),
+      grid(Axis[Feature].at(0)),
       grid.vmap(Axis[Sample])(model),
       _.title := "The model, over points it was never shown",
       _.encoding.color := Json.obj("value" -> Json.fromString("#f58518"))
@@ -161,8 +161,8 @@ case class Dataset(xs: Tensor2[Sample, Feature, Float32], ys: Tensor1[Sample, Fl
     Iterator.iterate(0)(_ + batchExtent.size).map: offset =>
       val ids = (0 until batchExtent.size).map(i => (offset + i) % count)
       BatchSample(
-        xs.slice(Axis[Sample].at(ids)).relabel(Axis[Sample], Axis[Batch]),
-        ys.slice(Axis[Sample].at(ids)).relabel(Axis[Sample], Axis[Batch])
+        xs(Axis[Sample].at(ids)).relabel(Axis[Sample], Axis[Batch]),
+        ys(Axis[Sample].at(ids)).relabel(Axis[Sample], Axis[Batch])
       )
 
 object NoisyCurve:
@@ -179,7 +179,7 @@ object NoisyCurve:
     val (positionKey, noiseKey) = key.split2()
     val shape = Shape(Axis[Sample] -> count, featureExtent)
     val xs = Uniform(Tensor(shape).fill(lower), Tensor(shape).fill(upper)).sample(positionKey)
-    val x = xs.slice(Axis[Feature].at(0))
+    val x = xs(Axis[Feature].at(0))
     val noise = Normal.standardNormal(x.shape).sample(noiseKey) *! noiseScale
     Dataset(xs, truth(x) + noise)
 
diff --git a/examples/src/main/scala/deepwit/examples/thinning/MoonsMLPEval.scala b/examples/src/main/scala/deepwit/examples/thinning/MoonsMLPEval.scala
index 7ca1502..21142db 100644
--- a/examples/src/main/scala/deepwit/examples/thinning/MoonsMLPEval.scala
+++ b/examples/src/main/scala/deepwit/examples/thinning/MoonsMLPEval.scala
@@ -51,7 +51,7 @@ def eval(): Unit =
     coordinates
       .vmap(Axis[Row]):
         _.vmap(Axis[Column]): point =>
-          softmax(MoonsMLP(params).logits(point)).slice(Axis[Output].at(1))
+          softmax(MoonsMLP(params).logits(point)).select(Axis[Output].at(1))
 
   val deployed = sweep(state.params)
   val draws = Key(7).split(numDraws).toList.map(key => sweep(state.params.thin(thinningProbability, key)))
@@ -70,8 +70,8 @@ def eval(): Unit =
 
 private def trainingPlot(dataset: TwoMoons.Dataset): VegaLiteSpec =
   plots.scatterPlot(
-    dataset.features.slice(Axis[Feature].at(0)),
-    dataset.features.slice(Axis[Feature].at(1)),
+    dataset.features(Axis[Feature].at(0)),
+    dataset.features(Axis[Feature].at(1)),
     dataset.labels.toArray.map(moon => s"moon $moon").toSeq,
     _.title := "The data the model was fitted to",
     _.encoding.x.title := "x",
@@ -109,28 +109,21 @@ private def painted[R: Label, C: Label](field: Tensor2[R, C, Float32], colours:
 private def shrunk(field: Tensor2[Row, Column, Float32], factor: Int): Tensor2[Row, Column, Float32] =
   def kept(extent: Int) = (0 until extent by factor).toIndexedSeq
   field
-    .slice(Axis[Row].at(kept(field.shape(Axis[Row]))))
-    .slice(Axis[Column].at(kept(field.shape(Axis[Column]))))
+    .select(Axis[Row].at(kept(field.shape(Axis[Row]))))
+    .select(Axis[Column].at(kept(field.shape(Axis[Column]))))
 
 /** Lays the fields out as a `perSide` x `perSide` tiling, as one field over a larger plane.
   *
-  * The axis they are stacked along is relabelled into a pair of axes and each folded into the grid
-  * axis beside it, so the tiles meet exactly: the gaps a plotting library leaves between separate
+  * The axis they are stacked along is unflattened into a pair of axes, each flattened together with the
+  * grid axis beside it, so the tiles meet exactly: the gaps a plotting library leaves between separate
   * panels never exist to be closed.
   */
 private def tiled(fields: Seq[Tensor2[Row, Column, Float32]], perSide: Int): Tensor2[Prime[Row] |*| Row, Prime[Column] |*| Column, Float32] =
   require(fields.size == perSide * perSide, s"A $perSide x $perSide tiling needs ${perSide * perSide} fields, but was given ${fields.size}.")
   stack(fields, Axis[Draw])
-    .relabel(Axis[Draw].as(Axis[Prime[Row] |*| Prime[Column]]))
-    .rearrange(
-      (Axis[Prime[Row] |*| Row], Axis[Prime[Column] |*| Column]),
-      (
-        Axis[Prime[Row]] -> perSide,
-        Axis[Prime[Column]] -> perSide,
-        Axis[Row] -> fields.head.shape(Axis[Row]),
-        Axis[Column] -> fields.head.shape(Axis[Column])
-      )
-    )
+    .unflatten(Axis[Draw], Shape(Axis[Prime[Row]] -> perSide, Axis[Prime[Column]] -> perSide))
+    .flatten((Axis[Prime[Row]], Axis[Row]))
+    .flatten((Axis[Prime[Column]], Axis[Column]))
 
 private type Colour = (Int, Int, Int)
 
diff --git a/examples/src/main/scala/deepwit/examples/variationalAutoencoder/VariationalAutoencoderEval.scala b/examples/src/main/scala/deepwit/examples/variationalAutoencoder/VariationalAutoencoderEval.scala
index 4904657..44e47ac 100644
--- a/examples/src/main/scala/deepwit/examples/variationalAutoencoder/VariationalAutoencoderEval.scala
+++ b/examples/src/main/scala/deepwit/examples/variationalAutoencoder/VariationalAutoencoderEval.scala
@@ -18,8 +18,7 @@ def eval(): Unit =
   require(iterations.nonEmpty, s"No checkpoint to render in ${checkpointer.rootPath}.")
   println(s"Loading checkpoints from: ${checkpointer.rootPath}, steps ${iterations.mkString(", ")}")
 
-  val originals = MNISTLoader.createTestDataset().get.images
-    .slice(Axis[TestSample].at(0 until 64))
+  val originals = MNISTLoader.createTestDataset().get.images(Axis[TestSample].at(0 until 64))
 
   // Fixed once for the whole rendering, so every checkpoint is judged on the same draws and the
   // panels differ by what was learned rather than by what was sampled.
@@ -47,11 +46,9 @@ private def toDigits(pixels: Tensor2[TestSample, ReconstructedPixel, Float32]):
 /** Lays 64 samples out as an 8x8 tiling of the digits, so a whole batch reads as one image. */
 private def toImgPanel(tensor: Tensor3[TestSample, Height, Width, Float32]): Tensor2[Prime[Width] |*| Width, Prime[Height] |*| Height, UInt8] =
   tensor
-    .relabel(Axis[TestSample].as(Axis[Prime[Height] |*| Prime[Width]]))
-    .rearrange(
-      (Axis[Prime[Height] |*| Height], Axis[Prime[Width] |*| Width]),
-      (Axis[Prime[Height]] -> 8, Axis[Prime[Width]] -> 8, Axis[Height] -> 28, Axis[Width] -> 28)
-    )
+    .unflatten(Axis[TestSample], Shape(Axis[Prime[Height]] -> 8, Axis[Prime[Width]] -> 8))
+    .flatten((Axis[Prime[Height]], Axis[Height]))
+    .flatten((Axis[Prime[Width]], Axis[Width]))
     .clip(0f, 1f)
     .scale(255f)
     .asInt(VType[UInt8])