Problem
prepare_4d_causal_attention_mask expects the padding-mask key dimension to match kv_len. In the Eagle3 TTT loop with a KV cache, kv_len grows with cached chunks, but the incoming 2D attention_mask still covers only the current chunk.
For seq_len = 4 at cached step k = 1:
causal mask: [1, 1, 4, 8]
expanded mask: [B, 1, 1, 4]
Adding these tensors fails on the key dimension.
Impact
The non-flex fallback attention path crashes whenever cached Eagle3 TTT steps require kv_len > seq_len.
Fix
Fixed in PR #30 by extending the 2D keep-mask to kv_len before converting it to the additive 4D mask. The final fix repeats the fixed chunk mask pattern, preserving padded cached columns instead of marking all cached positions valid.
Problem
prepare_4d_causal_attention_maskexpects the padding-mask key dimension to matchkv_len. In the Eagle3 TTT loop with a KV cache,kv_lengrows with cached chunks, but the incoming 2Dattention_maskstill covers only the current chunk.For
seq_len = 4at cached stepk = 1:Adding these tensors fails on the key dimension.
Impact
The non-flex fallback attention path crashes whenever cached Eagle3 TTT steps require
kv_len > seq_len.Fix
Fixed in PR #30 by extending the 2D keep-mask to
kv_lenbefore converting it to the additive 4D mask. The final fix repeats the fixed chunk mask pattern, preserving padded cached columns instead of marking all cached positions valid.