fix(engine): resolve moe-strategy auto to fused on unified-memory GPUs (GB10)

python/freetoken/engine/engine.py

@@ -999,6 +999,41 @@ def _profile_gpu(index: "int | None" = None) -> Tuple[str | None, str | None]:
     return ident["name"], ident["uuid"]
 
 
+def _is_unified_memory_gpu(index: "int | None" = None) -> bool:
+    """True when the GPU has no separate device memory (cudaDevAttrIntegrated): host
+    banks and the GPU slot cache are the same DRAM, so the offload family's pinned
+    staging + slot gather are DRAM-to-DRAM copies with no PCIe link to hide behind.
+
+    The attribute is reliable on true-UMA parts (Jetson, GB10/DGX Spark) but not on
+    C2C-linked discrete-HBM parts (GH200 reports integrated=0 despite coherent CPU
+    memory), and it has only been verified on GB10 so far --
+    FREETOKEN_UNIFIED_MEMORY=0/1 overrides the probe where the attribute lies."""
+    env = os.environ.get("FREETOKEN_UNIFIED_MEMORY")
+    if env is not None:
+        return env.strip().lower() not in {"0", "false", "no", "off"}
+    if not torch.cuda.is_available():
+        return False
+    try:
+        dev = torch.cuda.current_device() if index is None else index
+        return bool(torch.cuda.get_device_properties(dev).is_integrated)
+    except Exception:
+        return False
+
+
+def _fused_resident_ok(model_config) -> bool:
+    """Whether the resident ('fused') MoE path can hold this model's experts.
+
+    Mirrors the validation further down: fused requires expert_quant in
+    {none, fp8_block}. Legacy weight formats (mxfp4/q4_0 via moe_weight_format)
+    still dispatch on the offload cache's format tag, so they stay offload-only.
+    NVFP4 checkpoints have a resident_view since the quant-method refactor, but
+    the nvfp4-fused combination is not yet validated; keep it offload until it is."""
+    expert_quant = getattr(model_config, "expert_quant", "none")
+    if expert_quant not in ("none", "fp8_block"):
+        return False
+    return getattr(model_config, "moe_weight_format", None) in (None, "bf16")
+
+
 def _ensure_expandable_segments() -> None:
     """Default the CUDA allocator to expandable segments.
 
@@ -1451,6 +1486,21 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
         # -- auto never picks it, because nothing here knows whether the experts would fit in
         # HBM and a wrong guess is a weight-load OOM rather than a slower-but-working run.
         default_backend = "offload"
+        # Unified memory (GB10/DGX Spark, Jetson): there is no host/device boundary, so
+        # the offload family stages and gathers between two names for the same DRAM (on
+        # GB10 this added a measured ~130 s stall to every request, #369). Resident
+        # experts are the safe default here, not the risky one: the model and its banks
+        # page from the same pool, so the "wrong guess = weight-load OOM" rationale above
+        # does not apply. The benchbw hybrid upgrade is skipped too: CPU execution adds
+        # no bandwidth when both sides share one memory. Only formats the resident path
+        # can actually hold take this branch; the rest stay on offload as before.
+        unified_memory = _is_unified_memory_gpu()
+        if unified_memory and _fused_resident_ok(model_config):
+            default_backend = "fused"
+            logger.info_rank0(
+                "Unified-memory GPU detected; auto-selecting 'fused' MoE strategy "
+                "(resident experts) instead of offload"
+            )
         # Hardware-adaptive config: a cached `ft bench bw` profile can upgrade
         # the offload default to hybrid when this machine's CPU MoE bandwidth clears its PCIe
         # gather bandwidth by the bench threshold (default 2x). hybrid is VRAM-equivalent to
@@ -1464,7 +1514,14 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
         from freetoken.moe.bench_profile import load_backend_recommendation
 
         gpu_name, gpu_uuid = _profile_gpu()
-        if load_backend_recommendation(bench_fmt, gpu_name=gpu_name, gpu_uuid=gpu_uuid) == "hybrid":
+        if (
+            default_backend == "offload"
+            and not unified_memory
+            and 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")
@@ -1510,6 +1567,20 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
                 f"auto-selected strategy {config.moe_strategy!r}"
             )
 
+    if (
+        is_moe
+        and config.moe_strategy == "offload"
+        and _is_unified_memory_gpu()
+        and _fused_resident_ok(model_config)
+    ):
+        # An explicit offload pick is honored, but on unified memory the user is paying
+        # for copies between two names for the same DRAM; say so once at config time.
+        logger.warning_rank0(
+            "--moe-strategy offload on a unified-memory GPU: expert 'streaming' copies "
+            "DRAM to DRAM (there is no PCIe link to overlap it with). If the model fits, "
+            "--moe-strategy fused avoids the slot-cache machinery entirely."
+        )
+
     if is_moe and config.moe_strategy == "fused":
         # An explicit 'fused' keeps the experts resident, so there is no slot cache to size. The
         # sizing flags no longer redirect the backend, so ignore them here and say so -- the

tests/engine/test_moe_strategy_uma.py

@@ -0,0 +1,104 @@
+"""Unified-memory (UMA) MoE strategy resolution (DGX Spark / GB10, Jetson).
+
+CPU-only: exercises _is_unified_memory_gpu (env override) and
+_fused_resident_ok (format gating) without a GPU.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from freetoken.engine.engine import _fused_resident_ok, _is_unified_memory_gpu
+
+
+@pytest.fixture
+def uma_env(monkeypatch):
+    def set_env(value: str | None):
+        if value is None:
+            monkeypatch.delenv("FREETOKEN_UNIFIED_MEMORY", raising=False)
+        else:
+            monkeypatch.setenv("FREETOKEN_UNIFIED_MEMORY", value)
+
+    return set_env
+
+
+def test_env_override_forces_unified(uma_env):
+    uma_env("1")
+    assert _is_unified_memory_gpu() is True
+    uma_env("true")
+    assert _is_unified_memory_gpu() is True
+
+
+def test_env_override_forces_discrete(uma_env):
+    # must win even on real UMA hardware: this is the escape hatch when
+    # cudaDevAttrIntegrated lies
+    uma_env("0")
+    assert _is_unified_memory_gpu() is False
+    uma_env("off")
+    assert _is_unified_memory_gpu() is False
+
+
+def test_probe_falls_back_cleanly(uma_env, monkeypatch):
+    # no CUDA / probe failure -> discrete (offload stays the safe default)
+    uma_env(None)
+    import torch
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    assert _is_unified_memory_gpu() is False
+
+
+def test_fused_resident_ok_plain_formats():
+    assert _fused_resident_ok(SimpleNamespace(expert_quant="none")) is True
+    assert _fused_resident_ok(SimpleNamespace(expert_quant="none", moe_weight_format="bf16")) is True
+    assert _fused_resident_ok(SimpleNamespace(expert_quant="fp8_block")) is True
+
+
+def test_adjust_config_selects_fused_on_uma(monkeypatch):
+    from freetoken.distributed import DistributedInfo
+    from freetoken.engine.config import EngineConfig
+    import freetoken.engine.engine as engine_module
+
+    monkeypatch.setattr(engine_module, "_is_unified_memory_gpu", lambda index=None: True)
+    config = EngineConfig(
+        model_path="/tmp/freetoken-test-model",
+        tp_info=DistributedInfo(rank=0, size=1),
+        dtype=torch.float16,
+        attention_backend="fi",
+        moe_cache_rate=0.3,
+    )
+    object.__setattr__(
+        config,
+        "model_config",
+        SimpleNamespace(
+            has_swa_attention=False,
+            has_linear_attention=False,
+            is_moe=True,
+            num_layers=10,
+            num_moe_layers=10,
+            num_experts=8,
+            expert_quant="none",
+            moe_strategy="auto",
+        ),
+    )
+
+    engine_module._adjust_config(config)
+
+    assert config.moe_strategy == "fused"
+    assert config.moe_cache_size == 0
+    assert config.moe_cache_rate is None
+
+
+@pytest.mark.parametrize("fmt", ["nvfp4", "mxfp8"])
+def test_fused_resident_ok_rejects_quantized_experts(fmt):
+    # mirrors the engine's own fused guard; nvfp4 has resident_view but the
+    # combination is not yet validated
+    assert _fused_resident_ok(SimpleNamespace(expert_quant=fmt)) is False
+
+
+@pytest.mark.parametrize("fmt", ["mxfp4", "q4_0"])
+def test_fused_resident_ok_rejects_gguf_weight_formats(fmt):
+    # GGUF banks dispatch on the offload cache's format tag; no resident path
+    assert _fused_resident_ok(SimpleNamespace(expert_quant="none", moe_weight_format=fmt)) is False

tests/moe/test_offload.py

@@ -412,13 +412,18 @@ def test_lru_gpu_cache_assigns_unique_slots_for_large_miss_batch():
     assert cache.src_indices[:256].tolist() == list(range(256))
 
 
-def test_adjust_config_converts_moe_cache_rate_to_cache_size():
+def test_adjust_config_converts_moe_cache_rate_to_cache_size(monkeypatch):
     from types import SimpleNamespace
 
     from freetoken.distributed import DistributedInfo
     from freetoken.engine.config import EngineConfig
+    import freetoken.engine.engine as engine_module
     from freetoken.engine.engine import _adjust_config
 
+    # This test exercises the discrete-GPU offload path regardless of the host
+    # running the suite (GB10 reports cudaDevAttrIntegrated=1).
+    monkeypatch.setattr(engine_module, "_is_unified_memory_gpu", lambda index=None: False)
+
     config = EngineConfig(
         model_path="/tmp/freetoken-test-model",
         tp_info=DistributedInfo(rank=0, size=1),