----- 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.
Feature request
Currently, the
attention_maskpassed 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.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
/