Skip to content

Fix KV-cache generation in the HF DomynEdge model - #5

Merged
albertoveneri merged 1 commit into
mainfrom
fix/gen-hf-modeling
Oct 21, 2026
Merged

albertoveneri merged 1 commit into
mainfrom
fix/gen-hf-modeling

Conversation

@albertoveneri

Copy link
Copy Markdown
Collaborator

Since transformers 5.18 (#48814) and in particular, generate() passes attention_mask=None when there is no padding. Our default mask then covered only the current input tokens.

Changes

  • When no mask is given, the default mask now spans past_seen_tokens + seq_len, matching the keys.
  • Queries use only the mask columns of the new tokens (attention_mask[:, -sequence_length:]).
  • New tests/test_hf_generation.py: cached generate() and chunked prefill must match an uncached forward.

@francesco-bertolotti francesco-bertolotti left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@albertoveneri
albertoveneri merged commit 4a1d616 into main Oct 21, 2026
1 check passed
@albertoveneri
albertoveneri deleted the fix/gen-hf-modeling branch October 21, 2026 08:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants