[generate + static cache + torch.compile] ability to pass statically shaped 4D attention_mask to the model forward

[fxmarty]

Feature request

Currently, the attention_mask passed is 2D and of dynamic shapes.

This causes issues when using a compiled model with model.forward = torch.compile(model.forward, mode="reduce-overhead"), see pytorch/pytorch#120309 & #29114.

----- in forward 0
name=input_ids, shape=torch.Size([2, 7]), stride=(7, 1), dtype=torch.int64, device=cuda:0
name=position_ids, shape=torch.Size([2, 7]), stride=(7, 1), dtype=torch.int64, device=cuda:0
name=cache_position, shape=torch.Size([7]), stride=(1,), dtype=torch.int64, device=cuda:0
name=past_key_values, value=None
name=use_cache, value=True
name=attention_mask, shape=torch.Size([2, 7]), stride=(7, 1), dtype=torch.int64, device=cuda:0
forward call latency: 1784.737 ms      <---------------------------- EXTREMELY SLOW.
----- in forward 1
name=input_ids, shape=torch.Size([2, 1]), stride=(1, 1), dtype=torch.int64, device=cuda:0
name=position_ids, shape=torch.Size([2, 1]), stride=(1, 1), dtype=torch.int64, device=cuda:0
name=cache_position, shape=torch.Size([1]), stride=(1,), dtype=torch.int64, device=cuda:0
name=past_key_values, value=None
name=use_cache, value=True
name=attention_mask, shape=torch.Size([2, 8]), stride=(8, 1), dtype=torch.int64, device=cuda:0
forward call latency: 1851.579 ms      <---------------------------- EXTREMELY SLOW.
----- in forward 2
name=input_ids, shape=torch.Size([2, 1]), stride=(1, 1), dtype=torch.int64, device=cuda:0
name=position_ids, shape=torch.Size([2, 1]), stride=(1, 1), dtype=torch.int64, device=cuda:0
name=cache_position, shape=torch.Size([1]), stride=(1,), dtype=torch.int64, device=cuda:0
name=past_key_values, value=None
name=use_cache, value=True
name=attention_mask, shape=torch.Size([2, 9]), stride=(9, 1), dtype=torch.int64, device=cuda:0
forward call latency: 1421.504 ms      <---------------------------- EXTREMELY SLOW.
----- in forward 3
name=input_ids, shape=torch.Size([2, 1]), stride=(1, 1), dtype=torch.int64, device=cuda:0
name=position_ids, shape=torch.Size([2, 1]), stride=(1, 1), dtype=torch.int64, device=cuda:0
name=cache_position, shape=torch.Size([1]), stride=(1,), dtype=torch.int64, device=cuda:0
name=past_key_values, value=None
name=use_cache, value=True
name=attention_mask, shape=torch.Size([2, 10]), stride=(10, 1), dtype=torch.int64, device=cuda:0
forward call latency: 1740.283 ms      <---------------------------- EXTREMELY SLOW.

Instead, we may want to pass directly 4D masks of static shape to the model, avoiding cuda graphs recomputation. This may allow the compile time to go down from >3 min to a mere 50s.

Motivation

/

Your contribution

/

[amyeroberts]

cc @gante

[gante]

@fxmarty I had a quick look at this -- we still have models (like gpt2) that exclusively accept 2D masks. We would have to rework that before making generate prepare a 4D mask.

We may want to work with padded tensors instead? In our TF/PT XLA implementation we have static shapes everywhere.

[fxmarty]

@gante thank you! What do you mean by work with padded tensors?

[gante]

@fxmarty If we want to generate with max_length=512, the attention mask is always kept with sequence length = 512. Data is moved around inside the tensor as needed.

[fxmarty]

Oh yes, I think that's what I am suggesting here!