Repository: keras-team/keras-nlp
Pull request #780 (2023-02-25 00:59:43)

Simplify the cache decoding graph
This updates our `CachedMultiHeadAttention` layer to avoid slicing the input to a dynamic length during generative decoding. Instead, our cache is always a fixed length, as are the keys and values after the cache is applied.

Overall this ends up being both a nice simplification and speedup. After patching in a fix for #779, generating 25 sequences of length 256 with the base model goes from 52s -> 33s.

mattdangerw: Thanks for the review! Addressed all comments.

Patch:
diff --git a/keras_nlp/layers/cached_multi_head_attention.py b/keras_nlp/layers/cached_multi_head_attention.py
index 094d381e..a86bb93d 100644
--- a/keras_nlp/layers/cached_multi_head_attention.py
+++ b/keras_nlp/layers/cached_multi_head_attention.py
@@ -62,44 +62,24 @@ class CachedMultiHeadAttention(keras.layers.MultiHeadAttention):
         value,
         key=None,
         attention_mask=None,
-        use_causal_mask=False,
         cache=None,
-        cache_index=None,
+        cache_index=0,
     ):
         if not self._built_from_signature:
             self._build_from_signature(query=query, value=value, key=key)
         if key is None:
             key = value
 
-        attention_mask = self._compute_attention_mask(
-            query,
-            value,
-            key=key,
-            attention_mask=attention_mask,
-            use_causal_mask=use_causal_mask,
-        )
-
         query = self._query_dense(query)
         key = self._key_dense(key)
         value = self._value_dense(value)
 
         if cache is not None:
-            if cache_index is None:
-                seq_len = tf.shape(query)[1]
-                k_update_indices = [0, 0, 0, 0, 0]
-                v_update_indices = [0, 1, 0, 0, 0]
-                cache_index = seq_len - 1
-            else:
-                k_update_indices = [0, 0, cache_index, 0, 0]
-                v_update_indices = [0, 1, cache_index, 0, 0]
-            cache = dynamic_update_slice(
-                cache, key[:, tf.newaxis, ...], k_update_indices
-            )
-            cache = dynamic_update_slice(
-                cache, value[:, tf.newaxis, ...], v_update_indices
-            )
-            key = cache[:, 0, : cache_index + 1, :, :]
-            value = cache[:, 1, : cache_index + 1, :, :]
+            key_cache, value_cache = tf.unstack(cache, axis=1)
+            start = [0, cache_index, 0, 0]
+            key = dynamic_update_slice(key_cache, key, start)
+            value = dynamic_update_slice(value_cache, value, start)
+            cache = tf.stack((key, value), axis=1)
 
         query = tf.multiply(query, 1.0 / tf.math.sqrt(float(self._key_dim)))
         attention_scores = tf.einsum(self._dot_product_equation, key, query)
diff --git a/keras_nlp/layers/position_embedding.py b/keras_nlp/layers/position_embedding.py
index 755750c7..64b04173 100644
--- a/keras_nlp/layers/position_embedding.py
+++ b/keras_nlp/layers/position_embedding.py
@@ -17,8 +17,6 @@
 import tensorflow as tf
 from tensorflow import keras
 
-SEQUENCE_AXIS = -2
-
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
 class PositionEmbedding(keras.layers.Layer):
@@ -103,11 +101,12 @@ class PositionEmbedding(keras.layers.Layer):
 
         super().build(input_shape)
 
