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!