SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

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
1 changed files with 2 additions and 6 deletions
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]