fix(moe): fit Qwen3.8 Flash Next FP8 banks across host and GPU memory

python/freetoken/engine/engine.py

@@ -681,6 +681,9 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
                 )
         except PinFailed as exc:
             raise RuntimeError(f"{exc}; {_pin_hint(self._host_tables_bytes)}") from exc
+        resident_bytes = sum(t.nbytes for per_layer in banks.sources.values() for t in per_layer if t.is_cuda)
+        self._weights_bytes += resident_bytes
+        self._post_weights_free -= resident_bytes
         if config.moe_cache_auto:
             size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method)
             object.__setattr__(config, "moe_cache_size", size)

python/freetoken/moe/expert_banks.py

@@ -93,6 +93,19 @@ def build_expert_banks(
     specs = {role: ((E, *spec.shape), spec.dtype) for role, spec in layout.items() if not spec.resident}
     hb = alloc_layer_banks(specs, num_layers)
     banks = {role: [b.tensor for b in hb[role]] for role in specs}
+    gpu_layers = set()
+    fit_host_ram = layer_sink is None and not dummy and device.type == "cuda" and method.kind is QuantKind.FP8_BLOCK
+    layer_bytes = sum(per_layer[0].nbytes for per_layer in hb.values())
+    if fit_host_ram:
+        available = _host_available_bytes()
+        if available is not None and layer_bytes:
+            # Keep one layer of host staging space while packing; CMA cannot back registered banks.
+            gpu_layers = set(range(min(num_layers, max(0, (layer_bytes * (num_layers + 1) - available + layer_bytes - 1) // layer_bytes))))
+        if gpu_layers:
+            for role, per_layer in banks.items():
+                for layer_id in gpu_layers:
+                    per_layer[layer_id] = torch.zeros_like(per_layer[layer_id], device=device)
+            logger.info(f"expert banks: keeping {len(gpu_layers)} FP8 layers ({len(gpu_layers) * layer_bytes / 2**30:.2f} GiB) on GPU to fit host RAM")
     alphas = {
         role: torch.empty(num_layers * E, dtype=spec.dtype, device=device)
         for role, spec in layout.items() if spec.resident
@@ -122,6 +135,14 @@ def _fill(sink) -> None:
             # refuse before writing: a duplicate row would also complete the layer early and hand the sink a half-filled bank
             if written[layer_id, e0:e1].any():
                 raise ValueError(f"expert rows written more than once: layer {layer_id}, experts {e0}:{e1}")
+            if fit_host_ram and layer_id not in gpu_layers and not written[layer_id].any():
+                # Reclaim and host registration change usable RAM during a long load.
+                available = _host_available_bytes()
+                if available is not None and available < 2 * layer_bytes:
+                    for per_layer in banks.values():
+                        per_layer[layer_id] = torch.zeros_like(per_layer[layer_id], device=device)
+                    gpu_layers.add(layer_id)
+                    logger.info(f"expert banks: keeping FP8 layer {layer_id} on GPU after host RAM changed")
             written[layer_id, e0:e1] = 1
             out = {role: banks[role][layer_id][e0:e1] for role in specs}
             got = method.pack(piece, out)
@@ -138,7 +159,7 @@ def _fill(sink) -> None:
         _fill(layer_sink)
     elif torch.cuda.is_available():
         with PinPipeline() as pins:
-            _fill(pins)
+            _fill(lambda layer_id, layer_banks: pins(layer_id, layer_banks) if layer_id not in gpu_layers else None)
     else:
         _fill(None)
 
@@ -152,6 +173,15 @@ def _fill(sink) -> None:
 _PARALLEL_CHUNK = 8 << 20  # default O_DIRECT chunk for the parallel reader
 
 
+def _host_available_bytes() -> int | None:
+    try:
+        with open("/proc/meminfo") as f:
+            mem = {key: int(value.split()[0]) * 1024 for key, value in (line.split(":", 1) for line in f)}
+        return max(0, mem["MemAvailable"] - mem.get("CmaFree", 0))
+    except (OSError, KeyError, ValueError):
+        return None
+
+
 def _q4_0_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None) -> ExpertBanks:
     if parallel:
         raise NotImplementedError(
@@ -209,15 +239,7 @@ def _host_ram_fits_parallel(model_path: str) -> bool:
     extra (non-reclaimable) whole-shard buffer? Unknown (non-local path / no /proc) -> True,
     i.e. keep the fast path. Banks ~= checkpoint size (experts dominate); transient ~= the
     largest shard. Uses MemAvailable (counts reclaimable cache) -- the OOM-relevant figure."""
-    avail = None
-    try:
-        with open("/proc/meminfo") as f:
-            for line in f:
-                if line.startswith("MemAvailable:"):
-                    avail = int(line.split()[1]) * 1024
-                    break
-    except OSError:
-        pass
+    avail = _host_available_bytes()
     if avail is None:
         return True
     try:  # resolve a hub id to its local cache dir (no-op for a local path) so glob sees the shards

tests/moe/test_offload.py

@@ -76,6 +76,63 @@ def _bound(quant):
     assert torch.all(banks.sources["gate_up_global"][0].float() > 0)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+@pytest.mark.parametrize("available_layers, gpu_layers", [(4, ()), (3, (0,)), (1, (0, 1, 2)), ((4, 4, 1, 1), (1, 2))])
+def test_fp8_banks_fit_host_ram_without_changing_weight_bytes(monkeypatch, available_layers, gpu_layers):
+    from freetoken.layers.quantization.moe.base import MoEConfig
+    from freetoken.layers.quantization.moe.fp8_block import Fp8BlockMoEMethod
+    from freetoken.moe import expert_banks
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    method = Fp8BlockMoEMethod(MoEConfig(4, 256, 128, 2, strategy="offload"))
+    specs = method.layout()
+    layer_bytes = sum(4 * torch.empty(s.shape, dtype=s.dtype).nbytes for s in specs.values())
+    if isinstance(available_layers, tuple):
+        availability = iter(available_layers)
+        monkeypatch.setattr(expert_banks, "_host_available_bytes", lambda: next(availability) * layer_bytes)
+    else:
+        monkeypatch.setattr(expert_banks, "_host_available_bytes", lambda: available_layers * layer_bytes)
+    pieces = []
+    for layer in range(3):
+        piece = {
+            "gate_up": torch.full((4, 256, 256), layer + 1, dtype=torch.float32).to(torch.float8_e4m3fn),
+            "down": torch.full((4, 256, 128), layer + 2, dtype=torch.float32).to(torch.float8_e4m3fn),
+            "gate_up_scale": torch.full((4, 2, 2), layer + 1, dtype=torch.bfloat16),
+            "down_scale": torch.full((4, 2, 1), layer + 2, dtype=torch.bfloat16),
+        }
+        pieces.append((layer, 0, 4, piece))
+    banks = expert_banks.build_expert_banks(method, 3, iter(pieces), device=torch.device("cuda"))
+    cache = OffloadMoeCache(3, 4, 8, device=torch.device("cuda"), quant_format="fp8_block", prefill_overlap=True)
+    cache.set_bank_sources(banks.sources)
+    cache.begin_prefill()
+    for layer, _, _, piece in pieces:
+        got = cache.wait_prefill_layer(layer)
+        for (name, per_layer), copied in zip(cache.bank_sources.items(), got):
+            source = per_layer[layer]
+            assert source.is_cuda == (layer in gpu_layers)
+            reference = torch.zeros(source.shape, dtype=source.dtype)
+            value = piece[name]
+            reference[..., :value.shape[-1]].copy_(value)
+            assert torch.equal(copied.cpu().view(torch.uint8), reference.view(torch.uint8))
+        cache.release_prefill_layer(layer)
+    cache.reset()
+    for layer, _, _, _ in pieces:
+        ids = torch.tensor([3, 1], dtype=torch.int32, device="cuda")
+        cache.ensure_experts(layer, ids)
+        cache.copy_missing()
+        for name, per_layer in cache.bank_sources.items():
+            got = cache.bank_caches[name][ids.long()].cpu().view(torch.uint8)
+            assert torch.equal(got, per_layer[layer].cpu()[[3, 1]].view(torch.uint8))
+
+
+def test_host_available_bytes_excludes_cma(monkeypatch):
+    from io import StringIO
+    from freetoken.moe.expert_banks import _host_available_bytes
+
+    monkeypatch.setattr("builtins.open", lambda *a, **kw: StringIO("MemAvailable: 2048 kB\nCmaFree: 1792 kB\n"))
+    assert _host_available_bytes() == 256 * 1024
+
+
 def test_offload_moe_layer_prefill_forward_uses_single_layer_cache_view(monkeypatch):
     layer, cache = _make_layer_and_cache()
     topk_weights = torch.tensor([[0.7, 0.3]], dtype=torch.float32)