Fix Falcon ALiBi bias when using the KV cache

The ALiBi bias was built from the query window instead of the cache, so
it was shorter than the attention scores whenever a cache was in use:

- With more than one token per step it fails to broadcast, raising an
  error.
- With a single token it broadcasts silently and evaluates to all zeros,
  because `cumsum(ones(1)) - 1 == 0`. Since `generate()` decodes one
  token at a time, `FalconCausalLM.generate()` has been running with
  ALiBi effectively disabled at every decode step.

Build the bias over the cache length when a cache is present. Unused
cache positions are already masked out by the causal attention mask, so
an all-ones mask yields the correct absolute positions.

Verified: cached decoding now matches a full forward pass (max abs diff
2.98e-07), and the no-cache path is bit-identical before and after, both
with and without a padding mask.

Author: GargiGupta-io

keras_hub/src/models/falcon/falcon_causal_lm_test.py

@@ -148,6 +148,48 @@ def wrapper(*args, **kwargs):
             # We should immediately abort and output the prompt.
             self.assertEqual(prompt, output)
 
+    def test_cached_decoding_matches_full_forward(self):
+        # The ALiBi bias must span the cache, not just the query window.
+        # Building it from the query length breaks cached decoding two ways:
+        # it fails to broadcast when more than one token is passed, and it
+        # collapses to all zeros for a single token, silently disabling ALiBi.
+        causal_lm = FalconCausalLM(**self.init_kwargs)
+        token_ids = self.input_data["token_ids"]
+        batch_size, seq_length = token_ids.shape
+        num_layers = self.backbone.num_layers
+        num_heads = self.backbone.num_attention_heads
+        head_dim = self.backbone.hidden_dim // num_heads
+        cache_shape = [
+            batch_size,
+            num_layers,
+            2,
+            seq_length,
+            num_heads,
+            head_dim,
+        ]
+        dtype = causal_lm.compute_dtype
+
+        # One pass over the whole sequence, seeding an empty cache.
+        full_logits, _, _ = causal_lm.call_with_cache(
+            token_ids, ops.zeros(cache_shape, dtype=dtype), 0
+        )
+
+        # Prefill half the sequence, then decode a token at a time.
+        cache = ops.zeros(cache_shape, dtype=dtype)
+        prefill = seq_length // 2
+        logits, _, cache = causal_lm.call_with_cache(
+            token_ids[:, :prefill], cache, 0
+        )
+        pieces = [logits]
+        for i in range(prefill, seq_length):
+            logits, _, cache = causal_lm.call_with_cache(
+                token_ids[:, i : i + 1], cache, i
+            )
+            pieces.append(logits)
+        stepwise_logits = ops.concatenate(pieces, axis=1)
+
+        self.assertAllClose(stepwise_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,22 @@ 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 whole cache, so the ALiBi bias must
+            # too. Building it from the query window instead leaves the bias
+            # shorter than the attention scores: it fails to broadcast when
+            # more than one token is passed, and collapses to all zeros
+            # during single-token decoding, silently disabling ALiBi.
+            # Unused cache positions are already masked out by the causal
+            # attention mask, so an all-ones mask gives absolute positions.
+            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.