Skip to content

[Executorch][llm] Fix flakyness of quantized sdpa test #15653

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: gh/kimishpatel/204/base
Choose a base branch
from
+7 −0
Open

[Executorch][llm] Fix flakyness of quantized sdpa test #15653

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions extension/llm/custom_ops/TARGETS
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,5 +60,6 @@ runtime.python_test(
],
deps = [
"//caffe2:torch",
"//executorch/extension/pybindings:portable_lib",
],
)
6 changes: 6 additions & 0 deletions extension/llm/custom_ops/test_quantized_sdpa.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import torch.nn.functional as F

from executorch.extension.llm.custom_ops import custom_ops # noqa
from executorch.extension.pybindings.portable_lib import _unsafe_reset_threadpool


def is_fbcode():
Expand Down Expand Up @@ -40,6 +41,11 @@ def setUp(self):
self.q_shape = None
self.kv_shape = None
self.is_seq_at_dim_2 = True
# For some reason 4 threads doesnt work
Copy link
Contributor

@metascroy metascroy Nov 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a concern with users choose a thread count that doesn't work?

Copy link
Contributor Author

@kimishpatel kimishpatel Nov 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is specifically related to OMP which I am definitely a bit surprised why it occurs here. This hasnt been observed on-device

# This setting is needed to make this test not flaky due to OMP
# error of "OMP: Error #131: Thread identifier invalid"
# Not clear why that happens but having smaller threadpool resolves it
_unsafe_reset_threadpool(3)

def _scale_tensor(self, tensor, min_value, max_value, scale=True):
normalized_tensor = (tensor - tensor.min()) / (tensor.max() - tensor.min())
Expand Down
Loading