-    def call(self, inputs):
+    def call(self, inputs, start_index=0):
         if isinstance(inputs, tf.RaggedTensor):
             bounding_shape = inputs.bounding_shape()
             position_embeddings = self._trim_and_broadcast_position_embeddings(
                 bounding_shape,
+                start_index,
             )
             # then apply row lengths to recreate the same ragged shape as inputs
             return tf.RaggedTensor.from_tensor(
@@ -117,12 +116,18 @@ class PositionEmbedding(keras.layers.Layer):
         else:
             return self._trim_and_broadcast_position_embeddings(
                 tf.shape(inputs),
+                start_index,
             )
 
-    def _trim_and_broadcast_position_embeddings(self, shape):
-        input_length = shape[SEQUENCE_AXIS]
+    def _trim_and_broadcast_position_embeddings(self, shape, start_index):
+        feature_length = shape[-1]
+        sequence_length = shape[-2]
         # trim to match the length of the input sequence, which might be less
         # than the sequence_length of the layer.
-        position_embeddings = self.position_embeddings[:input_length, :]
+        position_embeddings = tf.slice(
+            self.position_embeddings,
+            (start_index, 0),
+            (sequence_length, feature_length),
+        )
         # then broadcast to add the missing dimensions to match "shape"
         return tf.broadcast_to(position_embeddings, shape)
diff --git a/keras_nlp/layers/transformer_decoder.py b/keras_nlp/layers/transformer_decoder.py
index afbdb672..469c0a4f 100644
--- a/keras_nlp/layers/transformer_decoder.py
+++ b/keras_nlp/layers/transformer_decoder.py
@@ -212,7 +212,7 @@ class TransformerDecoder(keras.layers.Layer):
         encoder_padding_mask=None,
         encoder_attention_mask=None,
         cache=None,
-        cache_index=None,
+        cache_index=0,
     ):
         """Forward pass of the TransformerDecoder.
 
@@ -236,8 +236,7 @@ class TransformerDecoder(keras.layers.Layer):
                 tokens. `cache` is of shape [B, 2, max_seq_len, num_heads,
                 key_dims].
             cache_index: a int or int Tensor, the index of the current token
-                being processed. If `cache_index=None` while `cache` is set, it
-                means it's the first pass to build the cache.
+                being processed.
         Returns:
             Either a tuple of (outputs, cache) if a cache was passed, or a
             single value outputs if a cache was not passed.
@@ -272,7 +271,19 @@ class TransformerDecoder(keras.layers.Layer):
         x = decoder_sequence  # Intermediate result.
 
         # Compute self attention mask.
-        self_attention_mask = compute_causal_mask(decoder_sequence)
+        batch_size = tf.shape(decoder_sequence)[0]
+        input_length = output_length = tf.shape(decoder_sequence)[1]
+        # We need to handle a rectangular causal mask when doing cached
+        # decoding. For generative inference, `decoder_sequence` will
+        # generally be length 1, and `cache` will be the full generation length.
+        if cache is not None:
+            input_length = tf.shape(cache)[2]
+        self_attention_mask = compute_causal_mask(
+            batch_size,
+            input_length,
+            output_length,
+            cache_index,
+        )
         decoder_mask = merge_padding_and_attention_mask(
             decoder_sequence, decoder_padding_mask, decoder_attention_mask
         )
diff --git a/keras_nlp/layers/transformer_layer_utils.py b/keras_nlp/layers/transformer_layer_utils.py
index 78b83997..2ce24d8a 100644
--- a/keras_nlp/layers/transformer_layer_utils.py
+++ b/keras_nlp/layers/transformer_layer_utils.py
@@ -18,21 +18,28 @@ import tensorflow as tf
 from absl import logging
 
 
-def compute_causal_mask(inputs):
-    input_shape = tf.shape(inputs)
-    batch_size, sequence_length = input_shape[0], input_shape[1]
-    i = tf.range(sequence_length)[:, tf.newaxis]
-    j = tf.range(sequence_length)
-    mask = tf.cast(i >= j, dtype="int32")
-    mask = tf.reshape(mask, (1, input_shape[1], input_shape[1]))
-    mult = tf.concat(
-        [
-            tf.expand_dims(batch_size, -1),
-            tf.constant([1, 1], dtype=tf.int32),
-        ],
-        axis=0,
-    )
-    return tf.tile(mask, mult)
+def compute_causal_mask(batch_size, input_length, output_length, cache_index=0):
+    """Compute a causal attention mask for a transformer decoder.
+
+    Args:
+        batch_size: batch size for the mask.
+        input_length: the length of key/value tensors in the attention layer.
+        output_length: the length of query tensors in the attention layer.
+        cache_index: the current index for cached generation. If passed, the
+            query sequence will be considered to start at `cache_index` rather
+            than zero. For example, a causal mask with `output_length=1` and
+            `cache_index=5` would allow the query tensor to attend to the first
+            five positions of the key/value tensors.
+
+    Return:
+        A causal attention mask with shape
+        `(batch_size, output_length, input_length)` that can be passed to a
+        attention layer.
+    """
+    i = tf.range(output_length)[:, tf.newaxis] + cache_index
+    j = tf.range(input_length)
+    mask = tf.cast(i >= j, dtype="int32")[tf.newaxis, :, :]
+    return tf.broadcast_to(mask, (batch_size, output_length, input_length))
 
 
 def merge_padding_and_attention_mask(
@@ -40,7 +47,7 @@ def merge_padding_and_attention_mask(
     padding_mask,
     attention_mask,
 ):
-    """Merge padding mask with users' customized mask.
+    """Merge the padding mask with a customized attention mask.
 
     Args:
         inputs: the input sequence.
diff --git a/keras_nlp/models/gpt2/gpt2_causal_lm.py b/keras_nlp/models/gpt2/gpt2_causal_lm.py
index 61526b85..63537fb6 100644
--- a/keras_nlp/models/gpt2/gpt2_causal_lm.py
+++ b/keras_nlp/models/gpt2/gpt2_causal_lm.py
@@ -17,7 +17,6 @@ import copy
 
 import tensorflow as tf
 from tensorflow import keras
-from tensorflow.compiler.tf2xla.python.xla import dynamic_update_slice
 
 import keras_nlp
 from keras_nlp.models.gpt2.gpt2_backbone import GPT2Backbone
@@ -201,7 +200,7 @@ class GPT2CausalLM(Task):
     def preprocessor_cls(cls):
         return GPT2CausalLMPreprocessor
 
-    def call_with_cache(self, token_ids, padding_mask, cache, cache_index=None):
+    def call_with_cache(self, token_ids, padding_mask, cache, cache_index):
         """Forward pass of `GPT2CausalLM` with cache.
 
         `call_with_cache` adds an additional forward pass for the model for
