From 4d79df59567694d1241c23f4ce7db94ae4b9985f Mon Sep 17 00:00:00 2001 From: Kimish Patel Date: Thu, 6 Nov 2025 12:51:07 -0800 Subject: [PATCH] [Executorch][llm] Fix flakyness of quantized sdpa test Without this fix when running diabled tets it fails with OMP: Error #131: Thread identifier invalid Differential Revision: [D86447540](https://our.internmc.facebook.com/intern/diff/D86447540/) [ghstack-poisoned] --- extension/llm/custom_ops/TARGETS | 1 + extension/llm/custom_ops/test_quantized_sdpa.py | 6 ++++++ 2 files changed, 7 insertions(+) diff --git a/extension/llm/custom_ops/TARGETS b/extension/llm/custom_ops/TARGETS index 9a437e7dad5..5dda2318f3f 100644 --- a/extension/llm/custom_ops/TARGETS +++ b/extension/llm/custom_ops/TARGETS @@ -60,5 +60,6 @@ runtime.python_test( ], deps = [ "//caffe2:torch", + "//executorch/extension/pybindings:portable_lib", ], ) diff --git a/extension/llm/custom_ops/test_quantized_sdpa.py b/extension/llm/custom_ops/test_quantized_sdpa.py index 87026d5c251..e6edf6ffbb1 100644 --- a/extension/llm/custom_ops/test_quantized_sdpa.py +++ b/extension/llm/custom_ops/test_quantized_sdpa.py @@ -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(): @@ -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 + # 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())