fix: use current GPU for e4m3 capability

[ousamabenyounes]

Summary

Decide native e4m3 availability from the CUDA device the current worker actually
uses, not by scanning every visible GPU. FreeToken runs one worker process per
GPU, so the old process-wide scan rejected valid heterogeneous hosts (issue #16)
even though each worker only ever touches its assigned device.

Fixes #16

The subtle part: host must match the compiled kernel branch

e4m3_native() is the host twin of the compile-time @constexpr_function
e4m3_native_cx(). The compiled kernels branch on e4m3_native_cx, and the host
side (e4m3_kernel_view / e4m3_act_dtype, plus the torch._scaled_mm path in
fp8_pertensor_linear) must pick the same representation (native fp8 vs the
uint8 view / bf16 act buffers) or you get illegal-fp8-pointer crashes or silent
corruption.

The first cut read torch.cuda.get_device_capability() (current device) and
memoized it into a module global on the first call. That has an ordering/cache
counterexample on a heterogeneous host: if the first e4m3_native() call lands
before the worker binds its GPU (default device 0), it pins a stale convention
that contradicts the device the kernel later compiles for — reintroducing the
exact host/kernel mismatch, just silently.

Fix: e4m3_native() now reads the same active-driver target as
e4m3_native_cx() (target_info.cuda_capability_geq) live on every call, so the
host decision and the compiled branch can never disagree. The _native
memoization is dropped. This keeps one process-global convention (sourced from
the worker's own compile target) — no per-device state. The sm_89 boundary is
extracted to NATIVE_FP8_CAPABILITY.

Test verification (RED → GREEN)

The new regression test_e4m3_native_tracks_worker_device_not_first_call encodes
the counterexample: a heterogeneous host where the first host decision is taken on
the non-native default device, then the worker binds its native GPU.

RED — PR-head code (current-device + memoization):

>       assert e4m3_compat.e4m3_native() is True
E       assert False is True            # stale first-call cache
1 failed, 3 deselected

GREEN — this fix:

....                                        4 passed
ruff: All checks passed!

Changed production lines are covered (the unified target_info read + the
NATIVE_FP8_CAPABILITY constant).

Validation note

The device-selection tests are pure-Python (they patch the capability sources) and
were run for real in a CPU container against the actual e4m3_compat module. The
GPU-only classes in test_e4m3_compat.py (bit-exact primitives, forced-emu A/B,
cross-arch compile gate) require an sm_89+ build host and are unchanged by this PR.

[ousamabenyounes]

Follow-up in c1e23cc: the first version decided native e4m3 from
torch.cuda.get_device_capability() (current device) memoized on first call.
On a heterogeneous one-worker-per-GPU host that first call can land before the
worker binds its GPU, pinning a stale convention that then contradicts the device
the kernel actually compiles for — so the host representation
(e4m3_kernel_view / e4m3_act_dtype) and the compiled e4m3_native_cx branch
could silently disagree.

e4m3_native() now reads the same active-driver target as its compile-time
twin (target_info.cuda_capability_geq) live on every call, and the _native
memoization is gone. One process-global convention, sourced from the worker's own
compile target — no per-device dict.

Validation (RED → GREEN): the new
test_e4m3_native_tracks_worker_device_not_first_call fails on the previous code
(assert False is True, stale cache) and passes after the fix; full
test_e4m3_device_selection.py = 4 passed, ruff clean. GPU-only classes in
test_e4m3_compat.py are unchanged and need an sm_89+ box.

[gdevenyi]

Merged and tested on 2× RTX 6000 Ada (sm_89), CUDA 13.0, torch 2.11. The change itself is sound and the reasoning is right — deriving host e4m3_native from the compile-time target rather than device 0 matters as soon as CUDA_VISIBLE_DEVICES or TP is in play, which is exactly the configuration I am running.

But this PR's own test fails on my hardware, in isolation:

tests/kernels/test_e4m3_compat.py::test_forced_emu_matches_native
E  AssertionError: blk_aq_y: EMU output differs from native
   assert torch.equal(a[k], b[k]), f"{k}: EMU output differs from native"
1 failed, 7 passed

I checked this against the PR branch alone (not my merge): git checkout upstream/pr/48, rebuild, run — same failure. So it is not an interaction with anything else I have merged.

The other 7 tests in the file pass, including the ones using _MMA_TOL. It is specifically blk_aq_y under exact torch.equal that diverges between the forced-emulation and native paths. Given the neighbouring assertions use a tolerance and this one demands bit-equality, my guess is that blk_aq_y legitimately differs in the last bits between the two paths and the exact comparison is too strong — but that is a guess, and you know the kernel. Worth checking whether it passes on your hardware, since if it is sm_89-specific that is a useful data point either way.

Happy to run any diagnostic you want on this box (sm_89, e4m3 native).

[gdevenyi]

Reviewed on sm_89 hardware (2x RTX 6000 Ada). The change is correct and I think it should land.

The change itself

Reading target_info in e4m3_native() is the right call. The host twin and e4m3_native_cx were two answers to the same question derived from two different sources, and only one of them tracked the device the kernel actually compiles for. Collapsing them onto target_info removes the disagreement by construction rather than by discipline, and NATIVE_FP8_CAPABILITY stops the (8, 9) literal from drifting between the two.

test_e4m3_native_tracks_worker_device_not_first_call is the test that earns the change — it names the actual failure (first call lands before the worker binds its GPU) instead of just asserting the new return value. And _patch_arch patching both capability sources is a nice touch: the test asserts the decision, not the implementation, so it stays honest if the source is swapped again.

One thing you dropped

The old code raised on a host with GPUs on both sides of the sm_89 boundary:

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"
    )

That guard is gone, and the new per-call answer is not equivalent to it. In one process spanning, say, an A100 and an L40S, a buffer allocated under e4m3_act_dtype() while the Ada device is current is fp8; a kernel later compiled for the Ampere device receives an fp8 pointer, which triton rejects in dtype.to_ir. Previously that was a named error at the first host decision. Now it surfaces as a compile failure inside an unrelated kernel — or, for e4m3_kernel_view, as a tensor viewed under the wrong convention.

I do not think this blocks the PR: the case is exotic and the bug you are fixing is not. But the guard is cheap to keep, and it is the difference between a one-line diagnosis and a confusing one. Worth a sentence in the docstring at minimum, saying the convention is now per-call and a heterogeneous host is out of contract.

The CI failure is not yours

tests/kernels/test_e4m3_compat.py::test_forced_emu_matches_native fails with blk_aq_y: EMU output differs from native (52 / 18944). I reverted this PR and reproduced it on main on the same hardware, so it is pre-existing, and the fact that it lands on a file you touched is a coincidence.

It is also a real bug, and the emulation is the correct side. Triton lowers fp32 -> tl.float8e4nv as a double-round (fp32 → fp16 RTZ → e4m3): a value a hair above a grid midpoint collapses onto the midpoint, then ties-to-even sends it down. Measured over 2^22 uniform samples in [-448, 448] against torch's RNE cast:

raw   x.to(tl.float8e4nv)        : 16127 / 4194304 differ  (0.3845%)
  of those:  16127 toward zero,  0 away from zero
round_e4m3(x).to(tl.float8e4nv)  :     0 / 4194304 differ

  x=400.199585  triton=384.0  torch=416.0     (midpoint 400, ulp 32)
  x=272.156555  triton=256.0  torch=288.0     (midpoint 272)
  x=136.002075  triton=128.0  torch=144.0     (midpoint 136)

fp_downcast_rounding="rtne" does not change it. So every native-fp8 activation quantizer shrinks ~0.4% of its values by 1 ULP, always toward zero — a 0.02% magnitude bias on the per-token-group quantizer (mean|dequant| / mean|x|: 0.998937 → 0.999169). Small, but one-sided, and free to remove.

round_e4m3's own docstring already warns about "an fp32 -> fp16 -> 3-bit chain [that] double-rounds when the fp16 result lands exactly on an e4m3 tie". The hazard was understood when the emulation was written; only the native branch skipped the guard. The fix is to round onto the grid in fp32 first on both paths, after which the downcast is exact and the lowering's rounding mode stops mattering.

I opened that as #85 (four quantizer sites + a direct regression test), so it does not become your problem. Once it lands, test_e4m3_compat.py goes 9 passed here.