@@ -213,101 +212,49 @@ class GPT2CausalLM(Task):
             token_ids: a dense int Tensor, input token ids.
             padding_mask: a dense bool Tensor, input padding mask.
             cache: a dense float Tensor, the cache of key and value.
-            cache_index: int, or int Tensor, defaults to None. If set, it
-                represents the index of current inputs in the whole sequence.
+            cache_index: int, or int Tensor. The index of current inputs in the
+                whole sequence.
 
         Returns:
             A (logits, cache) tuple. Where the first output is the language
             model logits for the input token_ids and the second output is the
             cache.
         """
-        transformer_layers = []
-        for i in range(self.backbone.num_layers):
-            transformer_layers.append(
-                self.backbone.get_layer(f"transformer_layer_{i}")
-            )
-        layer_norm = self.backbone.get_layer("layer_norm")
-        token_embeddings = self.backbone.get_layer("token_embedding")
-        position_embeddings = self.backbone.get_layer("position_embedding")
-        embeddings_add = self.backbone.get_layer("embeddings_add")
-        embeddings_dropout = self.backbone.get_layer("embeddings_dropout")
-
-        token_embedding = token_embeddings(token_ids)
-        if cache_index is None:
-            position_embedding = position_embeddings(token_embedding)
-        else:
-            position_embedding = position_embeddings.position_embeddings[
-                cache_index, :
-            ]
-        x = embeddings_add((token_embedding, position_embedding))
-        x = embeddings_dropout(x)
-        # Each `TransformerDecoder` layer has a cache, we update separately.
+        token_embedding = self.backbone.get_layer("token_embedding")(token_ids)
+        position_embedding = self.backbone.get_layer("position_embedding")(
+            token_embedding, start_index=cache_index
+        )
+        x = self.backbone.get_layer("embeddings_add")(
+            (token_embedding, position_embedding)
+        )
+        x = self.backbone.get_layer("embeddings_dropout")(x)
+        # Each decoder layer has a cache; we update them separately.
         caches = tf.unstack(cache, axis=1)
-        for i, transformer_layer in enumerate(transformer_layers):
+        for i in range(self.backbone.num_layers):
             current_cache = caches[i]
-            x, current_cache = transformer_layer(
+            x, next_cache = self.backbone.get_layer(f"transformer_layer_{i}")(
                 x,
                 decoder_padding_mask=padding_mask,
                 cache=current_cache,
                 cache_index=cache_index,
             )
-            caches[i] = current_cache
+            caches[i] = next_cache
         cache = tf.stack(caches, axis=1)
-        x = layer_norm(x)
+        x = self.backbone.get_layer("layer_norm")(x)
         x = tf.matmul(
             x,
-            self.backbone.token_embedding.embeddings,
+            self.backbone.get_layer("token_embedding").embeddings,
             transpose_b=True,
         )
         return x, cache
 
-    def build_initial_cache(self, initial_inputs, max_length):
-        """Build initial cache based on the prompt.
-
-        This method should be called before the decoding loop to build the
-        initial cache. The cache is of shape [batch_size, `self.num_layers`, 2
-        max_length, `self.num_heads`, `self.hidden_dim // self.num_heads`].
-        The first dim represents it's a key or value in multi-head attention.
-
-        Args:
-            initial_inputs: a dense Tensor, the initial inputs to the decoding
-                loop.
-            max_length: int, the max length of the generated sequence.
-
-        Returns:
-            A (logits, cache) tuple. The first output is the language
-            model logits for the input token_ids and the second output is the
-            cache.
-        """
-        token_ids = initial_inputs["token_ids"]
-        padding_mask = initial_inputs["padding_mask"]
-
-        if max_length < self.backbone.max_sequence_length:
-            token_ids = token_ids[:, :max_length]
-            padding_mask = padding_mask[:, :max_length]
-
-        batch_size = tf.shape(token_ids)[0]
-        outputs = tf.zeros(
-            [batch_size, max_length, self.backbone.vocabulary_size]
-        )
-        cache = tf.zeros(
-            [
-                batch_size,
-                self.backbone.num_layers,
-                2,
-                max_length,
-                self.backbone.num_heads,
-                self.backbone.hidden_dim // self.backbone.num_heads,
-            ],
-        )
-
-        output, cache = self.call_with_cache(
-            token_ids,
-            padding_mask,
-            cache=cache,
-        )
-        outputs = dynamic_update_slice(outputs, output, [0, 0, 0])
-        return outputs, cache
+    def build_empty_cache(self, batch_size, max_length):
+        """Build an empty cache for use with `call_with_cache()`."""
+        num_layers = self.backbone.num_layers
+        num_heads = self.backbone.num_heads
+        head_dim = self.backbone.hidden_dim // self.backbone.num_heads
+        shape = [batch_size, num_layers, 2, max_length, num_heads, head_dim]
+        return tf.zeros(shape)
 
     def _get_token_probability(
         self,
@@ -316,23 +263,9 @@ class GPT2CausalLM(Task):
         cache=None,
         cache_index=None,
     ):
-        model_inputs = {
-            "token_ids": prompt,
-            "padding_mask": mask,
-        }
-        if cache_index is None and cache is None:
-            return self(model_inputs)
-        if cache_index is not None:
-            batch_size = tf.shape(prompt)[0]
-            prompt = tf.slice(prompt, [0, cache_index], [batch_size, 1])
-            mask = mask[:, : cache_index + 1]
-        output, cache = self.call_with_cache(
-            prompt,
-            mask,
-            cache,
-            cache_index,
-        )
-        return output, cache
+        batch_size = tf.shape(prompt)[0]
+        prompt = tf.slice(prompt, [0, cache_index], [batch_size, 1])
+        return self.call_with_cache(prompt, mask, cache, cache_index)
 
     def generate(
         self,
@@ -356,28 +289,38 @@ class GPT2CausalLM(Task):
         """
         if self.preprocessor is None:
             raise ValueError(
-                "`self.preprocessor` is None, please make sure `preprocessor` "
-                "is set before calling `generate`."
+                "`self.preprocessor` is `None`, please make sure "
+                "`preprocessor` is set before calling `generate`."
             )
