Fix empty tensor shape issue in DynamicCache for torch.compile

[yashwantbezawada]

What does this PR do?

Fixes #42027

This PR fixes a regression where torch.cat receives incorrectly shaped empty tensors during GPT2 model tracing with torch.compile, causing compilation failures.

Background

The issue was introduced in commit dc11a3c (PR #39797) where empty cache tensors were initialized as 1D tensors with shape [0] using torch.tensor([]). When these are concatenated with 4D key/value tensors [batch_size, num_heads, seq_len, head_dim] along dim=-2, torch.compiles tracing fails with empty tensor errors.

Changes

Modified DynamicLayer.lazy_initialization()

• Changed from: torch.tensor([], dtype=..., device=...) → shape [0] (1D)
• Changed to: torch.zeros((batch_size, num_heads, 0, head_dim), dtype=..., device=...) → shape [batch, heads, 0, dim] (4D)

Modified QuantizedLayer.update()

• Applied same fix when resetting cache after quantization
• Ensures empty tensors have proper 4D shape matching key_states dimensions

Testing

The fix ensures:

• torch.cat([empty_4d_tensor, key_states], dim=-2) works correctly
• Compatible with torch.compile tracing
• Maintains backward compatibility with eager mode
• Works for both DynamicLayer and QuantizedLayer caches

Impact

• Fixes regression from v4.52.4 to v4.57.1
• Affects models using DynamicCache with torch.compile (GPT2, and others)
• No breaking changes to API or behavior

[yashwantbezawada]

I see that the CI tests are failing (tests_exotic_models, tests_generate, tests_torch), while code quality checks pass. I'm unable to access the detailed CircleCI logs to understand the specific test failures.

The changes I made:

• Changed empty tensor initialization from torch.tensor([]) (1D) to torch.zeros((batch_size, num_heads, 0, head_dim)) (4D with 0 seq_len)
• Applied this to both DynamicLayer.lazy_initialization and QuantizedLayer.update

This approach ensures torch.cat works correctly in torch.compile mode by providing properly shaped 4D tensors.

Could someone help me understand what tests are failing and why? I'd be happy to adjust the approach if needed. I'm aware of PR #40328 which takes a more comprehensive approach to torch.compile + DynamicCache compatibility.

cc @huggingface/transformers

[Rocketknight1]

This is an update to a PR from @Cyrilvallez, so I'll wait for him to approve it!

[Cyrilvallez]

How are you using torch.compile? In general, it should not really be used with DynamicCache, as the shapes will keep on changing each iteration

[yashwantbezawada]

The original issue was about using a custom torch.compile backend with aot_export_joint_simple for tracing - not really standard generation with DynamicCache.

I noticed PR #40328 is doing the proper fix with symbolic shapes and mark_dynamic. That seems like the right long-term solution. Should I just close this one?

[Cyrilvallez]

Hey @yashwantbezawada! Indeed the PR you linked will provide much better support for compile options! I'll go back into it asap when things will be calmer after the v5 release! Feel free to close in the meantime!

[yashwantbezawada]

Thanks for the update! Closing this one - looking forward to #40328.