Skip to content

Fix causal attention masks for uneven layer caches - #302

Open
Entity0-1 wants to merge 1 commit into
NVIDIA:mainfrom
Entity0-1:fix-uneven-cache
Open

Entity0-1 wants to merge 1 commit into
NVIDIA:mainfrom
Entity0-1:fix-uneven-cache

Conversation

@Entity0-1

Copy link
Copy Markdown

PR description

Fixes #298.

Corrects attention-mask handling for unequal per-layer KV caches, addressing Transformers 5.x crashes while preserving causality during multi-token question processing.

  • Right-aligns causal masks with each layer’s cache length and safely extends ordinary causal masks for longer caches.
  • Wraps model-specific eager attention functions and passes dropout by keyword.
  • Preserves broadcast and two-dimensional padding masks. Restricted masks require a layer-specific mask when extending.
  • Adds regression coverage for pipeline generation, model.generate, and batched-versus-token-by-token decoding.

Validation

  • Latest CPU run on Transformers 5.2.0: all 94 relevant tests passed normally, including pipeline tests.
  • Earlier runs: all 94 relevant tests also passed on Transformers 4.57.6 and 5.0.0.
  • A broader offline run reached 220 passed and 115 skipped, then stopped with two failures and one setup error caused by unavailable external test assets.
  • Black, isort, flake8, mypy, copyright-header checks, and git diff --check pass.
  • Full make test validation remains incomplete. GPU testing was not performed.

Checklist

  • Tests are working (make test) — full-suite validation remains incomplete
  • Code is formatted correctly (make style, on errors try fix with make format) — equivalent checks passed as documented above
  • Copyright header is included
  • All commits are signed-off using git commit -s — correct the placeholder author/sign-off email before checking

The following items are N/A: this PR does not add a new press.

  • (new press) mypress_press.py is in the presses directory
  • (new press) MyPress is in __init__.py
  • (new press) README.md is updated with a 1 liner about the new press in the Available presses section
  • (new press) New press is in the default_presses list in tests/default_presses.py
  • (new press) A docstring is provided that follows the same structure as the existing ones

Signed-off-by: Nate <you@example.com>
@copy-pr-bot

copy-pr-bot Bot commented Oct 14, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

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.

PyramidKVPress (and other presses with uneven per-layer caches) crash in the pipeline under transformers 5.x 🤖🤖🤖

1 participant