Skip to content

Commit 2aab563

Browse files
authored
[SDPA] Fix decode under compile (#50388)
* fix * remove comment
1 parent b55e076 commit 2aab563

1 file changed

Lines changed: 2 additions & 2 deletions

File tree

‎src/transformers/integrations/sdpa_attention.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
from ..generation.continuous_batching.cache import PagedAttentionCache
66
from ..utils import is_torch_mps_available, is_torch_npu_available, is_torch_xpu_available, logging
7-
from ..utils.import_utils import is_torch_greater_or_equal
7+
from ..utils.import_utils import is_torch_greater_or_equal, is_torchdynamo_compiling
88

99

1010
logger = logging.get_logger(__name__)
@@ -129,7 +129,7 @@ def sdpa_attention_forward(
129129
if hasattr(module, "num_key_value_groups") and module.num_key_value_groups > 1:
130130
if use_gqa_in_sdpa(attention_mask, key, value):
131131
sdpa_kwargs["enable_gqa"] = True
132-
elif q_length == 1:
132+
elif q_length == 1 and not is_torchdynamo_compiling():
133133
query = pack_gqa(query, module.num_key_value_groups)
134134
attention_mask = pack_gqa(attention_mask, module.num_key_value_groups)
135135
position_bias = pack_gqa(position_bias, module.num_key_value_groups)

0 commit comments

Comments
 (0)