feat(server)!: add --gpu to choose the GPU on multi-GPU machines (#117)

* feat(server)!: add --gpu to choose the GPU on multi-GPU machines

--gpu takes a GPU UUID (as nvidia-smi -L prints) or an nvidia-smi index.

It is applied as CUDA_VISIBLE_DEVICES in the parent before the workers spawn, so the engine still binds cuda:<rank>.

/v1/stats reports the engine's GPU under "gpus".

ft bench bw takes the same --gpu and writes one profile per GPU (benchbw/<gpu-uuid>.json); the legacy benchbw.json is still read by GPU name.

The daemon's /bench/profile returns the running serve's GPU profile.

BREAKING CHANGE: ft checkpoint --device is removed; use --gpu.

* refactor!: resolve --gpu via NVML and bind by UUID, not CUDA_VISIBLE_DEVICES

BREAKING CHANGE: ft bench bw --device is removed; use --gpu.

* refactor: name the id namespaces and drop the multi-device e4m3 scan

Splits the published id into _assigned_physical (UUID) and
_assigned_visible (CUDA ordinal). One process runs on one GPU, so
e4m3_native() judges that card instead of scanning every visible device.

benchmarks/bench_decode_moe.py

@@ -100,6 +100,8 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
         help="hybrid: max PCIe fetches/layer; -1 = auto (benched pcie/cpu bandwidth fraction)",
     )
     p.add_argument("--mem-ratio", type=float, default=0.9, help="target VRAM utilization")
+    p.add_argument("--gpu", default=None,
+                   help="GPU for the serve: a UUID or nvidia-smi index (as ft serve --gpu)")
     p.add_argument("--no-graph", action="store_true", help="eager decode instead of CUDA graph")
     p.add_argument(
         "--greedy",
@@ -183,6 +185,8 @@ def serve_cmd(args: argparse.Namespace, backend: str, port: int) -> list[str]:
         "--cuda-graph-max-bs", "0" if args.no_graph else "1",
         "--moe-hybrid-max-fetch", str(args.hybrid_fetch),
     ]
+    if args.gpu:
+        cmd += ["--gpu", args.gpu]
     if args.cache > 0:
         cmd += ["--moe-cache-size", str(args.cache)]
     elif args.cache_rate is not None:

benchmarks/bench_load_weight_generic.py

@@ -140,8 +140,10 @@ def _model_config(model_path: str):
 
     if try_get_tp_info() is None:
         set_tp_info(rank=0, size=1)
-    torch.cuda.set_device(0)
-    torch.zeros(1, device="cuda")  # init CUDA context (pinning / nvfp4 backend pick)
+    from freetoken.gpu_select import bind_assigned_gpu
+
+    dev = bind_assigned_gpu()
+    torch.zeros(1, device=dev)  # init CUDA context (pinning / nvfp4 backend pick)
     cfg = EngineConfig(model_path=model_path, tp_info=DistributedInfo(0, 1),
                        dtype=torch.bfloat16, moe_backend="offload")
     return cfg.model_config
