[None][chore] Make torch compiling and piecewise running flags thread-safe

Signed-off-by: Jonas Li <[email protected]>

Author: longlee0622

tensorrt_llm/_torch/utils.py

@@ -11,8 +11,8 @@
 from tensorrt_llm.math_utils import ceil_div, pad_up
 from tensorrt_llm.quantization.utils import fp4_utils
 
-is_torch_compiling_flag = False
-is_piecewise_running_flag = False
+_torch_compiling = threading.local()
+_piecewise_running = threading.local()
 
 aux_stream_name_list = [
     'Attention',
@@ -46,23 +46,19 @@ class ActivationType(IntEnum):
 
 
 def set_torch_compiling(enable: bool):
-    global is_torch_compiling_flag
-    is_torch_compiling_flag = enable
+    _torch_compiling.flag = enable
 
 
 def is_torch_compiling() -> bool:
-    global is_torch_compiling_flag
-    return is_torch_compiling_flag
+    return getattr(_torch_compiling, 'flag', False)
 
 
 def set_piecewise_running(enable: bool):
-    global is_piecewise_running_flag
-    is_piecewise_running_flag = enable
+    _piecewise_running.flag = enable
 
 
 def is_piecewise_running() -> bool:
-    global is_piecewise_running_flag
-    return is_piecewise_running_flag
+    return getattr(_piecewise_running, 'flag', False)
 
 
 _global_attrs = threading.local()