From 55144259cd14f31ab679a6ba22d65535ded41be3 Mon Sep 17 00:00:00 2001 From: Ben Younes <2910651+ousamabenyounes@users.noreply.github.com> Date: Sat, 22 Aug 2026 17:46:25 +0000 Subject: [PATCH 1/2] fix: use current GPU for e4m3 capability --- python/freetoken/kernel/triton/e4m3_compat.py | 9 +------ tests/kernels/test_e4m3_device_selection.py | 24 +++++++++++++++++++ 2 files changed, 25 insertions(+), 8 deletions(-) create mode 100644 tests/kernels/test_e4m3_device_selection.py diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 61d3a0e..9c203ac 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -60,14 +60,7 @@ def e4m3_native() -> bool: 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) + _native = torch.cuda.get_device_capability() >= (8, 9) return _native diff --git a/tests/kernels/test_e4m3_device_selection.py b/tests/kernels/test_e4m3_device_selection.py new file mode 100644 index 0000000..85a3944 --- /dev/null +++ b/tests/kernels/test_e4m3_device_selection.py @@ -0,0 +1,24 @@ +import pytest + +from freetoken.kernel.triton import e4m3_compat + + +PRE_FP8_CAPABILITY = (8, 0) +NATIVE_FP8_CAPABILITY = (8, 9) + + +@pytest.mark.parametrize( + ("current_capability", "expected"), + [(PRE_FP8_CAPABILITY, False), (NATIVE_FP8_CAPABILITY, True)], +) +def test_e4m3_native_uses_current_device(monkeypatch, current_capability, expected): + capabilities = [PRE_FP8_CAPABILITY, NATIVE_FP8_CAPABILITY] + + def get_device_capability(device=None): + return current_capability if device is None else capabilities[device] + + monkeypatch.setattr(e4m3_compat, "_native", None) + monkeypatch.setattr(e4m3_compat.torch.cuda, "device_count", lambda: len(capabilities)) + monkeypatch.setattr(e4m3_compat.torch.cuda, "get_device_capability", get_device_capability) + + assert e4m3_compat.e4m3_native() is expected From c1e23cc105bf30667ec18686fbce84953c53d415 Mon Sep 17 00:00:00 2001 From: Ben Younes <2910651+ousamabenyounes@users.noreply.github.com> Date: Sat, 22 Aug 2026 18:24:47 +0000 Subject: [PATCH 2/2] fix(e4m3): derive host e4m3_native from the compile-time target (#48) e4m3_native() memoized torch.cuda.get_device_capability() of the current device on its first call. On a heterogeneous one-worker-per-GPU host that first call can land before the worker binds its GPU (default device 0), pinning a stale convention that contradicts the device the kernel actually compiles for -- e4m3_kernel_view/e4m3_act_dtype then pick the wrong representation for the compiled e4m3_native_cx branch. Read the same active-driver target as the compile-time twin e4m3_native_cx (target_info.cuda_capability_geq) live on every call, so the host and the kernel branch can never disagree, and drop the _native memoization. One process-global convention, sourced from the worker's own compile target -- no per-device state. Extract the sm_89 boundary to NATIVE_FP8_CAPABILITY. RED -> GREEN: test_e4m3_native_tracks_worker_device_not_first_call fails on the PR-head code (assert False is True -- stale first-call cache) -> 4 passed. Co-Authored-By: Claude Opus 4.8 (1M context) --- python/freetoken/kernel/triton/e4m3_compat.py | 22 +++--- tests/kernels/test_e4m3_device_selection.py | 67 +++++++++++++++---- 2 files changed, 67 insertions(+), 22 deletions(-) diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 9c203ac..3e24cb0 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -43,25 +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() >= (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: @@ -82,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 diff --git a/tests/kernels/test_e4m3_device_selection.py b/tests/kernels/test_e4m3_device_selection.py index 85a3944..5a0d496 100644 --- a/tests/kernels/test_e4m3_device_selection.py +++ b/tests/kernels/test_e4m3_device_selection.py @@ -3,22 +3,65 @@ from freetoken.kernel.triton import e4m3_compat -PRE_FP8_CAPABILITY = (8, 0) -NATIVE_FP8_CAPABILITY = (8, 9) +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( - ("current_capability", "expected"), - [(PRE_FP8_CAPABILITY, False), (NATIVE_FP8_CAPABILITY, True)], + ("arch", "expected"), + [(PRE_FP8_ARCH, False), (NATIVE_FP8_ARCH, True)], ) -def test_e4m3_native_uses_current_device(monkeypatch, current_capability, expected): - capabilities = [PRE_FP8_CAPABILITY, NATIVE_FP8_CAPABILITY] +def test_e4m3_native_matches_active_target(monkeypatch, arch, expected): + _patch_arch(monkeypatch, {"arch": arch}) + assert e4m3_compat.e4m3_native() is expected - def get_device_capability(device=None): - return current_capability if device is None else capabilities[device] - monkeypatch.setattr(e4m3_compat, "_native", None) - monkeypatch.setattr(e4m3_compat.torch.cuda, "device_count", lambda: len(capabilities)) - monkeypatch.setattr(e4m3_compat.torch.cuda, "get_device_capability", get_device_capability) +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) - assert e4m3_compat.e4m3_native() is expected + # 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