@@ -176,7 +178,7 @@ def _bench_load(mode: str, model_path: str, *, parallel: bool, workers: int, chu
     s.start()
     t = time.perf_counter()
     try:
-        banks = load_expert_banks(model_path, mc, device=torch.device("cuda:0"),
+        banks = load_expert_banks(model_path, mc, device=torch.device("cuda", torch.cuda.current_device()),
                                   dtype=torch.bfloat16, parallel=parallel,
                                   workers=workers, chunk=chunk)
     except NotImplementedError as e:
@@ -220,7 +222,8 @@ def worker_build(ns):
     s = MemSampler()
     s.start()
     t = time.perf_counter()
-    idx = convert_checkpoint(ns.model, ns.ftw_dir, moe_backend="offload", shard_limit=shard_limit)
+    dev = f"cuda:{torch.cuda.current_device()}" if ns.gpu else None
+    idx = convert_checkpoint(ns.model, ns.ftw_dir, moe_backend="offload", shard_limit=shard_limit, device=dev)
     build_s = time.perf_counter() - t
     s.stop()
     print("@@RESULT@@" + json.dumps({
@@ -237,6 +240,8 @@ def _spawn(worker: str, ns):
     cmd = [sys.executable, os.path.abspath(__file__), "--_worker", worker, "--model", ns.model,
            "--workers", str(ns.workers), "--chunk-mib", str(ns.chunk_mib),
            "--shard-gib", str(ns.shard_gib), "--ftw-dir", ns.ftw_dir]
+    if ns.gpu:
+        cmd += ["--gpu", ns.gpu]
     if ns.no_drop_cache:
         cmd.append("--no-drop-cache")
     # Capture ONLY stdout (the @@RESULT@@ line); let stderr inherit the terminal so the
@@ -260,6 +265,10 @@ def main():
     p.add_argument("--no-drop-cache", action="store_true",
                    help="don't evict page cache before each read (warm comparison)")
     p.add_argument("--keep-ftw", action="store_true", help="keep + reuse the FTW dir across runs")
+    from freetoken.gpu_select import single_gpu_arg
+
+    p.add_argument("--gpu", type=single_gpu_arg, default=None,
+                   help="GPU UUID or nvidia-smi index (default: the first visible GPU)")
     p.add_argument("--_worker", default="")
     ns = p.parse_args()
 
@@ -268,7 +277,17 @@ def main():
     if not ns.ftw_dir:
         ns.ftw_dir = _default_ftw_dir(ns.model)
 
+    from freetoken.gpu_select import assign_gpu
+
+    try:
+        assign_gpu(ns.gpu)
+    except ValueError as e:
+        p.error(str(e))
+
     if ns._worker:
+        from freetoken.gpu_select import bind_assigned_gpu
+
+        bind_assigned_gpu()
         return _WORKERS[ns._worker](ns)
 
     from freetoken.checkpoint.ftw import is_ftw_checkpoint

benchmarks/bench_offload_cache_copy.py

@@ -16,6 +16,7 @@ from dataclasses import dataclass
 
 import torch
 
+from freetoken.gpu_select import assign_gpu, bind_assigned_gpu, single_gpu_arg
 from freetoken.moe.offload_cache import _BANK_SCHEMAS, OffloadMoeCache
 
 
@@ -80,7 +81,8 @@ def expert_bytes(profile: ModelProfile) -> int:
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
-    parser.add_argument("--device", type=int, default=0)
+    parser.add_argument("--gpu", type=single_gpu_arg, default=None,
+                        help="GPU UUID or nvidia-smi index (default: the first visible GPU)")
     parser.add_argument("--repeat", type=int, default=25)
     parser.add_argument("--models", type=str, nargs="+", default=list(MODELS), choices=list(MODELS))
     parser.add_argument(
@@ -228,8 +230,11 @@ def print_table(
 def main() -> None:
     args = parse_args()
     assert torch.cuda.is_available(), "CUDA is required"
-    torch.cuda.set_device(args.device)
-    device = torch.device("cuda")
+    try:
+        assign_gpu(args.gpu)
+        device = bind_assigned_gpu()
+    except (ValueError, RuntimeError) as e:
+        raise SystemExit(f"error: {e}") from e
 
     print("gpu", torch.cuda.get_device_name(device), flush=True)
     for name in args.models:

docs/cli.md

@@ -40,6 +40,7 @@ parsers all resolve automatically from the checkpoint and the GPU.
 |---|---|---|
 | `--host` | 127.0.0.1 | Bind address |
 | `--port` | 1919 | Bind port |
+| `--gpu` | GPU 0 | GPU to run on: a UUID from `nvidia-smi -L` or an `nvidia-smi` index; see [below](#choosing-a-gpu) |
 | `--max-running-requests` | 4 | Max concurrently running requests |
 | `--max-output-tokens` | 32768 | Default output budget for requests that omit one |
 | `--max-seq-len-override` | from checkpoint | Max sequence length |
@@ -47,6 +48,21 @@ parsers all resolve automatically from the checkpoint and the GPU.
 | `--cuda-graph-max-bs`, `--graph` | = max running requests | Max batch size captured as CUDA graphs |
 | `--decode-log-interval` | 40 | Scheduler status line every N decode steps |
 
+### Choosing a GPU
+
+For example, a machine with an RTX 5090 and an RTX 3060 Ti:
+
+```console
+$ nvidia-smi -L
+GPU 0: NVIDIA GeForce RTX 3060 Ti (UUID: GPU-2f3a9b1c-8d7e-4a05-b6c1-0e5f9a3d7b42)
+GPU 1: NVIDIA GeForce RTX 5090 (UUID: GPU-9e8d7c6b-5a49-4f13-8207-c1b0a4e6d3f5)
+```
+
+```bash
+ft serve --model ... --gpu 1             # by nvidia-smi index -- the 5090
+ft serve --model ... --gpu GPU-9e8d7c6b  # the same card by UUID (a unique prefix is enough)
+```
+
 ### KV cache & memory
 
 | Flag | Default | Meaning |
@@ -130,7 +146,7 @@ environment so the agent cannot silently fall back to a paid endpoint.
 ## ft checkpoint
 
 ```bash
-ft checkpoint --model <hf_dir> --out <ftw_dir> [--dtype bfloat16] [--moe-backend offload] [--shard-gib 8] [--device cuda:0]
+ft checkpoint --model <hf_dir> --out <ftw_dir> [--dtype bfloat16] [--moe-backend offload] [--shard-gib 8] [--gpu <uuid-or-index>]
 ```
 
 Converts an HF safetensors checkpoint to FTW, FreeToken's self-contained
@@ -142,15 +158,18 @@ keeps them dense for resident serving. See the FTW caveats in
 ## ft bench bw
 
 ```bash
-ft bench bw                       # once per machine
+ft bench bw                       # once per GPU
 ft bench bw --dtype nvfp4,bf16    # only the formats you serve
+ft bench bw --gpu 1               # a specific GPU (UUID or nvidia-smi index, as for ft serve)
 ```
 
-Measures host-RAM vs PCIe bandwidth with the real cpu/offload MoE kernels and
-writes a profile (`~/.cache/freetoken/benchbw.json`) that `ft serve
---moe-backend auto` and `--moe-hybrid-max-fetch -1` read. Profiles are keyed on
-expert format + GPU name, so a profile from different hardware is ignored
-rather than misapplied. Selection flags: `--dtype`, `--model`, `--formats`,
-`--isa`; decision rule: `--threshold` (default 2.0 — recommend hybrid when CPU
-bandwidth > 2× PCIe).
+Measures host-RAM vs PCIe bandwidth with the real cpu/offload MoE kernels and writes a
+profile that `ft serve --moe-backend auto` and `--moe-hybrid-max-fetch -1` then read.
+
+- One profile per GPU, at `~/.cache/freetoken/benchbw/<gpu-uuid>.json`.
+- Keyed on expert format + GPU, so a profile from other hardware is ignored rather than
+  misapplied. An older single `benchbw.json` still counts if its GPU name matches.
+- What to measure: `--dtype`, `--model`, `--formats`, `--isa`.
+- `--threshold` (default 2.0) sets the call: recommend hybrid when CPU bandwidth beats PCIe
+  by that factor.
 

python/freetoken/checkpoint/__main__.py

@@ -1,7 +1,7 @@
 """CLI: convert an HF safetensors checkpoint to a FreeToken Weight (FTW) checkpoint.
 
     ft checkpoint --model <hf_dir> --out <ftw_dir> \
-        [--dtype bfloat16] [--moe-backend offload] [--shard-gib 8]
+        [--dtype bfloat16] [--moe-backend offload] [--shard-gib 8] [--gpu <uuid-or-index>]
 
 The output dir is self-contained: point the server's ``--model`` at it to load via the FTW
 fast path (auto-detected).
@@ -14,6 +14,8 @@ import time
 
 import torch
 
+from freetoken.gpu_select import assign_gpu, bind_assigned_gpu, single_gpu_arg
+
 from .convert import convert_checkpoint
 
 _DTYPES = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}
@@ -27,15 +29,24 @@ def main(argv: list[str] | None = None, prog: str = "freetoken.checkpoint") -> i
     p.add_argument("--moe-backend", default="offload",
                    help="offload (experts -> banks) or e.g. triton (experts stay dense)")
     p.add_argument("--shard-gib", type=float, default=8.0, help="max shard size in GiB")
-    p.add_argument("--device", default=None, help="CUDA device for repack (default cuda:0)")
+    p.add_argument("--gpu", type=single_gpu_arg, default=None,
+                   help="GPU for the repack: a GPU UUID (GPU-xxxx..., as nvidia-smi -L prints) or "
+                        "an nvidia-smi index (default: the first visible GPU)")
     ns = p.parse_args(argv)
 
+    # same as ft serve --gpu: resolve, then bind by UUID at CUDA init
+    try:
+        assign_gpu(ns.gpu)
+        device = f"cuda:{bind_assigned_gpu().index}"
+    except (ValueError, RuntimeError) as e:
+        p.error(str(e))
+
     shard_limit = int(ns.shard_gib * (1 << 30))
     shard_limit -= shard_limit % 4096  # keep aligned
     t = time.perf_counter()
     index = convert_checkpoint(
         ns.model, ns.out, dtype=_DTYPES[ns.dtype],
-        moe_backend=ns.moe_backend, shard_limit=shard_limit, device=ns.device,
+        moe_backend=ns.moe_backend, shard_limit=shard_limit, device=device,
     )
     dt = time.perf_counter() - t
     c = index["counts"]

python/freetoken/daemon/app.py

@@ -58,16 +58,41 @@ class BenchBody(BaseModel):
     args: list[str] = []
 
 
-def _bench_profile_path() -> str:
-    from freetoken.moe.bench_profile import default_profile_path  # torch-free
+def _bench_profile_path(gpu_uuid: str | None) -> str | None:
+    # per-GPU profiles and no torch here: the serve's own card when its --gpu names one, else the newest file
+    from freetoken.moe.bench_profile import default_profile_path, latest_profile_path  # torch-free
 
-    return default_profile_path()
+    if gpu_uuid:
+        path = default_profile_path(gpu_uuid)
+        if os.path.isfile(path):
+            return path
+    return latest_profile_path()
 
 
-def _read_bench_profile() -> dict | None:
-    """The engine host's cached benchbw.json (this is where the serve reads it too), or None."""
+def _serve_gpu_uuid(args: list[str]) -> str | None:
+    """The full UUID a serve's `--gpu` pins, or None when there is none or it cannot be resolved."""
+    for i, a in enumerate(args):
+        val = a[len("--gpu="):] if a.startswith("--gpu=") else (args[i + 1] if a == "--gpu" and i + 1 < len(args) else None)
+        if not val:
+            continue
+        from freetoken.gpu_select import resolve_gpu_uuids
+
+        try:
+            resolved = resolve_gpu_uuids([val])
+        except ValueError:
+            return None
+        if resolved:
+            return resolved[0]
+        # no NVML: a UUID value still keys the profile file (canonical prefix), an index cannot
+        return "GPU-" + val[len("GPU-"):] if val.upper().startswith("GPU-") else None
+    return None
+
+
+def _read_bench_profile(path: str | None) -> dict | None:
+    if path is None:
+        return None
     try:
-        with open(_bench_profile_path()) as f:
+        with open(path) as f:
             return json.load(f)
     except (OSError, ValueError):
         return None
@@ -331,19 +356,23 @@ def build_app(
                 yield _bench_sse("error", {"message": f"failed to spawn bench: {exc}"})
                 return
             tail: collections.deque = collections.deque(maxlen=8)  # last non-progress lines (errors)
+            out_path: str | None = None
             assert proc.stdout is not None
             async for raw in proc.stdout:
                 line = raw.decode(errors="replace").rstrip()
                 prog = _parse_ftbench(line)
                 if prog is not None:
                     yield _bench_sse("progress", prog)
+                elif line.startswith("FTBENCH_OUT "):
+                    out_path = line[len("FTBENCH_OUT "):]
                 elif line:
                     tail.append(line)
             rc = await proc.wait()
             if rc != 0:
                 yield _bench_sse("error", {"message": "\n".join(tail) or f"bench exited {rc}"})
                 return
-            prof = _read_bench_profile()
+            # the file this run wrote (an older engine prints no FTBENCH_OUT: newest file, as before)
+            prof = _read_bench_profile(out_path or _bench_profile_path(None))
             if prof is None:
                 yield _bench_sse("error", {"message": "bench finished but no profile was written"})
             else:
@@ -353,7 +382,23 @@ def build_app(
 
     @app.get("/bench/profile", dependencies=auth)
     async def bench_profile():
-        return await run(proxy_pool, _read_bench_profile)
+        def read() -> dict | None:
+            return _read_bench_profile(_bench_profile_path(serve_gpu_uuid()))
+
+        def serve_gpu_uuid() -> str | None:
+            # the running serve reports the full UUID of its card (/v1/stats gpus); a --gpu given as
+            # a UUID prefix would not match the profile file name
+            st = manager.status()
+            if st.get("running"):
+                try:
+                    gpus = probe.stats(st.get("port") or default_serve_port).get("gpus") or []
+                    if gpus and gpus[0].get("uuid"):
+                        return gpus[0]["uuid"]
+                except Exception:  # noqa: BLE001 -- the arg below is the fallback
+                    pass
+            return _serve_gpu_uuid(manager.serve_args())
+
+        return await run(proxy_pool, read)
 
     return app
 

python/freetoken/daemon/serve_manager.py

@@ -862,6 +862,11 @@ class ServeManager:
         with self._cond:
             return self._child.pid if self._child is not None else None
 
+    def serve_args(self) -> list[str]:
+        """Engine args of the running (or last started) serve."""
+        with self._cond:
+            return list(self._args)
+
     def status(self) -> dict:
         with self._cond:
             child = self._child

python/freetoken/engine/engine.py

@@ -10,6 +10,7 @@ import torch
 from freetoken.attention import AttnType, attention_backend_info, create_attention_backend
 from freetoken.core import Batch, Context, Req, set_global_ctx
 from freetoken.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info
+from freetoken.gpu_select import gpu_identity
 from freetoken.layers import set_rope_device
 from freetoken.models import create_model, load_weight
 from freetoken.moe import create_moe_backend, is_offload_moe_backend
@@ -294,10 +295,11 @@ class Engine:
         assert not torch.cuda.is_initialized()
         set_tp_info(rank=config.tp_info.rank, size=config.tp_info.size)
         _ensure_expandable_segments()  # before the first CUDA allocation below
-        _adjust_config(config)
 
-        self.device = torch.device(f"cuda:{config.tp_info.rank}")
-        torch.cuda.set_device(self.device)
+        from freetoken.gpu_select import bind_assigned_gpu
+
+        self.device = bind_assigned_gpu(config.tp_info.rank)
+        _adjust_config(config)
         torch.manual_seed(42)
         self.stream = torch.cuda.Stream()
         torch.cuda.set_stream(self.stream)
@@ -649,8 +651,10 @@ class Engine:
             return  # explicit fixed cap
         from freetoken.moe.bench_profile import load_hybrid_fetch_fraction
 
-        gpu_name = torch.cuda.get_device_name(self.device) if torch.cuda.is_available() else None
-        fraction = load_hybrid_fetch_fraction(cache.quant_format, gpu_name=gpu_name)
+        gpu_name, gpu_uuid = _profile_gpu(self.device.index)
+        fraction = load_hybrid_fetch_fraction(
+            cache.quant_format, gpu_name=gpu_name, gpu_uuid=gpu_uuid
+        )
         if fraction is None:
             cache.hybrid_max_fetch = 1
             logger.warning_rank0(
@@ -993,6 +997,14 @@ class Engine:
         destroy_distributed()
 
 
+def _profile_gpu(index: "int | None" = None) -> Tuple[str | None, str | None]:
+    """(name, uuid) of visible device ``index`` (default: the current, i.e. bound, device); (None, None) without CUDA."""
+    if not torch.cuda.is_available():
+        return None, None
+    ident = gpu_identity(torch.cuda.current_device() if index is None else index)
+    return ident["name"], ident["uuid"]
+
+
 def _ensure_expandable_segments() -> None:
     """Default the CUDA allocator to expandable segments.
 
@@ -1358,8 +1370,8 @@ def _adjust_config(config: EngineConfig):
         bench_fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16")
         from freetoken.moe.bench_profile import load_backend_recommendation
 
-        gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else None
-        if load_backend_recommendation(bench_fmt, gpu_name=gpu_name) == "hybrid":
+        gpu_name, gpu_uuid = _profile_gpu()
+        if load_backend_recommendation(bench_fmt, gpu_name=gpu_name, gpu_uuid=gpu_uuid) == "hybrid":
             from freetoken.moe.cpu_executor import compiled_extension_supports
 
             _act = getattr(model_config, "hidden_act", "silu")

python/freetoken/gpu_select.py

@@ -0,0 +1,287 @@
+"""--gpu for ft serve / bench bw / checkpoint: resolve entries to GPU UUIDs, bind by UUID at CUDA init.
+
+Three device-id namespaces, converted explicitly:
+- logical: position in the --gpu list == TP rank; each worker takes its own entry by rank and the id ends there.
+- physical: NVML / nvidia-smi order, carried as a GPU UUID (_assigned_physical); not affected by CUDA_VISIBLE_DEVICES.
+- visible: CUDA ordinal in this process (_assigned_visible), what torch.device("cuda", n) means.
+
+The parent resolves --gpu entries to full UUIDs via NVML (resolve_gpu_uuids) and fails fast on a typo.
+Each worker publishes its own entry (set_assigned_gpu / assign_gpu) and binds it when CUDA comes up (bind_assigned_gpu) by matching the UUID against CUDA's visible devices.
+One process runs on one GPU. Binding is unconditional: a process that publishes nothing binds a default ordinal and records it, so assigned_visible_gpu() names that card in every case.
+No process mutates CUDA_VISIBLE_DEVICES, and the UUID match holds under any CUDA_DEVICE_ORDER.
+
+Stdlib only (torch is imported lazily); not under freetoken.utils, which imports transformers.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from typing import Sequence
+
+UUID_PREFIX = "GPU-"
+
+
+def is_gpu_uuid(spec: str) -> bool:
+    return spec[: len(UUID_PREFIX)].upper() == UUID_PREFIX
+
+
+def is_gpu_index(spec: str) -> bool:
+    # not str.isdigit(): that also accepts superscripts and other Unicode digits
+    return spec.isascii() and spec.isdecimal()
+
+
+def _canonical(entry: str) -> str:
+    """A UUID in the exact form the driver matches (upper-case GPU- prefix), an index as-is."""
+    if not (is_gpu_uuid(entry) or is_gpu_index(entry)):
+        raise ValueError(
+            f"{entry!r} is neither a GPU UUID (GPU-xxxx..., as `nvidia-smi -L` prints) "
+            f"nor an nvidia-smi index"
+        )
+    return UUID_PREFIX + entry[len(UUID_PREFIX):] if is_gpu_uuid(entry) else entry
+
+
+def parse_gpu_spec(value: str) -> tuple[str, ...]:
+    """Split a --gpu value; ValueError on a bad entry, an empty value, or a mix of UUIDs and indices."""
+    entries = tuple(_canonical(e.strip()) for e in value.split(",") if e.strip())
+    if not entries:
+        raise ValueError("--gpu needs at least one GPU")
+    if len({is_gpu_uuid(e) for e in entries}) > 1:
+        # the driver parses CUDA_VISIBLE_DEVICES as all-UUID or all-index
+        raise ValueError("--gpu entries must be all UUIDs or all indices")
+    return entries
+
+
+def gpu_arg(value: str) -> tuple[str, ...]:
+    """argparse type for a --gpu list."""
+    try:
+        return parse_gpu_spec(value)
+    except ValueError as exc:
+        raise argparse.ArgumentTypeError(str(exc)) from exc
+
+
+def single_gpu_arg(value: str) -> str:
+    """argparse type for a single-GPU --gpu."""
+    entries = gpu_arg(value)
+    if len(entries) != 1:
+        raise argparse.ArgumentTypeError("takes exactly one GPU")
+    return entries[0]
+
+
+def _nvml_uuids() -> "list[str] | None":
+    """Full GPU UUIDs in physical (nvidia-smi) order, or None when NVML is unavailable.
+
+    Own ctypes loader instead of torch's _raw_device_uuid_nvml: that helper only knows the Linux library name, raises (not None) when the library is missing, and is private API.
+    NVML exports are cdecl on every platform, so CDLL is right on Windows too (same as nvidia-ml-py).
+    None on any failure -- no library, a stub library without the _v2 symbols, WSL, a dead device -- and callers fall back.
+    """
+    import ctypes
+
+    if os.name == "nt":
+        candidates = [
+            "nvml.dll",
+            os.path.join(os.environ.get("SystemRoot", r"C:\\Windows"), "System32", "nvml.dll"),
+            os.path.join(os.environ.get("ProgramFiles", r"C:\\Program Files"), "NVIDIA Corporation", "NVSMI", "nvml.dll"),
+        ]
+    else:
+        candidates = ["libnvidia-ml.so.1"]
+    try:
+        for name in candidates:
+            try:
+                lib = ctypes.CDLL(name)
+                break
+            except OSError:
+                continue
+        else:
+            return None
+        if lib.nvmlInit() != 0:
+            return None
+        try:
+            count = ctypes.c_int()
+            if lib.nvmlDeviceGetCount_v2(ctypes.byref(count)) != 0:
+                return None
+            uuids = []
+            for i in range(count.value):
+                handle = ctypes.c_void_p()
+                if lib.nvmlDeviceGetHandleByIndex_v2(i, ctypes.byref(handle)) != 0:
+                    return None
+                buf = ctypes.create_string_buffer(96)
+                if lib.nvmlDeviceGetUUID(handle, buf, 96) != 0:
+                    return None
+                uuids.append(buf.value.decode("ascii", "replace"))
+            return uuids
+        finally:
+            lib.nvmlShutdown()
+    except (OSError, AttributeError):
+        return None
+
+
+def _match_uuid(spec: str, uuids: "list[str]", where: str) -> str:
+    """The unique full UUID that ``spec`` prefixes, else ValueError."""
+    hits = [u for u in uuids if u.upper().startswith(spec.upper())]
+    if len(hits) != 1:
+        raise ValueError(f"--gpu {spec}: not found or not a unique prefix {where}; run `nvidia-smi -L` to list GPUs")
+    return hits[0]
+
+
+def resolve_gpu_uuids(specs: Sequence[str]) -> "tuple[str, ...] | None":
+    """--gpu entries -> full GPU UUIDs, one per TP rank; raises ValueError on a bad entry.
+
+    A preset CUDA_VISIBLE_DEVICES is a quota to stay inside: an index counts within that list, a UUID must name one of its entries.
+    Returns None when NVML is unavailable -- the worker then interprets the raw entries against CUDA's own enumeration (see bind_assigned_gpu).
+    """
+    specs = parse_gpu_spec(",".join(specs))
+    if len({s.upper() for s in specs}) != len(specs):
+        raise ValueError(f"--gpu {','.join(specs)}: the same GPU appears twice")
+    uuids = _nvml_uuids()
+    if uuids is None:
+        return None
+    preset_raw = os.environ.get("CUDA_VISIBLE_DEVICES")
+    preset = None if preset_raw is None else [e.strip() for e in preset_raw.split(",") if e.strip()]
+
+    resolved: list[str] = []
+    for spec in specs:
+        if preset is None:
+            if is_gpu_uuid(spec):
+                resolved.append(_match_uuid(spec, uuids, "on this machine"))
+            elif int(spec) < len(uuids):
+                resolved.append(uuids[int(spec)])
+            else:
+                raise ValueError(f"--gpu {spec}: only {len(uuids)} GPU(s) on this machine; run `nvidia-smi -L` to list GPUs")
+        else:
+            entry = _preset_entry(spec, preset, preset_raw)
+            # an integer entry is read in physical order, as under CUDA_DEVICE_ORDER=PCI_BUS_ID; a negative or MIG-form entry cannot name a whole GPU
+            if is_gpu_uuid(entry):
+                resolved.append(_match_uuid(entry, uuids, f"(from CUDA_VISIBLE_DEVICES={preset_raw!r})"))
+            elif is_gpu_index(entry) and int(entry) < len(uuids):
+                resolved.append(uuids[int(entry)])
+            else:
+                raise ValueError(
+                    f"--gpu {spec}: cannot resolve CUDA_VISIBLE_DEVICES entry {entry!r} "
+                    f"({len(uuids)} GPU(s) on this machine)"
+                )
+    if len(set(resolved)) != len(resolved):
+        raise ValueError(f"--gpu {','.join(specs)}: the same GPU appears twice")
+    return tuple(resolved)
+
+
+def _preset_entry(spec: str, preset: "list[str]", preset_raw: str) -> str:
+    """The CUDA_VISIBLE_DEVICES entry ``spec`` selects, else ValueError."""
+    if not is_gpu_uuid(spec):
+        idx = int(spec)
+        if idx >= len(preset):
+            raise ValueError(
+                f"--gpu {spec}: only {len(preset)} GPU(s) are visible through "
+                f"CUDA_VISIBLE_DEVICES={preset_raw!r} (indices count within that list)"
+            )
+        return preset[idx]
+    if not all(is_gpu_uuid(p) for p in preset):
+        raise ValueError(
+            f"--gpu {spec}: CUDA_VISIBLE_DEVICES={preset_raw!r} lists GPUs by index; "
+            f"give --gpu as an index into that list"
+        )
+    hits = [p for p in preset if p.upper().startswith(spec.upper()) or spec.upper().startswith(p.upper())]
+    if len(hits) != 1:
+        raise ValueError(
+            f"--gpu {spec}: not one of the GPUs visible through CUDA_VISIBLE_DEVICES={preset_raw!r}"
+        )
+    return hits[0]
+
+
+# The GPU this process was assigned, in whichever namespace it arrived in; bind_assigned_gpu fills in the visible one.
+# Process-global on purpose: publishing is torch-free so a worker can do it before heavy imports, and kernel-compat checks (e4m3_native) need the device this process will use.
+_assigned_physical: "str | None" = None
+_assigned_visible: "int | None" = None
+
+
+def set_assigned_gpu(target: str) -> None:
+    """Publish this process's GPU before CUDA init; second call must agree.
+
+    A UUID names a physical GPU and is converted at bind time; a bare index is already a visible ordinal (a preset CUDA_VISIBLE_DEVICES has narrowed to it).
+    """
+    global _assigned_physical, _assigned_visible
+    physical = target if is_gpu_uuid(target) else None
+    visible = None if physical is not None else int(target)
+    current = (_assigned_physical, _assigned_visible)
+    if current not in ((None, None), (physical, visible)):
+        raise RuntimeError(f"set_assigned_gpu called twice: {current} then {target!r}")
+    _assigned_physical, _assigned_visible = physical, visible
+
+
+def assign_gpu(spec: "str | None") -> None:
+    """Resolve one --gpu value and publish it for bind_assigned_gpu; no-op when the flag was not given."""
+    if spec is None:
+        return
+    resolved = resolve_gpu_uuids([spec])
+    set_assigned_gpu(resolved[0] if resolved else parse_gpu_spec(spec)[0])
+
+
+def _visible_of_physical(uuid: str) -> int:
+    """CUDA ordinal of the physical GPU ``uuid`` (or unique prefix) among this process's visible devices."""
+    import torch
+
+    seen: list[str] = []
+    hits: list[int] = []
+    for v in range(torch.cuda.device_count()):
+        u = format_gpu_uuid(getattr(torch.cuda.get_device_properties(v), "uuid", None))
+        seen.append(u or "?")
+        if u is not None and u.upper().startswith(uuid.upper()):
+            hits.append(v)
+    if len(hits) == 1:
+        return hits[0]
+    if hits:
+        raise RuntimeError(f"--gpu {uuid}: not a unique prefix (visible: {', '.join(seen)})")
+    raise RuntimeError(
+        f"GPU {uuid} is not visible to CUDA in this process "
+        f"(CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')!r}, "
+        f"visible: {', '.join(seen) or 'none'})"
+    )
+
+
+def bind_assigned_gpu(default: int = 0):
+    """torch.cuda.set_device this process's GPU and return the device.
+
+    ``default`` is a visible ordinal, used and recorded when nothing was published, so the process always knows which card it runs on.
+    A published UUID (or prefix) is matched against CUDA's own device list, so the result is right under any CUDA_DEVICE_ORDER.
+    """
+    global _assigned_visible
+    import torch
+
+    if _assigned_visible is None:
+        _assigned_visible = default if _assigned_physical is None else _visible_of_physical(_assigned_physical)
+    if not 0 <= _assigned_visible < torch.cuda.device_count():
+        raise RuntimeError(
+            f"cannot use CUDA device {_assigned_visible}: only {torch.cuda.device_count()} device(s) visible "
+            f"(CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')!r})"
+        )
+    device = torch.device("cuda", _assigned_visible)
+    torch.cuda.set_device(device)
+    return device
+
+
+def assigned_visible_gpu() -> "int | None":
+    """Visible ordinal this process is pinned to, or None before it publishes or binds a GPU (= the current device).
+
+    Published-but-not-yet-bound still counts: compat checks in the window between publish and bind must judge the assigned card, not whatever the calling thread happens to sit on.
+    """
+    if _assigned_visible is not None:
+        return _assigned_visible
+    return None if _assigned_physical is None else _visible_of_physical(_assigned_physical)
+
+
+def format_gpu_uuid(raw) -> str | None:
+    """nvidia-smi form GPU-<uuid> from a uuid.UUID."""
+    return None if raw is None else f"{UUID_PREFIX}{raw}"
+
+
+def gpu_identity(index: int) -> dict:
+    """{index, name, uuid, total_bytes} of visible device ``index``."""
+    import torch
+
+    props = torch.cuda.get_device_properties(index)
+    return {
+        "index": index,
+        "name": props.name,
+        "uuid": format_gpu_uuid(getattr(props, "uuid", None)),
+        "total_bytes": int(props.total_memory),
+    }

python/freetoken/kernel/fla/utils.py

@@ -271,13 +271,13 @@ is_intel = device_platform == "intel"
 is_nvidia = device_platform == "nvidia"
 is_intel_alchemist = is_intel and "Intel(R) Arc(TM) A" in torch.xpu.get_device_name(0)
 is_nvidia_hopper = is_nvidia and (
-    "NVIDIA H" in torch.cuda.get_device_name(0)
+    "NVIDIA H" in torch.cuda.get_device_name(torch.cuda.current_device())
     or torch.cuda.get_device_capability()[0] >= 9
 )
 use_cuda_graph = is_nvidia and os.environ.get("FLA_USE_CUDA_GRAPH", "0") == "1"
 
 # Nvidia Ampere or newer, haven't check AMD and intel yet.
-is_tf32_supported = is_nvidia and torch.cuda.get_device_capability(0)[0] >= 8
+is_tf32_supported = is_nvidia and torch.cuda.get_device_capability()[0] >= 8
 is_gather_supported = hasattr(triton.language, "gather")
 
 
@@ -309,8 +309,10 @@ class Backend(Enum):
 
 
 @lru_cache(maxsize=None)
-def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool:
+def check_shared_mem(arch: str = "none", tensor_idx: "int | None" = None) -> bool:
     try:
+        if tensor_idx is None:
+            tensor_idx = device_torch_lib.current_device()
         device_shared_mem_list = get_all_max_shared_mem()
         max_shared_memory = device_shared_mem_list[tensor_idx]
         return max_shared_memory >= Backend.get_shared_memory(arch)

python/freetoken/kernel/triton/e4m3_compat.py

@@ -60,14 +60,10 @@ 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)
+            from freetoken.gpu_select import assigned_visible_gpu
+
+            # one process runs on one GPU, so its convention is that GPU's; None (-> the current device) only before the process binds
+            _native = torch.cuda.get_device_capability(assigned_visible_gpu()) >= (8, 9)
     return _native
 
 

python/freetoken/kernel/triton/sampling.py

@@ -28,7 +28,7 @@ import triton.language as tl
 
 from freetoken.kernel.triton.autotune_cache import autotune_cache_kwargs
 
-_NUM_SM = torch.cuda.get_device_properties(0).multi_processor_count
+_NUM_SM = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
 _MIN_CHUNK = 4096  # do not split a row finer than this
 
 

python/freetoken/moe/bench_profile.py

@@ -1,4 +1,4 @@
-"""Torch-free reader for the ``ft bench bw`` hardware profile (``benchbw.json``).
+"""Torch-free reader for the ``ft bench bw`` hardware profile (``benchbw/<gpu-uuid>.json``).
 
 The engine consults this at MoE-backend *auto* resolution (``engine.py``) to make the
 offload-vs-hybrid choice hardware-adaptive without importing the (torch-heavy) benchmark
@@ -31,10 +31,42 @@ _QUANT_TO_BENCH_FORMAT = {
 }
 
 
-def default_profile_path() -> str:
-    """``$XDG_CACHE_HOME/freetoken/benchbw.json`` (mirrors ``benchbw.default_out_path``)."""
+def _cache_dir() -> str:
     cache = os.environ.get("XDG_CACHE_HOME") or os.path.expanduser("~/.cache")
-    return os.path.join(cache, "freetoken", "benchbw.json")
+    return os.path.join(cache, "freetoken")
+
+
+def default_profile_path(gpu_uuid: str | None = None) -> str:
+    """``$XDG_CACHE_HOME/freetoken/benchbw/<gpu-uuid>.json``, or the legacy ``benchbw.json`` without a uuid.
+
+    One file per GPU: bandwidth differs between slots.
+    """
+    if gpu_uuid:
+        return os.path.join(_cache_dir(), "benchbw", f"{gpu_uuid}.json")
+    return os.path.join(_cache_dir(), "benchbw.json")
+
+
+def latest_profile_path() -> str | None:
+    """Newest ``benchbw/*.json``, else the legacy ``benchbw.json``, else None."""
+    per_gpu = os.path.join(_cache_dir(), "benchbw")
+    newest: tuple[float, str] | None = None
+    try:
+        for name in os.listdir(per_gpu):
+            if not name.endswith(".json"):
+                continue
+            path = os.path.join(per_gpu, name)
+            try:
+                mtime = os.path.getmtime(path)
+            except OSError:
+                continue
+            if newest is None or mtime > newest[0]:
+                newest = (mtime, path)
+    except OSError:
+        pass
+    if newest is not None:
+        return newest[1]
+    legacy = default_profile_path()
+    return legacy if os.path.isfile(legacy) else None
 
 
 def _load(path: str) -> dict | None:
@@ -45,14 +77,28 @@ def _load(path: str) -> dict | None:
         return None
 
 
-def _usable_profile(gpu_name: str | None, path: str | None) -> dict | None:
+def _usable_profile(
+    gpu_name: str | None, path: str | None, gpu_uuid: str | None = None
+) -> dict | None:
     """The cached profile, or ``None`` when there is no file / it was benched on another GPU
     (bandwidths are hardware-specific, so a mismatch is ignored rather than trusted).
 
-    ``path`` overrides the profile location (else ``FREETOKEN_BENCHBW_PATH`` then the default).
+    Lookup: explicit ``path`` (else ``FREETOKEN_BENCHBW_PATH``) -> ``benchbw/<gpu_uuid>.json`` -> legacy ``benchbw.json``.
     """
-    src = path or os.environ.get("FREETOKEN_BENCHBW_PATH") or default_profile_path()
-    prof = _load(src)
+    explicit = path or os.environ.get("FREETOKEN_BENCHBW_PATH")
+    if explicit:
+        candidates = [explicit]
+    else:
+        candidates = [default_profile_path(gpu_uuid)] if gpu_uuid else []
+        candidates.append(default_profile_path())
+    prof = None
+    for src in candidates:
+        prof = _load(src)
+        if isinstance(prof, dict):
+            break
+        if os.path.exists(src):
+            # unreadable profile for this card: stay on the safe default, do not borrow the legacy file
+            return None
     if not isinstance(prof, dict):
         return None
     prof_gpu = (prof.get("gpu") or {}).get("name")
@@ -66,7 +112,10 @@ def _usable_profile(gpu_name: str | None, path: str | None) -> dict | None:
 
 
 def load_backend_recommendation(
-    quant_format: str, gpu_name: str | None = None, path: str | None = None
+    quant_format: str,
+    gpu_name: str | None = None,
+    path: str | None = None,
+    gpu_uuid: str | None = None,
 ) -> str | None:
     """Bench-recommended offload-family backend for ``quant_format`` on this GPU, or ``None``.
 
@@ -77,7 +126,7 @@ def load_backend_recommendation(
     default (offload) on ``None``.
     """
     fmt = _QUANT_TO_BENCH_FORMAT.get(quant_format, quant_format)
-    prof = _usable_profile(gpu_name, path)
+    prof = _usable_profile(gpu_name, path, gpu_uuid)
     if prof is None:
         return None
 
@@ -105,7 +154,10 @@ def load_backend_recommendation(
 
 
 def load_hybrid_fetch_fraction(
-    quant_format: str, gpu_name: str | None = None, path: str | None = None
+    quant_format: str,
+    gpu_name: str | None = None,
+    path: str | None = None,
+    gpu_uuid: str | None = None,
 ) -> float | None:
     """Benched hybrid fetch fraction for ``quant_format``, or ``None``.
 
@@ -119,7 +171,7 @@ def load_hybrid_fetch_fraction(
     ``None`` = no usable profile; clamped to [0, 1].
     """
     fmt = _QUANT_TO_BENCH_FORMAT.get(quant_format, quant_format)
-    prof = _usable_profile(gpu_name, path)
+    prof = _usable_profile(gpu_name, path, gpu_uuid)
     if prof is None:
         return None
     entries = [(prof.get("dtype_kernels") or {}).get(fmt)] + [

python/freetoken/moe/benchbw.py

@@ -21,13 +21,15 @@ the CPU, so they always resolve to ``offload``.
 The hybrid-vs-offload choice is dtype-dominated, so the default is a **per-dtype tuning bench**
 (``--dtype``): one bench per expert format against a canonical geometry -- the minimal set the
 runtime backend pick matches on. ``--model`` additionally benches specific model geometries for
-per-model detail. Results are written to a JSON file (default ``$XDG_CACHE_HOME/freetoken/
-benchbw.json``) so the choice is reproducible.
+per-model detail. Results are written to a JSON file per GPU (default ``$XDG_CACHE_HOME/
+freetoken/benchbw/<gpu-uuid>.json``) so the choice is reproducible and a multi-GPU box keeps
+one profile per card.
 
     ft bench bw                              # per-dtype tuning bench (default: all formats)
     ft bench bw --dtype nvfp4,bf16           # only these formats
     ft bench bw --model qwen3.6-moe          # per-model detail instead
     ft bench bw --dtype all --model all      # both
+    ft bench bw --gpu GPU-2f3a...            # bench a specific GPU (UUID or nvidia-smi index)
 """
 
 from __future__ import annotations
@@ -47,6 +49,12 @@ from types import SimpleNamespace
 
 import torch
 
+from freetoken.gpu_select import (
+    assign_gpu,
+    bind_assigned_gpu,
+    gpu_identity,
+    single_gpu_arg,
+)
 from freetoken.kernel.pinned import alloc_pinned_tensor
 from freetoken.moe.cpu_executor import physical_core_cpus, resolve_threads_and_affinity
 from freetoken.utils import init_logger
@@ -134,11 +142,11 @@ DTYPE_WORKLOADS: dict[str, Workload] = {
 }
 
 
-def default_out_path() -> str:
+def default_out_path(gpu_uuid: str | None = None) -> str:
     # Single source of truth with the (torch-free) reader the engine consults.
     from freetoken.moe.bench_profile import default_profile_path
 
-    return default_profile_path()
+    return default_profile_path(gpu_uuid)
 
 
 def _cgroup_mem_headroom() -> int | None:
@@ -696,12 +704,17 @@ def run_benchbw(
             "benchbw needs a CUDA device to measure PCIe bandwidth (both offload and "
             "hybrid serve experts to the GPU)."
         )
+    if torch.cuda.device_count() == 0:
+        raise RuntimeError(
+            f"no CUDA device visible (CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')!r})."
+        )
     if not 0 <= device_index < torch.cuda.device_count():
         raise RuntimeError(
-            f"--device {device_index} out of range (found {torch.cuda.device_count()} CUDA devices)."
+            f"device_index {device_index} out of range (found {torch.cuda.device_count()} CUDA devices)."
         )
     device = torch.device("cuda", device_index)
     torch.cuda.set_device(device)
+    gpu = gpu_identity(device_index)
 
     # Machine-parseable progress on stdout (opt-in), so the daemon/Desktop can stream feedback
     # while the bench runs -- mirrors ft checkpoint's FTCONVERT lines. `done`/`total` count the
@@ -766,7 +779,8 @@ def run_benchbw(
         "timestamp": datetime.now(timezone.utc).astimezone().isoformat(timespec="seconds"),
         "epoch": int(time.time()),
         "host": socket.gethostname(),
-        "gpu": {"index": device_index, "name": torch.cuda.get_device_name(device)},
+        # index is the CUDA ordinal the bench ran on; uuid keys the profile file
+        "gpu": {"index": device_index, "name": gpu["name"], "uuid": gpu["uuid"]},
         "cpu": {"physical_cores": len(physical_core_cpus()), "threads_used": cpu["threads"]},
         "threshold": threshold,
         "ceilings": {
@@ -780,9 +794,11 @@ def run_benchbw(
         "workloads": workloads_out,
     }
 
-    out_path = os.path.expanduser(out_path or default_out_path())
+    out_path = os.path.expanduser(out_path or default_out_path(gpu["uuid"]))
     _atomic_write_json(out_path, result)
     result["out_path"] = out_path
+    if prog_on:
+        print(f"FTBENCH_OUT {out_path}", flush=True)
     return result
 
 
@@ -927,10 +943,12 @@ def main(argv: list[str] | None = None, prog: str = "ft bench bw") -> int:
                    help=f"CPU MoE ISA: 'auto' (default, best), 'all', or a subset of "
                         f"{list(_ISA_TIERS)} to sweep (kernel caps down to hw support)")
     p.add_argument("-o", "--out", default=None,
-                   help=f"JSON output path (default {default_out_path()})")
+                   help=f"JSON output path (default {default_out_path('<gpu-uuid>')})")
     p.add_argument("--threshold", type=_positive_float, default=2.0,
                    help="recommend hybrid when CPU BW > threshold x PCIe BW (default 2.0)")
-    p.add_argument("--device", type=_nonneg_int, default=0, help="CUDA device index (default 0)")
+    p.add_argument("--gpu", type=single_gpu_arg, default=None,
+                   help="GPU to bench: a GPU UUID (GPU-xxxx..., as nvidia-smi -L prints) or an "
+                        "nvidia-smi index (default: the first visible GPU)")
     p.add_argument("--cpu-threads", type=_nonneg_int, default=0,
                    help="CPU worker threads (0 = one per physical core)")
     p.add_argument("--cpu-iters", type=_positive_int, default=8, help="STREAM read passes to time")
@@ -942,6 +960,13 @@ def main(argv: list[str] | None = None, prog: str = "ft bench bw") -> int:
                    help="fast_index_copy gather passes to time")
     ns = p.parse_args(argv)
 
+    # same as ft serve --gpu: resolve, then bind by UUID at CUDA init
+    try:
+        assign_gpu(ns.gpu)
+        device_index = bind_assigned_gpu().index
+    except (ValueError, RuntimeError) as e:
+        p.error(str(e))
+
     models = ns.model or ()
     dtypes = ns.dtype
     # Default (no selection): the full per-dtype tuning bench -- the minimal set the runtime
@@ -951,7 +976,7 @@ def main(argv: list[str] | None = None, prog: str = "ft bench bw") -> int:
 
     try:
         result = run_benchbw(
-            out_path=ns.out, threshold=ns.threshold, device_index=ns.device, models=models,
+            out_path=ns.out, threshold=ns.threshold, device_index=device_index, models=models,
             dtypes=dtypes, formats=ns.formats, isas=ns.isas, cpu_threads=ns.cpu_threads,
             cpu_iters=ns.cpu_iters, pcie_bytes=ns.pcie_mib << 20, pcie_iters=ns.pcie_iters,
             kernel_cpu_iters=ns.kernel_cpu_iters, kernel_pcie_iters=ns.kernel_pcie_iters,

python/freetoken/moe/host_banks.py

@@ -292,10 +292,14 @@ class PinPipeline:
     def __init__(self) -> None:
         self._q: queue.SimpleQueue = queue.SimpleQueue()
         self._exc: BaseException | None = None
+        # the current device is thread-local: a fresh thread sits on device 0 and cudaHostRegister would build its context there -- carry the creator's (bound) device into the worker
+        self._device = torch.cuda.current_device() if torch.cuda.is_available() else None
         self._thread = threading.Thread(target=self._run, daemon=True)
         self._thread.start()
 
     def _run(self) -> None:
+        if self._device is not None:
+            torch.cuda.set_device(self._device)
         while True:
             item = self._q.get()
             if item is None:

python/freetoken/scheduler/scheduler.py

@@ -6,6 +6,7 @@ import torch
 from freetoken.attention.linear import build_fla_metadata
 from freetoken.core import Batch, Req
 from freetoken.env import ENV
+from freetoken.gpu_select import gpu_identity
 from freetoken.message import (
     AbortBackendMsg,
     BaseBackendMsg,
@@ -68,6 +69,8 @@ class Scheduler(SchedulerIOMixin):
         self.stream = torch.cuda.Stream(device=self.device)
         self.engine_stream_ctx = torch.cuda.stream(self.engine.stream)
         torch.cuda.set_stream(self.stream)
+        # sent on the readiness ack for /v1/stats gpus; a list so TP can add one entry per rank
+        self.gpus = [gpu_identity(self.device.index)] if self.device.type == "cuda" else []
 
         # initialize other managers
         self.table_manager = TableManager(config.max_running_req, self.engine.page_table)

python/freetoken/server/api_server.py

@@ -174,6 +174,8 @@ class FrontendManager:
     # "num_mamba_slots"}, from the same ack. Seeds geometry before the first generation reply
     # (the running snapshot channel) has anything. None until meta arrives.
     cache_pools: Dict[str, int] | None = None
+    # one {index, name, uuid, total_bytes} per TP rank, from the same ack; /v1/stats gpus
+    gpus: List[Dict[str, Any]] = field(default_factory=list)
     # Backend worker Process handles (TP schedulers + tokenizer/detokenizer), captured from the
     # BackendHandle after start_backend(). The orderly-shutdown path (lifespan / shell signal
     # handler) tears these down itself, AFTER setting _SHUTTING_DOWN, so the supervisor observes
@@ -1010,6 +1012,7 @@ def run_api_server(config: ServerArgs, start_backend: Callable[[], "Any"], run_s
         _GLOBAL_STATE.cache_pools = meta.pop("pools", None)
         _GLOBAL_STATE.swa_full_tokens_ratio = float(meta.pop("swa_full_tokens_ratio", 0.0) or 0.0)
         _GLOBAL_STATE.cache_budget_bytes = int(meta.pop("cache_budget_bytes", 0) or 0)
+        _GLOBAL_STATE.gpus = list(meta.pop("gpus", None) or [])
         _GLOBAL_STATE.unit_bytes = meta
 
     # Early-bind: supervise the backend on a daemon thread so uvicorn can bind

python/freetoken/server/args.py

@@ -39,6 +39,10 @@ class ServerArgs(SchedulerConfig):
     # Comma-separated CORS allow-list for browser/webview clients (e.g. the desktop
     # app). Empty string disables CORS headers entirely; "*" allows any origin.
     cors_origins: str = "tauri://localhost,http://tauri.localhost,http://localhost:1420"
+    # --gpu entries in TP-rank order, empty = not given
+    gpu: tuple[str, ...] = ()
+    # full UUIDs resolved from --gpu, entry i = TP rank i; None = NVML unavailable, each worker then resolves its raw entry against CUDA's own enumeration
+    gpu_assigned: "tuple[str, ...] | None" = None
 
     @property
     def share_tokenizer(self) -> bool:
@@ -109,6 +113,11 @@ def parse_args(
             raise argparse.ArgumentTypeError("must be >= 1")
         return n
 
+    def _lazy_gpu_arg(value: str) -> tuple[str, ...]:
+        from freetoken.gpu_select import gpu_arg
+
+        return gpu_arg(value)
+
     def _infer_tool_call_parser(model_path: str) -> str:
         try:
             from freetoken.utils import cached_load_hf_config
@@ -221,6 +230,16 @@ def parse_args(
         help="The tensor parallelism size.",
     )
 
+    parser.add_argument(
+        "--gpu",
+        type=_lazy_gpu_arg,
+        default=ServerArgs.gpu,
+        help=(
+            "GPU(s) to run on, comma-separated; entry i is TP rank i. Each entry is a GPU "
+            "UUID (GPU-xxxx..., as nvidia-smi -L prints) or an nvidia-smi index"
+        ),
+    )
+
     parser.add_argument(
         "--max-running-requests",
         type=int,
@@ -611,6 +630,15 @@ def parse_args(
     # Parse arguments
     kwargs = parser.parse_args(args).__dict__.copy()
 
+    # reject a too-long list here with a clear reason, not as a dead rank later
+    if len(kwargs["gpu"]) not in (0, kwargs["tensor_parallel_size"]):
+        if kwargs["tensor_parallel_size"] == 1 and len(kwargs["gpu"]) > 1:
+            parser.error("tensor parallelism is not supported yet: --gpu takes one entry")
+        parser.error(
+            f"--gpu has {len(kwargs['gpu'])} entries but --tensor-parallel-size is "
+            f"{kwargs['tensor_parallel_size']}; give one entry per TP rank"
+        )
+
     # resolve some arguments
     run_shell |= kwargs.pop("shell_mode")
     kwargs["shell_mode"] = run_shell

python/freetoken/server/launch.py

@@ -59,6 +59,13 @@ def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None:
     if args.shell_mode:
         _detach_process_group()
 
+    # published (not bound) here: the engine binds it after the allocator setup
+    from freetoken.gpu_select import set_assigned_gpu
+
+    # resolved UUIDs when we have them, the raw --gpu entries when NVML could not resolve them, else one CUDA ordinal per rank
+    targets = args.gpu_assigned or args.gpu or tuple(str(r) for r in range(args.tp_info.size))
+    set_assigned_gpu(targets[args.tp_info.rank])
+
     import torch
     from freetoken.scheduler import Scheduler
 
@@ -90,7 +97,10 @@ def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None:
             try:
                 from freetoken.kvcache.cache_status import compute_cache_status_meta
 
-                ack_queue.put(("meta", compute_cache_status_meta(scheduler.engine)))
+                meta = compute_cache_status_meta(scheduler.engine)
+                # the parent must not touch CUDA to learn this
+                meta["gpus"] = scheduler.gpus
+                ack_queue.put(("meta", meta))
             except Exception:  # noqa: BLE001 -- metadata is a nicety; readiness is not
                 pass
             ack_queue.put("Scheduler is ready")
@@ -127,6 +137,19 @@ def launch_server(
     )
     logger = init_logger(__name__, "initializer")
 
+    if server_args.gpu:
+        # resolve here so a typo is one clear error before any worker spawns
+        from freetoken.gpu_select import resolve_gpu_uuids
+
+        try:
+            server_args = replace(server_args, gpu_assigned=resolve_gpu_uuids(server_args.gpu))
+        except ValueError as exc:
+            raise SystemExit(f"{prog or 'ft serve'}: error: {exc}") from exc
+        logger.info(
+            f"--gpu {','.join(server_args.gpu)} -> "
+            f"{', '.join(server_args.gpu_assigned) if server_args.gpu_assigned else 'resolved at CUDA init (no NVML)'}"
+        )
+
     def start_subprocess() -> "BackendHandle":
         import multiprocessing as mp
 

python/freetoken/server/stats.py

@@ -130,7 +130,9 @@ def _swa_page_size(config: Any) -> int:
 def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict:
     """Full /v1/stats doc. throughput is 0 when idle; kv/mamba/swa are null
     when their total is 0 (owned-KV / non-hybrid / non-SWA). kv and swa share one shape:
-    pages + the pool's own page_size (tokens = pages x page_size)."""
+    pages + the pool's own page_size (tokens = pages x page_size). gpus: the engine's GPU as
+    [{index, name, uuid, total_bytes}] (the primary rank's; a list so TP can extend it), []
+    until the readiness meta arrives."""
     tr: StatsTracker = state.stats
     config = state.config
     ready_at = getattr(state, "ready_at", None)
@@ -158,6 +160,7 @@ def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict:
         "mamba": mamba,
         "swa": swa,
         "vram_bytes": tr.vram_bytes,
+        "gpus": list(getattr(state, "gpus", None) or []),
         "throughput": {
             "decode_tps": round(tr.decode_tps(), 1),
             "prefill_tps": round(tr.prefill_tps(), 1),

tests/moe/test_hybrid_fetch.py

@@ -6,11 +6,12 @@ per-step integer split (GPU kernel vs CPU reference mirror, and the balance rule
 """
 
 import json
+import os
 
 import pytest
 import torch
 
-from freetoken.moe.bench_profile import load_hybrid_fetch_fraction
+from freetoken.moe.bench_profile import default_profile_path, load_backend_recommendation, load_hybrid_fetch_fraction
 from freetoken.moe.offload_cache import OffloadMoeCache
 
 Q = 1 << 16
@@ -65,6 +66,25 @@ def test_load_hybrid_fetch_fraction(tmp_path):
     assert load_hybrid_fetch_fraction("bf16", gpu_name="OTHER", path=str(path)) is None
 
 
+def test_profile_lookup_prefers_the_gpu_uuid_file(tmp_path, monkeypatch):
+    monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path))
+    monkeypatch.delenv("FREETOKEN_BENCHBW_PATH", raising=False)
+    uuid = "GPU-2f3a9b1c-0000-1111-2222-333344445555"
+
+    def write(path, name, verdict):
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        with open(path, "w") as f:
+            json.dump({"gpu": {"name": name}, "dtypes": {"bf16": verdict}}, f)
+
+    # legacy single file only: used when the name matches, ignored otherwise
+    write(default_profile_path(), "FAKE GPU", "hybrid")
+    assert load_backend_recommendation("bf16", gpu_name="FAKE GPU", gpu_uuid=uuid) == "hybrid"
+    assert load_backend_recommendation("bf16", gpu_name="OTHER", gpu_uuid=uuid) is None
+    # this card's own file wins over the legacy one
+    write(default_profile_path(uuid), "FAKE GPU", "offload")
+    assert load_backend_recommendation("bf16", gpu_name="FAKE GPU", gpu_uuid=uuid) == "offload"
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
 def test_hybrid_fraction_gpu_matches_cpu_reference():
     torch.manual_seed(0)