-        end_token_id = self.preprocessor.tokenizer.end_token_id
-
         sampler = keras_nlp.samplers.get(sampler)
-        if hasattr(self, "jit_compile"):
-            # `jit_compile` is a public property as of tf 2.12. hasattr is for
-            # backward compat.
-            sampler.jit_compile = self.jit_compile
+        # `jit_compile` is a property of keras.Model after tf 2.12.
+        # Use `getattr()` for backwards compatibility.
+        sampler.jit_compile = getattr(self, "jit_compile", True)
         sampler.run_eagerly = self.run_eagerly
-        x, _, _ = self.preprocessor(prompt)
-        if len(x["token_ids"].shape) == 1:
-            x["token_ids"] = x["token_ids"][tf.newaxis, :]
-            x["padding_mask"] = x["padding_mask"][tf.newaxis, :]
-        _, cache = self.build_initial_cache(x, max_length)
-        next_token_probability = self._get_token_probability
+
+        # Tokenize.
+        prompt = self.preprocessor.tokenizer(prompt)
+
+        # Create and seed the cache before generation.
+        token_ids = prompt
+        if prompt.shape.rank == 1:
+            token_ids = tf.RaggedTensor.from_tensor(prompt[tf.newaxis, :])
+        token_ids = token_ids.to_tensor(shape=(None, max_length))
+        # Pass a padding mask of all ones when seeing the cache. The mask will
+        # not affect cached key/values for input tokens we care about.
+        padding_mask = tf.ones_like(token_ids, dtype=tf.bool)
+        batch_size = tf.shape(token_ids)[0]
+        cache = self.build_empty_cache(batch_size, max_length)
+        _, cache = self.call_with_cache(token_ids, padding_mask, cache, 0)
+
+        # Run generation.
         generated = sampler(
-            self.preprocessor.tokenizer(prompt),
-            next_token_probability,
+            prompt,
+            self._get_token_probability,
             max_length=max_length,
-            end_token_id=end_token_id,
+            end_token_id=self.preprocessor.tokenizer.end_token_id,
             cache=cache,
         )
