From 948df9221091af0d25e4ceb8349709fc9194f6c2 Mon Sep 17 00:00:00 2001 From: iamanishx Date: Fri, 11 Sep 2026 13:34:50 +0530 Subject: [PATCH] fix(engine): resolve moe-strategy auto to fused on unified-memory GPUs (GB10) --- python/freetoken/engine/engine.py | 73 +++++++++++++++++- tests/engine/test_moe_strategy_uma.py | 104 ++++++++++++++++++++++++++ tests/moe/test_offload.py | 7 +- 3 files changed, 182 insertions(+), 2 deletions(-) create mode 100644 tests/engine/test_moe_strategy_uma.py diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 006c19089..38eba5c20 100644 --- a/python/freetoken/engine/engine.py +++ b/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 diff --git a/tests/engine/test_moe_strategy_uma.py b/tests/engine/test_moe_strategy_uma.py new file mode 100644 index 000000000..0f316c871 --- /dev/null +++ b/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 diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index b29ce531d..cd5f5dc77 100644 --- a/tests/moe/test_offload.py +++ b/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),