diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 61d3a0e..3e24cb0 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/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 diff --git a/tests/kernels/test_e4m3_device_selection.py b/tests/kernels/test_e4m3_device_selection.py new file mode 100644 index 0000000..5a0d496 --- /dev/null +++ b/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