+
+        # Detokenize.
         return self.preprocessor.tokenizer.detokenize(generated)
diff --git a/keras_nlp/samplers/sampler.py b/keras_nlp/samplers/sampler.py
index 03df9e35..50c58e56 100644
--- a/keras_nlp/samplers/sampler.py
+++ b/keras_nlp/samplers/sampler.py
@@ -172,18 +172,6 @@ class Sampler:
         prompt = tf.ragged.boolean_mask(prompt, mask)
         return prompt, mask
 
-    def _validate_token_probability_fn(
-        self, token_probability_fn, prompt, mask
-    ):
-        """Helper method to validate `token_probability_fn` output."""
-        test_pred = token_probability_fn(prompt, mask=mask)
-        if len(test_pred.shape) != 3:
-            raise ValueError(
-                "Output of `token_probability_fn` is not a 3D tensor, "
-                "please provide a function with the output shape "
-                "[batch_size, sequence_length, vocab_size]."
-            )
-
     def _pad_prompt(self, prompt, max_length):
         """Pad prompt to `max_length`."""
         mask = tf.ones_like(prompt, dtype=tf.bool)
@@ -237,7 +225,6 @@ class Sampler:
         # current prompt.
         prompt, mask = self._pad_prompt(prompt, max_length)
         original_padding_mask = tf.identity(mask)
-        self._validate_token_probability_fn(token_probability_fn, prompt, mask)
 
         # Convert `sample` method to a `tf.function` if `self.run_eagerly=False`
         # , and turn on `jit_compile` accordingly.
@@ -359,7 +346,7 @@ class Sampler:
             return [current_index, prompt, mask, cache]
 
         if cache is None:
