Fix Falcon ALiBi bias when using the KV cache

[pctablet505]

Summary

Fixes #2861.

Falcon's ALiBi bias was built over the current query length, but the attention scores span the KV-cache length. When the two differ, this breaks in two ways:

1. Crash. Prefilling a cache longer than the prompt (call_with_cache with query length < cache length) raises a broadcast error: The size of tensor a (32) must match the size of tensor b (8) at non-singleton dimension 3.
2. Silent corruption during generate(). At every single-token decode step the bias tensor has shape [batch, heads, 1, 1], which broadcasts over the cache length without error — but its value is slope * 0 = 0. So generate() runs to completion while adding no ALiBi bias at all at decode time, and the generated tokens diverge from the correct non-cached output.

Fix

In FalconTransformerDecoder.call, when a KV cache is present, build the ALiBi bias over the full cache length instead of the query length. ALiBi positions are absolute (position j gets bias slope * j), so an all-ones mask over the cache yields exactly the same per-position values as a full forward pass; unused cache slots are excluded by the causal attention mask. The no-cache code path is untouched.

Numeric verification (before → after)

Setup: torch backend, CPU, float32. "Full-forward" = a normal forward pass with no cache (training/inference path).

Check	Before	After
Full-forward logits, small random model	—	bit-identical to before (max diff 0.0)
Full-forward logits, falcon_refinedweb_1b_en architecture (load_weights=False)	—	bit-identical to before (max diff 0.0)
Cached prefill (query len 8, cache len 32)	crash (broadcast error)	works; max diff vs full-forward 3.3e-07
Teacher-forced cached decode steps	crash	works; max diff vs full-forward 4.8e-07
Greedy generate() vs non-cached greedy reference	mismatch (silent corruption)	exact match

So: the only behavior that changes is the cached path, which was previously broken (crash or wrong outputs); everything that worked before produces bit-identical outputs.

The regression test added in falcon_causal_lm_test.py (test_cached_decoding_logits_match_full_forward) fails on master with the broadcast error and passes with this fix. Full falcon test suite: 19 passed, 4 skipped.

[gemini-code-assist]

Code Review

This pull request addresses an issue where the ALiBi bias in the Falcon model did not span the full KV-cache length during cached decoding. The FalconTransformerDecoder class in falcon_transformer_decoder.py has been updated to build the ALiBi mask using the full key/value length from the attention_cache when available. A regression test, test_cached_decoding_logits_match_full_forward, has been added to falcon_causal_lm_test.py to verify that cached decoding logits match those of a full forward pass. No review comments were provided, and the implementation looks robust and correct.

[pctablet505]

/gemini review

[gemini-code-assist]

Code Review

This pull request updates the Falcon transformer decoder to correctly handle ALiBi bias when using an attention cache. Specifically, when attention_cache is provided, the mask used for building the ALiBi tensor is sized to the full cache length (kv_length) rather than the input sequence length. Additionally, a unit test test_cached_decoding_logits_match_full_forward has been added to falcon_causal_lm_test.py to verify that cached decoding logits match those of a full forward pass. I have no feedback to provide as there are no review comments.

[divyashreepathihalli]

have you checked the numerics after your edits?