saymrwulf/transformers
1
0
Fork 0
mirror of https://github.com/saymrwulf/transformers.git synced 2026-05-15 21:01:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Use sdpa_kernel in tests (#35472)

Browse source 
* update: use sdpa_kernel

* update: rerun test
This commit is contained in:
Jacky Lee 2025-01-03 05:39:52 -08:00 committed by GitHub
parent cba49cb2a6
commit 30a9971632
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 2 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
tests/test_modeling_common.py
View file

@ -4255,7 +4255,7 @@ class ModelTesterMixin:
if isinstance(inp, torch.Tensor) and inp.dtype in [torch.float32, torch.float16]:
inputs_dict[name] = inp.to(torch.float16)
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
with sdpa_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
_ = model(**inputs_dict)
@require_non_xpu
@ -4347,11 +4347,7 @@ class ModelTesterMixin:
model_sdpa = model_sdpa.eval()
with torch.no_grad():
with torch.backends.cuda.sdp_kernel(
enable_flash=False,
enable_math=True,
enable_mem_efficient=False,
):
with sdpa_kernel(enable_flash=False, enable_math=True, enable_mem_efficient=False):
res_eager = model_eager(**inputs_dict, return_dict=False)[0]
res_sdpa = model_sdpa(**inputs_dict, return_dict=False)[0]