-            current_index, prompt, mask = tf.while_loop(
+            _, prompt, _ = tf.while_loop(
                 cond=lambda current_index, prompt, mask: tf.less(
                     current_index, max_length
                 ),
@@ -368,7 +355,7 @@ class Sampler:
             )
             return prompt
         # Run a while loop till `max_length` of tokens has been generated.
-        current_index, prompt, mask, cache = tf.while_loop(
+        _, prompt, _, _ = tf.while_loop(
             cond=lambda current_index, prompt, mask, cache: tf.less(
                 current_index, max_length
             ),


Test patch:
diff --git a/keras_nlp/layers/cached_multi_head_attention_test.py b/keras_nlp/layers/cached_multi_head_attention_test.py
index d2e21bc2..0fc20c6f 100644
--- a/keras_nlp/layers/cached_multi_head_attention_test.py
+++ b/keras_nlp/layers/cached_multi_head_attention_test.py
@@ -28,7 +28,11 @@ class CachedMultiHeadAttentionTest(tf.test.TestCase, parameterized.TestCase):
         x = tf.random.uniform(shape=[2, 2, 8])
         layer(query=x, value=x)
 
-    def test_cache_call_is_correct(self):
+    @parameterized.named_parameters(
+        ("graph", False),
+        ("eager", True),
+    )
+    def test_cache_call_is_correct(self, eager):
         batch_size = 2
         seq_len = 5
         num_heads = 2
@@ -37,43 +41,36 @@ class CachedMultiHeadAttentionTest(tf.test.TestCase, parameterized.TestCase):
         layer = CachedMultiHeadAttention(num_heads=num_heads, key_dim=key_dim)
         x = tf.random.uniform(shape=[batch_size, seq_len, num_heads * key_dim])
         cache = tf.zeros([batch_size, 2, seq_len, num_heads, key_dim])
-        # Build the intial cache.
-        intial_seq_len = 2
-        initial_inputs = x[:, :intial_seq_len, :]
+        # Use a causal mask.
+        mask = tf.linalg.band_part(tf.ones((seq_len, seq_len)), -1, 0)
         outputs = tf.zeros_like(x)
-        output, cache = layer(
-            query=initial_inputs,
-            value=initial_inputs,
-            use_causal_mask=True,
-            cache=cache,
-        )
-        # Update the outputs in place.
-        outputs = dynamic_update_slice(outputs, output, [0, 0, 0])
 
-        def call(i, cache, outputs):
-            def loop_body(i, cache, outputs):
+        def call(outputs, cache):
+            def loop_body(i, outputs, cache):
                 # Compute the rest tokens.
-                current_input = x[:, i : i + 1, :]
-                output, cache = layer(
-                    query=current_input,
-                    value=current_input,
-                    use_causal_mask=True,
+                next_input = x[:, i : i + 1, :]
+                next_mask = mask[i : i + 1, :]
+                next_output, cache = layer(
+                    query=next_input,
+                    value=next_input,
                     cache=cache,
                     cache_index=i,
+                    attention_mask=next_mask,
                 )
-                outputs = dynamic_update_slice(outputs, output, [0, i, 0])
-                return i + 1, cache, outputs
+                outputs = dynamic_update_slice(outputs, next_output, [0, i, 0])
+                return i + 1, outputs, cache
 
-            i, cache, cached_outputs = tf.while_loop(
-                cond=lambda i, cache, outputs: i < seq_len,
+            _, outputs, cache = tf.while_loop(
+                cond=lambda i, outputs, cache: i < seq_len,
                 body=loop_body,
-                loop_vars=[i, cache, outputs],
+                loop_vars=[0, outputs, cache],
             )
-            return cached_outputs
+            return outputs, cache
 
-        cached_outputs = call(intial_seq_len, cache, outputs)
-        graph_call = tf.function(call)
-        graph_cached_outputs = graph_call(intial_seq_len, cache, outputs)
-        normal_outputs, _ = layer(query=x, value=x, use_causal_mask=True)
-        self.assertAllClose(cached_outputs, normal_outputs)
-        self.assertAllClose(graph_cached_outputs, normal_outputs)
+        call = call if eager else tf.function(call)
+        output, cache = call(outputs, cache)
+
+        no_loop_outputs, _ = layer(x, x, attention_mask=mask)
+        _, no_loop_cache = layer(x, x, cache=cache, attention_mask=mask)
+        self.assertAllClose(output, no_loop_outputs)
+        self.assertAllClose(cache, no_loop_cache)
diff --git a/keras_nlp/layers/transformer_decoder_test.py b/keras_nlp/layers/transformer_decoder_test.py
index 4896c23f..6b1929b2 100644
--- a/keras_nlp/layers/transformer_decoder_test.py
+++ b/keras_nlp/layers/transformer_decoder_test.py
@@ -266,56 +266,50 @@ class TransformerDecoderTest(tf.test.TestCase, parameterized.TestCase):
         outputs = decoder(decoder_sequence)
         self.assertAllEqual(outputs._keras_mask, mask)
 
-    def test_cached_decoding_is_correct(self):
+    @parameterized.named_parameters(
+        ("graph", False),
+        ("eager", True),
+    )
+    def test_cached_decoding_is_correct(self, eager):
         batch_size = 2
         seq_len = 5
         num_heads = 2
-        hidden_dim = 8
+        head_dim = 4
 
         layer = transformer_decoder.TransformerDecoder(
             intermediate_dim=4,
             num_heads=num_heads,
         )
-        x = tf.random.uniform(shape=[batch_size, seq_len, hidden_dim])
-        cache = tf.zeros(
-            [batch_size, 2, seq_len, num_heads, hidden_dim // num_heads]
-        )
-        # Build the intial cache.
-        initial_seq_len = 2
-        initial_inputs = x[:, :initial_seq_len, :]
+        x = tf.random.uniform(shape=[batch_size, seq_len, num_heads * head_dim])
+        cache = tf.zeros([batch_size, 2, seq_len, num_heads, head_dim])
         outputs = tf.zeros_like(x)
-        output, cache = layer(
-            decoder_sequence=initial_inputs,
-            cache=cache,
-        )
-        # Update the outputs in place.
-        outputs = dynamic_update_slice(outputs, output, [0, 0, 0])
 
-        def call(i, cache, outputs):
-            def loop_body(i, cache, outputs):
+        def call(outputs, cache):
+            def loop_body(i, outputs, cache):
                 # Compute the rest tokens.
-                current_input = x[:, i : i + 1, :]
-                output, cache = layer(
-                    decoder_sequence=current_input,
+                next_input = x[:, i : i + 1, :]
+                next_output, cache = layer(
+                    decoder_sequence=next_input,
                     cache=cache,
                     cache_index=i,
                 )
-                outputs = dynamic_update_slice(outputs, output, [0, i, 0])
-                return i + 1, cache, outputs
+                outputs = dynamic_update_slice(outputs, next_output, [0, i, 0])
+                return i + 1, outputs, cache
 
-            i, cache, cached_outputs = tf.while_loop(
-                cond=lambda i, cache, outputs: i < seq_len,
+            _, outputs, cache = tf.while_loop(
+                cond=lambda i, outputs, cache: i < seq_len,
                 body=loop_body,
-                loop_vars=[i, cache, outputs],
+                loop_vars=[0, outputs, cache],
             )
-            return cached_outputs
-
-        cached_outputs = call(initial_seq_len, cache, outputs)
-        graph_call = tf.function(call)
-        graph_cached_outputs = graph_call(initial_seq_len, cache, outputs)
-        normal_outputs = layer(decoder_sequence=x)
-        self.assertAllClose(cached_outputs, normal_outputs)
-        self.assertAllClose(graph_cached_outputs, normal_outputs)
+            return outputs, cache
+
+        call = call if eager else tf.function(call)
+        output, cache = call(outputs, cache)
+
+        no_loop_outputs = layer(x)
+        _, no_loop_cache = layer(x, cache=cache)
+        self.assertAllClose(output, no_loop_outputs)
+        self.assertAllClose(cache, no_loop_cache)
 
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
diff --git a/keras_nlp/layers/transformer_layer_utils_test.py b/keras_nlp/layers/transformer_layer_utils_test.py
index 73ef7bcf..82ee1d93 100644
--- a/keras_nlp/layers/transformer_layer_utils_test.py
+++ b/keras_nlp/layers/transformer_layer_utils_test.py
@@ -19,8 +19,7 @@ import keras_nlp.layers.transformer_layer_utils as utils
 
 class TransformerEncoderTest(tf.test.TestCase):
     def test_compute_causal_mask(self):
-        inputs = tf.random.uniform(shape=[1, 2, 2])
-        mask = utils.compute_causal_mask(inputs)
+        mask = utils.compute_causal_mask(1, 2, 2)
         self.assertTrue((mask.numpy() == [[1, 0], [1, 1]]).all())
 
     def test_merge_padding_and_attention_mask(self):