Fix Falcon ALiBi bias when using the KV cache

keras_hub/src/models/falcon/falcon_causal_lm_test.py

@@ -148,6 +148,40 @@ def wrapper(*args, **kwargs):
             # We should immediately abort and output the prompt.
             self.assertEqual(prompt, output)
 
+    def test_cached_decoding_logits_match_full_forward(self):
+        # Cached prefill followed by single-token decode steps must produce
+        # the same logits as a full forward pass without a cache.
+        causal_lm = FalconCausalLM(**self.init_kwargs)
+        batch_size, seq_length, kv_length = 2, 8, 16
+        prefill_length = 4
+        token_ids = ops.zeros((batch_size, seq_length), dtype="int32")
+        # Full-forward pass without a cache.
+        full_logits = causal_lm(
+            {
+                "token_ids": token_ids,
+                "padding_mask": ops.ones_like(token_ids),
+            }
+        )
+        # Cached pass: prefill, then decode one token at a time.
+        num_layers = self.backbone.num_layers
+        num_heads = self.backbone.num_attention_heads
+        head_dim = self.backbone.hidden_dim // num_heads
+        cache = ops.zeros(
+            (batch_size, num_layers, 2, kv_length, num_heads, head_dim),
+            dtype="float32",
+        )
+        logits, _, cache = causal_lm.call_with_cache(
+            token_ids[:, :prefill_length], cache, 0
+        )
+        cached_logits = [logits]
+        for i in range(prefill_length, seq_length):
+            logits, _, cache = causal_lm.call_with_cache(
+                token_ids[:, i : i + 1], cache, i
+            )
+            cached_logits.append(logits)
+        cached_logits = ops.concatenate(cached_logits, axis=1)
+        self.assertAllClose(cached_logits, full_logits, atol=1e-5, rtol=1e-5)
+
     def test_generate_compilation(self):
         causal_lm = FalconCausalLM(**self.init_kwargs)
         # Assert we do not recompile with successive calls.

keras_hub/src/models/falcon/falcon_transformer_decoder.py

@@ -118,10 +118,19 @@ def call(
 
         x = self.input_layernorm(inputs)
 
-        mask = decoder_padding_mask
-        if mask is None:
-            batch_size, seq_length = ops.shape(inputs)[:2]
-            mask = ops.ones((batch_size, seq_length), dtype="int32")
+        if attention_cache is not None:
+            # Keys and values span the full cache length, so the ALiBi
+            # bias must as well. Positions are absolute: an all-ones mask
+            # yields `arange` positions over the cache, and unused cache
+            # slots are masked out by the causal attention mask.
+            batch_size = ops.shape(inputs)[0]
+            kv_length = ops.shape(attention_cache)[2]
+            mask = ops.ones((batch_size, kv_length), dtype="int32")
+        else:
+            mask = decoder_padding_mask
+            if mask is None:
+                batch_size, seq_length = ops.shape(inputs)[:2]
+                mask = ops.ones((batch_size, seq_length), dtype="int32")
         alibi = self._build_alibi_tensor(self.num_attention_heads, mask)
 
         # Attention block.