[DSV4.1][gfx1250] single launch: thread-local trace slots for FlyDSL >= 0.5

atom/models/deepseek_v41/mono/gfx1250/kernels/device.py

@@ -11,6 +11,7 @@
 from __future__ import annotations
 
 import re as _re
+import threading
 
 import flydsl.expr as fx
 from flydsl._mlir import ir
@@ -297,10 +298,32 @@ def unlikely(b):
     return fx.Boolean(llvm.call_intrinsic(i1, "llvm.expect.i1", [_v(b), f], [], []))
 
 
+class TraceSlot:
+    """One trace-time value private to the tracing thread (used as ``slot[0]``).
+
+    The step kernel sets such slots while it traces (the poll-timeout counter pointer, a stamp hook) and clears them
+    when its body ends. FlyDSL >= 0.5 traces every kernel in its own MLIRContext, and a launcher may trace several
+    ranks' kernels in parallel threads; a module-level ``[None]`` list then hands one thread's IR value to another
+    thread's trace ("operand type from a different MLIRContext"). A thread-local slot keeps each trace's values in
+    its own context and behaves the same when tracing is serial.
+    """
+
+    def __init__(self):
+        self._local = threading.local()
+
+    def __getitem__(self, i):
+        assert i == 0
+        return getattr(self._local, "value", None)
+
+    def __setitem__(self, i, value):
+        assert i == 0
+        self._local.value = value
+
+
 # Poll timeouts made visible. Every bounded poll (cross-rank tag waits, device handoffs) gives up after 2 s and the
 # kernel goes on with whatever it read; a timeout therefore means wrong results. TOFLAG[0] (set by the step kernel to its
 # 'cyc' scratch) counts them: int32 [0] += 1, int32 [1] = the last site id. The runner reads it after a step.
-TOFLAG = [None]
+TOFLAG = TraceSlot()
 POLL_TICKS = 200000000   # a bounded poll gives up after this many clock ticks (~2 s) and records a timeout
 
 

atom/models/deepseek_v41/mono/gfx1250/kernels/stamps.py

@@ -6,7 +6,9 @@
 
 import flydsl.expr as fx
 
-HOOK = [None]
+from atom.models.deepseek_v41.mono.gfx1250.kernels.device import TraceSlot
+
+HOOK = TraceSlot()
 
 
 def point(k):