fix: apply dynamo config on GPU worker thread (#167)

mstar/engine/__init__.py

@@ -1,6 +1,22 @@
 import torch
 
-torch._dynamo.config.recompile_limit = 84
-torch._dynamo.config.allow_unspec_int_on_nn_module = True
-torch._dynamo.config.specialize_int = False
-torch.set_float32_matmul_precision('high')
+# torch._dynamo ConfigModule stores user overrides in a ContextVar (thread-local
+# as of torch 2.13+). Import-time assignments only affect the importing thread;
+# the dedicated GPU executor thread that actually compiles never sees them.
+# Call apply_dynamo_config() on any thread that may trigger dynamo tracing.
+RECOMPILE_LIMIT = 84
+
+
+def apply_dynamo_config() -> None:
+    """Apply dynamo settings to the *calling* thread.
+
+    torch's ConfigModule keeps user overrides in a ContextVar, so these are
+    thread-local — any thread that may trigger a compile must call this.
+    """
+    torch._dynamo.config.recompile_limit = RECOMPILE_LIMIT
+    torch._dynamo.config.allow_unspec_int_on_nn_module = True
+    torch._dynamo.config.specialize_int = False
+
+
+apply_dynamo_config()
+torch.set_float32_matmul_precision("high")

mstar/engine/cuda_graph_runner.py

@@ -609,11 +609,14 @@ def _capture_slots(
                 # graph (finished in the submodule ``postprocess``).
                 forward = getattr(submodule, config.capture_forward_method)
                 if config.compile:
+                    from mstar.engine import RECOMPILE_LIMIT
+
                     forward = torch.compile(
                         forward,
                         mode="max-autotune-no-cudagraphs",
                         fullgraph=False,
                         dynamic=False,
+                        recompile_limit=RECOMPILE_LIMIT,
                     )
                 spec.engine_inputs.sampler.applied_penalty_in_graph = False
 
@@ -2184,11 +2187,14 @@ def _capture_one(
 
         fwd = submodule.forward_batched
         if config.compile:
+            from mstar.engine import RECOMPILE_LIMIT
+
             fwd = torch.compile(
                 fwd,
                 mode="max-autotune-no-cudagraphs",
                 fullgraph=False,
                 dynamic=False,
+                recompile_limit=RECOMPILE_LIMIT,
             )
 
         # Warmup and capture on the shared side stream (see warmup_and_capture
@@ -2509,11 +2515,14 @@ def _capture_one(self, shape: PiecewiseCaptureShape) -> None:
 
         fn = self.config.capture_fn
         if self.config.compile:
+            from mstar.engine import RECOMPILE_LIMIT
+
             fn = torch.compile(
                 fn,
                 mode="max-autotune-no-cudagraphs",
                 fullgraph=False,
                 dynamic=False,
+                recompile_limit=RECOMPILE_LIMIT,
             )
 
         def run_fn():

mstar/engine/kv_cache_engine.py

@@ -332,15 +332,19 @@ def _compile_submodules(self) -> None:
                 continue
 
             try:
+                from mstar.engine import RECOMPILE_LIMIT
+
                 submodule.forward = torch.compile(
                     submodule.forward,
                     fullgraph=False,
                     dynamic=None,
+                    recompile_limit=RECOMPILE_LIMIT,
                 )
                 submodule.forward_batched = torch.compile(
                     submodule.forward_batched,
                     fullgraph=False,
                     dynamic=None,
+                    recompile_limit=RECOMPILE_LIMIT,
                 )
                 logger.info("KVCacheEngine: torch.compile applied to %s language_model", node_name)
             except Exception:

mstar/engine/stateless_engine.py

@@ -545,9 +545,12 @@ def _apply_torch_compile(self, node_name: str, submodule: NodeSubmodule) -> None
             return
         try:
             if hasattr(submodule, "forward"):
+                from mstar.engine import RECOMPILE_LIMIT
+
                 submodule.forward = torch.compile(
                     submodule.forward,
                     fullgraph=False,
+                    recompile_limit=RECOMPILE_LIMIT,
                 )
                 logger.info(
                     "StatelessEngine[%s]: torch.compile applied to %s.forward",

mstar/worker/worker.py

@@ -19,6 +19,7 @@
 from mstar.conductor.request_info import CurrentForwardPassInfo
 from mstar.distributed.base import ShardingConfig
 from mstar.distributed.communication import WorkerParallelGroups
+from mstar.engine import apply_dynamo_config
 from mstar.engine.base import EngineType, NodeBatch, NodeOutput
 from mstar.engine.kv_store import KVCacheConfig, StoreWritePolicy, TransferEngineInfo
 from mstar.graph.base import GraphEdge, GraphNode
@@ -1974,8 +1975,12 @@ def run(self) -> None:
         # the main loop can overlap queue/tensor polling and post-processing
         # with GPU execution. Run the engine unconditionally on a dedicated
         # 1-worker GPU thread.
+        # initializer re-applies dynamo config on the GPU thread — ConfigModule
+        # overrides are thread-local (ContextVar) as of torch 2.13+ (#167).
         gpu_executor = ThreadPoolExecutor(
-            max_workers=1, thread_name_prefix=f"mstar-gpu-{self.worker_id}"
+            max_workers=1,
+            thread_name_prefix=f"mstar-gpu-{self.worker_id}",
+            initializer=apply_dynamo_config,
         )
         logger.info(
             "Worker %s: engine runs on dedicated GPU thread",

test/modular/test_dynamo_thread_config.py

@@ -0,0 +1,59 @@
+"""Regression: dynamo config must apply on worker threads (#167).
+
+torch._dynamo ConfigModule stores user overrides in a ContextVar (thread-local
+as of torch 2.13+). Import-time assignments on the main thread do not propagate
+to the dedicated GPU executor thread. apply_dynamo_config() must be called on
+any thread that may trigger dynamo tracing.
+"""
+
+from concurrent.futures import ThreadPoolExecutor
+
+import torch
+
+from mstar.engine import RECOMPILE_LIMIT, apply_dynamo_config
+
+_KEYS = ("recompile_limit", "allow_unspec_int_on_nn_module", "specialize_int")
+
+
+def _read_dynamo_flags() -> dict:
+    return {k: getattr(torch._dynamo.config, k) for k in _KEYS}
+
+
+def test_apply_dynamo_config_on_worker_thread():
+    """GPU-thread initializer path: flags match RECOMPILE_LIMIT / expected defaults."""
+    # Main thread already applied at import; re-apply for a clean baseline.
+    apply_dynamo_config()
+    main = _read_dynamo_flags()
+    assert main["recompile_limit"] == RECOMPILE_LIMIT
+    assert main["allow_unspec_int_on_nn_module"] is True
+    assert main["specialize_int"] is False
+
+    with ThreadPoolExecutor(max_workers=1, initializer=apply_dynamo_config) as ex:
+        worker = ex.submit(_read_dynamo_flags).result()
+
+    assert worker["recompile_limit"] == RECOMPILE_LIMIT
+    assert worker["allow_unspec_int_on_nn_module"] is True
+    assert worker["specialize_int"] is False
+
+
+def test_import_time_flags_do_not_leak_to_fresh_thread_without_initializer():
+    """Document the torch 2.13+ ContextVar pitfall the fix addresses.
+
+    On builds where ConfigModule is thread-local, a fresh thread without
+    apply_dynamo_config() does not see main-thread overrides. Skip when the
+    installed torch still shares config across threads (pre-2.13 behavior).
+    """
+    apply_dynamo_config()
+    main = _read_dynamo_flags()
+
+    with ThreadPoolExecutor(max_workers=1) as ex:
+        worker = ex.submit(_read_dynamo_flags).result()
+
+    if worker == main:
+        # Older torch: config is process-global — nothing to regress against.
+        return
+
+    # Thread-local: worker should have fallen back away from our overrides.
+    assert worker["recompile_limit"] != RECOMPILE_LIMIT or worker[
+        "allow_unspec_int_on_nn_module"
+    ] is not True