fix: use current GPU for e4m3 capability

python/freetoken/kernel/triton/e4m3_compat.py

@@ -43,32 +43,27 @@ def _env_force() -> bool:
     os.environ["TRITON_CACHE_DIR"] = os.path.join(
         os.path.expanduser("~/.triton"), "cache-e4m3emu")
 
-_native: bool | None = None
+# The sm_89 (Ada) boundary: fp8e4nv is a native tensor-core type from here up.
+NATIVE_FP8_CAPABILITY = (8, 9)
 
 
 def e4m3_native() -> bool:
     """Host-side twin of :func:`e4m3_native_cx`: True when kernels take fp8e4nv
-    tensors directly. False: pass ``.view(torch.uint8)`` and bf16 act buffers."""
-    global _native
+    tensors directly. False: pass ``.view(torch.uint8)`` and bf16 act buffers.
+
+    Reads the *same* active-driver target as :func:`e4m3_native_cx` (via
+    ``target_info``) on every call, so the host representation and the compiled
+    kernel branch can never disagree. It is deliberately NOT memoized on the
+    first call: a memoized snapshot taken before the worker binds its GPU (or on
+    the default device of a heterogeneous host) would pin a stale convention that
+    contradicts the device the kernel actually compiles for."""
     if _env_force() != FORCE_EMU:
         raise RuntimeError(
             "FREETOKEN_FORCE_E4M3_EMU changed after import: the flag is read once at "
             "import and is not part of triton's compile cache key -- set it before "
             "the process starts (with its own TRITON_CACHE_DIR)"
         )
-    if _native is None:
-        if FORCE_EMU:
-            _native = False
-        else:
-            native = {torch.cuda.get_device_capability(i) >= (8, 9)
-                      for i in range(torch.cuda.device_count())}
-            if len(native) > 1:
-                raise NotImplementedError(
-                    "GPUs on both sides of the sm_89 fp8 boundary in one process: "
-                    "the host-side e4m3 convention is process-global"
-                )
-            _native = native.pop() if native else torch.cuda.get_device_capability() >= (8, 9)
-    return _native
+    return not FORCE_EMU and target_info.cuda_capability_geq(*NATIVE_FP8_CAPABILITY)
 
 
 def e4m3_kernel_view(t: torch.Tensor) -> torch.Tensor:
@@ -89,7 +84,7 @@ def e4m3_native_cx():
     Delegates to ``target_info`` (reads the active driver's target, so
     cross-compilation tests that patch ``driver.active.get_current_target``
     resolve consistently)."""
-    return not FORCE_EMU and target_info.cuda_capability_geq(8, 9)
+    return not FORCE_EMU and target_info.cuda_capability_geq(*NATIVE_FP8_CAPABILITY)
 
 
 @jit

tests/kernels/test_e4m3_device_selection.py

@@ -0,0 +1,67 @@
+import pytest
+
+from freetoken.kernel.triton import e4m3_compat
+
+
+PRE_FP8_ARCH = 80
+NATIVE_FP8_ARCH = 89
+
+
+def _patch_arch(monkeypatch, arch_box):
+    """Point both capability sources at ``arch_box['arch']`` so the test exercises
+    the real decision regardless of which source the implementation reads: the
+    triton active target (``target_info.cuda_capability_geq``, used by the host
+    twin and its compile-time twin) and ``torch.cuda.get_device_capability``."""
+    monkeypatch.setattr(
+        e4m3_compat.target_info,
+        "cuda_capability_geq",
+        lambda major, minor=0: arch_box["arch"] >= major * 10 + minor,
+    )
+    monkeypatch.setattr(
+        e4m3_compat.torch.cuda,
+        "get_device_capability",
+        lambda device=None: (arch_box["arch"] // 10, arch_box["arch"] % 10),
+    )
+
+
+@pytest.mark.parametrize(
+    ("arch", "expected"),
+    [(PRE_FP8_ARCH, False), (NATIVE_FP8_ARCH, True)],
+)
+def test_e4m3_native_matches_active_target(monkeypatch, arch, expected):
+    _patch_arch(monkeypatch, {"arch": arch})
+    assert e4m3_compat.e4m3_native() is expected
+
+
+def test_e4m3_native_tracks_worker_device_not_first_call(monkeypatch):
+    """Regression: the host convention must follow the device the kernel actually
+    compiles for, matching :func:`e4m3_native_cx`. On a heterogeneous host the
+    first ``e4m3_native()`` call can happen before the worker binds its GPU (or
+    while the non-native default device is current); a memoized snapshot taken
+    then would pin a stale convention and force ``e4m3_kernel_view`` /
+    ``e4m3_act_dtype`` to the wrong representation for the compiled kernel."""
+    arch_box = {"arch": PRE_FP8_ARCH}
+    _patch_arch(monkeypatch, arch_box)
+
+    # First host decision, taken before the worker binds its native GPU.
+    assert e4m3_compat.e4m3_native() is False
+
+    # Worker binds its actual compute device (native fp8).
+    arch_box["arch"] = NATIVE_FP8_ARCH
+
+    # The host twin must now agree with the active compile target, not a stale
+    # first-call cache.
+    assert e4m3_compat.e4m3_native() is True
+
+
+def test_e4m3_native_force_emu_short_circuits(monkeypatch):
+    """FREETOKEN_FORCE_E4M3_EMU pins the emulated convention regardless of the
+    device capability."""
+    monkeypatch.setattr(e4m3_compat, "FORCE_EMU", True)
+    monkeypatch.setattr(e4m3_compat, "_env_force", lambda: True)
+    monkeypatch.setattr(
+        e4m3_compat.target_info,
+        "cuda_capability_geq",
+        lambda major, minor=0: True,
+    )
+    assert e4m3_compat.e4m3_native() is False