Fix Falcon ALiBi bias when using the KV cache

The ALiBi bias was built over the current query length while attention
scores span the KV-cache length. With cache length > query length this
either crashes (broadcast mismatch, e.g. prefilling a longer cache)
or silently adds a zero bias during single-token decode steps
([b, n, 1, 1] broadcasts over the cache), corrupting generate() output.

Build the bias over the full cache length when a cache is present;
positions are absolute and unused cache slots are masked out by the
causal mask. The no-cache path is unchanged and produces bit-identical
outputs.

Fixes #2861

Author: pctablet505

keras_hub/src/models/falcon/falcon_causal_lm_test.py

@@ -148,6 +148,42 @@ 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):
+        # Regression test for
+        # https://github.com/keras-team/keras-hub/issues/2861 — the ALiBi
+        # bias must span the full KV-cache length, not just the current
+        # query length.
+        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,20 @@ 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:
+            # With a KV cache, the ALiBi bias must span the full key/value
+            # length (the cache), not just the current query length. ALiBi
+            # positions are absolute, so an all-ones mask yields `arange`
+            # positions over the cache; 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.