Windows/ROCm: make GGUF MoE offload work end-to-end on gfx1201 (24.4 -> 37.9 tok/s)

dist/install.ps1

@@ -43,7 +43,6 @@ $tag = ($Py -replace '\s', '') + "-" + $Arch.Replace(',', '+')
 if (-not $WheelDir) { $WheelDir = Join-Path $REPO "rocm-wheels\$tag" }
 $VENV = Join-Path $REPO ".venv"
 $PYEXE = "$VENV\Scripts\python.exe"
-$PIP   = "$VENV\Scripts\python.exe -m pip"
 
 Write-Host ""
 Write-Host "  FreeToken installer for Windows + AMD" -ForegroundColor Cyan
@@ -52,6 +51,10 @@ Write-Host "  --------------------------------------"
 # ---- Step 1: private python environment -------------------------------
 Write-Host "`n[1/5] Creating a private Python environment (.venv) ..." -ForegroundColor Yellow
 Invoke-Expression "$Py -m venv `"$VENV`""
+# A fresh venv ships pip but no setuptools on 3.12+, and step 2 installs the `rocm`
+# metapackage sdist with --no-build-isolation -- without setuptools that dies with
+# "Cannot import 'setuptools.build_meta'", leaving torch unable to import rocm_sdk.
+& $PYEXE -m pip install --upgrade pip setuptools wheel
 
 # ---- Step 2: AMD GPU wheels -------------------------------------------
 Write-Host "[2/5] AMD GPU wheels (torch / ROCm) ..." -ForegroundColor Yellow
@@ -83,6 +86,11 @@ if ($rocmSdist) { & $PYEXE -m pip install $rocmSdist.FullName --no-deps --no-bui
 Write-Host "[3/5] Installing FreeToken + helpers ..." -ForegroundColor Yellow
 & $PYEXE -m pip install "triton-windows>=3.7.1" apache-tvm-ffi==0.1.13.post3 msgpack pyzmq psutil requests aiohttp partial_json_parser gguf `
     einops fastapi uvicorn pydantic openai prompt_toolkit "transformers>=5.5,<6" huggingface_hub safetensors `
+    "numpy>=2.0,<2.5" tqdm modelscope tornado ninja setuptools wheel numba
+# flashlib (MoE expert-cache slot_cache kernel) is a real runtime dep, but it declares
+# torch>=2.0 -- resolving that would pull the CUDA torch from PyPI over the ROCm wheel
+# installed in step 2. Its other deps (triton-windows, numpy, numba, tqdm) are above.
+& $PYEXE -m pip install flashlib==0.3.0 --no-deps
     "numpy>=2.0,<2.5" tqdm modelscope tornado ninja numba setuptools wheel
 $env:FREETOKEN_SKIP_CUDA_EXT = "1"
 & $PYEXE -m pip install -e "$REPO" --no-deps --no-build-isolation
@@ -95,6 +103,9 @@ Write-Host "[4/5] Applying 3 small compatibility patches ..." -ForegroundColor Y
 # ---- Step 5: verify -----------------------------------------------------
 Write-Host "[5/5] Checking your GPU ..." -ForegroundColor Yellow
 & $PYEXE -c "import torch; print('      torch', torch.__version__, '| HIP', torch.version.hip); print('      GPU:', torch.cuda.get_device_name(0)); arch=torch.cuda.get_device_properties(0).gcnArchName.split(':')[0]; print('      arch:', arch); assert arch=='$Arch', f'GPU arch {arch} != installed device wheels ($Arch) - rerun with -Arch {arch}'"
+# $ErrorActionPreference does not apply to native exit codes, so check it explicitly --
+# otherwise a failed GPU check still prints "All done!" over a broken install.
+if ($LASTEXITCODE -ne 0) { throw "GPU check failed - see the traceback above; the install is not usable." }
 
 Write-Host ""
 Write-Host "  All done! To chat with a model:" -ForegroundColor Green

dist/run-server.ps1

@@ -10,7 +10,15 @@
 #                           (not needed if HIP_PATH is already set)
 #    -Port 1919             API port
 #    -KVPages 4096          cap KV cache (needed for big dense models)
-#    -ExtraArgs "--flag"    anything else for `ft serve`
+#    -ExtraArgs "..."       anything else for `ft serve`, as ONE space-separated
+#                           string:  -ExtraArgs "--moe-backend offload --num-pages 4096"
+#                           (`powershell -File` does no PowerShell parsing, so the
+#                           array form "--a","b" arrives as the literal string
+#                           "--a,b" -- both are split apart below)
+#
+#  Pass --cuda-graph-max-bs 0 in -ExtraArgs. On gfx1201 CUDA graphs are a large
+#  LOSS, not a win: measured on Gemma-4-26B-A4B QAT q4_0 with identical settings,
+#  eager decodes at 37.9 tok/s and graph replay at 10.4.
 #
 #  When it says READY, open http://localhost:1420 and chat.
 # ============================================================
@@ -26,6 +34,12 @@ $ErrorActionPreference = "Stop"
 $REPO = Split-Path -Parent $PSScriptRoot
 $LogDir = "$env:TEMP\freetoken-logs"
 New-Item -ItemType Directory -Force -Path $LogDir | Out-Null
+# Clear last run's logs up front: the redirect below only truncates them once cmd
+# reaches the `ft serve` line (after vcvarsall), and until then the readiness poll
+# is reading the PREVIOUS run - a stale traceback aborts this launch instantly.
+Remove-Item "$LogDir\serve.log", "$LogDir\serve_err.log" -ErrorAction SilentlyContinue
+if (Test-Path "$LogDir\serve.log") {
+    throw "$LogDir\serve.log is still locked by a previous run, so this launch would read the old log and report its error as this one's. A scheduler worker outlived a failed server and still holds the inherited handle: it is a python.exe running ``multiprocessing.spawn``, NOT anything matching ``freetoken``. Run dist\stop-server.ps1 - it sweeps those - then retry."
 "$LogDir\serve.log", "$LogDir\serve_err.log" | ForEach-Object {
     if (Test-Path $_) {
         try {
@@ -39,14 +53,62 @@ New-Item -ItemType Directory -Force -Path $LogDir | Out-Null
 
 if (-not $RocmPath) {
     if ($env:HIP_PATH) { $RocmPath = $env:HIP_PATH }
-    else { throw "Where is the AMD ROCm runtime? Pass -RocmPath or set HIP_PATH once:`n       [Environment]::SetEnvironmentVariable('HIP_PATH','C:\\ROCm\\...','User')" }
+    else {
+        # The ROCm runtime ships INSIDE the venv install.ps1 built -- the TheRock wheels
+        # unpack it to site-packages\_rocm_sdk_core (lib\llvm\..., the layout the env vars
+        # below assume). Ask Python where that is before telling the user to go find a
+        # system ROCm install they never had to make; HIP_PATH is set by nothing here.
+        $pyExe = Join-Path $REPO ".venv\Scripts\python.exe"
+        if (-not (Test-Path $pyExe)) { $pyExe = "python" }
+        $found = & $pyExe -c "import _rocm_sdk_core, os; print(os.path.dirname(_rocm_sdk_core.__file__))" 2>$null
+        if ($found) { $found = ([string]$found).Trim() }
+        if ($found -and (Test-Path (Join-Path $found "lib\llvm\bin\clang.exe"))) {
+            $RocmPath = $found
+            Write-Host "ROCm runtime: $RocmPath (from the venv)" -ForegroundColor DarkGray
+        }
+    }
+    if (-not $RocmPath) { throw "Where is the AMD ROCm runtime? It is not in this repo's .venv and HIP_PATH is unset. Pass -RocmPath, or re-run dist\install.ps1 to build the venv." }
+}
+# `powershell -File` hands every argument through as a literal string -- there is no
+# PowerShell parser on that path -- so -ExtraArgs "--moe-backend","offload" arrives as
+# the single element "--moe-backend,offload", which `ft serve` rejects as one
+# unrecognized argument. Split every element on commas and whitespace so the array form
+# (dot-sourced, or -Command) and the string form (-File) both reach ft the same way.
+$ExtraArgs = @($ExtraArgs | ForEach-Object { $_ -split '[,\s]+' } | Where-Object { $_ })
+if ($KVPages -gt 0) { $ExtraArgs += "--num-pages", "$KVPages" }
+
+# The generated runner does `cd /d %TEMP%` before calling ft, so a relative -Model
+# (the form the README and every note use: modelsoo.gguf) resolves against %TEMP%,
+# misses, and transformers then treats it as a HUGGING FACE REPO ID -- the error names
+# repo-id character rules and never mentions the path. Resolve it here, against the
+# caller's cwd and then the repo, before it can turn into a download attempt.
+if (-not [System.IO.Path]::IsPathRooted($Model)) {
+    $candidate = Join-Path (Get-Location) $Model
+    if (-not (Test-Path $candidate)) { $candidate = Join-Path $REPO $Model }
+    if (-not (Test-Path $candidate)) {
+        throw "Model not found: '$Model' (looked in $(Get-Location) and $REPO). Pass a full path."
+    }
+    $Model = (Resolve-Path $candidate).Path
 }
-if ($KVPages -gt 0) { $ExtraArgs += "--num-pages $KVPages" }
 
 # engine binary: prefer the repo venv this installer created, fall back to PATH
 $ft = Join-Path $REPO ".venv\Scripts\ft.exe"
 if (-not (Test-Path $ft)) { $ft = "ft" }
 
+# vcvarsall lives under either Program Files root (Build Tools installs land in the
+# x86 one), so ask vswhere first and only then fall back to scanning both roots.
+$vswhere = "${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer\vswhere.exe"
+$vcvars = $null
+if (Test-Path $vswhere) {
+    $vsRoot = & $vswhere -latest -products * -find "VC\Auxiliary\Build\vcvarsall.bat" 2>$null | Select-Object -First 1
+    if ($vsRoot) { $vcvars = $vsRoot }
+}
+if (-not $vcvars) {
+    $vcvars = Get-ChildItem "$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio" `
+                  -Recurse -Filter vcvarsall.bat -ErrorAction SilentlyContinue |
+              Select-Object -First 1 -ExpandProperty FullName
+}
+if (-not $vcvars) { Write-Warning "vcvarsall.bat not found - JIT DLL links may fail to find the MSVC CRT." }
 $vcvars = @("${env:ProgramFiles(x86)}\Microsoft Visual Studio", "$env:ProgramFiles\Microsoft Visual Studio") |
           Where-Object { Test-Path $_ } |
           ForEach-Object { Get-ChildItem $_ -Recurse -Filter vcvarsall.bat -ErrorAction SilentlyContinue } |
@@ -58,7 +120,17 @@ if (-not $vcvars) {
 $cmd = @"
 call `"$vcvars`" x64 >nul
 set HIP_PATH=$RocmPath
+rem tvm_ffi's JIT resolves the toolchain from ROCM_HOME (not HIP_PATH) and dies at
+rem the first kernel build without it -- long after the model has loaded.
+rem ROCM_HOME only: setting ROCM_PATH too makes clang look for the device bitcode
+rem at %ROCM_PATH%\amdgcn\bitcode, which is not the TheRock wheel layout (it lives
+rem under lib\llvm\), and every kernel build then fails to find it.
+set ROCM_HOME=$RocmPath
 set TVM_FFI_ROCM_ARCH_LIST=$Arch
+rem Without this torch's JIT compiles every extension for EVERY visible card --
+rem including the CPU's iGPU (gfx1036 on a 9800X3D), which is not the serving
+rem device and whose build failure kills the backend worker after a full load.
+set PYTORCH_ROCM_ARCH=$Arch
 set TRITON_OVERRIDE_ARCH=$Arch
 set ROCM_SDK_TARGET_FAMILY=$Arch
 set PYTORCH_ROCM_ARCH=$Arch
@@ -69,6 +141,7 @@ set TVM_FFI_CACHE_DIR=$REPO\.tvm-ffi-cache
 set PATH=$RocmPath\bin;%PATH%
 set "CC=$RocmPath\lib\llvm\bin\clang-cl.exe"
 cd /d %TEMP%
+"$ft" serve --model "$Model" --port $Port $($ExtraArgs -join ' ') > "$LogDir\serve.log" 2> "$LogDir\serve_err.log"
 "$ft" serve --model-path "$Model" --port $Port $($ExtraArgs -join ' ') > "$LogDir\serve.log" 2> "$LogDir\serve_err.log"
 "@
 $runner = Join-Path $env:TEMP "freetoken_serve.cmd"
@@ -87,7 +160,14 @@ for ($i = 1; $i -le 120; $i++) {
         Write-Host "  Logs: $LogDir\serve.log / serve_err.log"
         exit 0
     }
-    if (Select-String -Path "$LogDir\serve.log","$LogDir\serve_err.log" -Pattern "AssertionError|Traceback|exited during load" -ErrorAction SilentlyContinue) {
+    # torch LOGS tracebacks as warnings and keeps going -- cpp_extension's
+    # "Error checking compiler version" probe prints a full Traceback on every ROCm
+    # start, and matching it aborted a load that went on to serve fine. Warning lines
+    # carry torch's rank/severity stamp ("[rank0]:W0830 ..."); a real crash does not.
+    $fatal = Select-String -Path "$LogDir\serve.log","$LogDir\serve_err.log" `
+                 -Pattern "AssertionError|Traceback|exited during load" -ErrorAction SilentlyContinue |
+             Where-Object { $_.Line -notmatch '\]:[WI]\d{4} ' }
+    if ($fatal) {
         Write-Host "`n  The server hit an error while loading. Last lines:" -ForegroundColor Red
         Get-Content "$LogDir\serve_err.log","$LogDir\serve.log" -Tail 6 -ErrorAction SilentlyContinue
         exit 1

dist/stop-server.ps1

@@ -1,3 +1,75 @@
+# ============================================================
+#  Stop any running FreeToken engine / web UI -- and their children.
+#
+#  Matching command lines alone is not enough. The engine's scheduler runs in
+#  multiprocessing SPAWN workers, whose command line is
+#      python.exe -c "from multiprocessing.spawn import spawn_main; ..."
+#  which names neither `serve` nor `freetoken` nor the model. Killing only the
+#  roots left those workers alive holding ~20 GiB of pinned expert banks AND the
+#  inherited serve.log / serve_err.log handles, which then trips run-server.ps1's
+#  "log is still locked" guard on the next launch.
+#
+#  So: kill the roots' whole process tree, and separately sweep up spawn workers
+#  already orphaned by an earlier partial stop. Those are identified by their
+#  EXECUTABLE, not their command line -- only this repo's .venv python is ours, so
+#  an unrelated multiprocessing app on the box is never touched.
+# ============================================================
+$REPO = Split-Path -Parent $PSScriptRoot
+$snapshot = Get-CimInstance Win32_Process |
+            Select-Object ProcessId, ParentProcessId, Name, CommandLine, CreationDate, ExecutablePath
+
+$isOurs = { $_.ExecutablePath -and $_.ExecutablePath.StartsWith($REPO, 'OrdinalIgnoreCase') }
+
+$roots = $snapshot | Where-Object {
+    ($_.Name -eq 'ft.exe' -or $_.Name -like '*python*') -and
+    $_.CommandLine -match 'serve|http\.server 1420|freetoken'
+}
+# Workers orphaned by a previous run: our venv's python, sitting in a spawn loop.
+$orphans = $snapshot | Where-Object {
+    $_.CommandLine -match 'multiprocessing\.spawn' -and (& $isOurs)
+}
+
+# Children of a pid, transitively. A dead PID gets reused, so a stale ParentProcessId
+# can point at an unrelated live process -- only follow a child that started after its
+# claimed parent did.
+function Get-Descendants($processId, $startedAt) {
+    $kids = $snapshot | Where-Object {
+        $_.ParentProcessId -eq $processId -and $_.ProcessId -ne $processId -and
+        (-not $startedAt -or -not $_.CreationDate -or $_.CreationDate -ge $startedAt)
+    }
+    foreach ($k in $kids) { $k; Get-Descendants $k.ProcessId $k.CreationDate }
+}
+
+$targets = @()
+foreach ($r in @($roots) + @($orphans)) {
+    $targets += $r
+    $targets += Get-Descendants $r.ProcessId $r.CreationDate
+}
+$targets = @($targets | Sort-Object ProcessId -Unique)
+if (-not $targets) { "no FreeToken processes running"; return }
+
+# Depth in the kill set, so we can go deepest-first: a parent killed early can spawn a
+# replacement worker on its way out.
+$depth = @{}
+foreach ($t in $targets) {
+    $d = 0; $cur = $t
+    while ($cur -and $d -lt 32) {
+        $cur = @($targets | Where-Object { $_.ProcessId -eq $cur.ParentProcessId })[0]
+        if ($cur) { $d++ }
+    }
+    $depth[[int]$t.ProcessId] = $d
+}
+foreach ($t in ($targets | Sort-Object { -$depth[[int]$_.ProcessId] })) {
+    $mb = 0
+    try { $mb = [int]((Get-Process -Id $t.ProcessId -ErrorAction Stop).WorkingSet64 / 1MB) } catch {}
+    try {
+        Stop-Process -Id $t.ProcessId -Force -ErrorAction Stop
+        "stopped $($t.ProcessId) $($t.Name) ($mb MB)"
+    } catch {
+        # Usually already gone because we killed its parent -- that is the point.
+        if (Get-Process -Id $t.ProcessId -ErrorAction SilentlyContinue) {
+            Write-Warning "could not stop $($t.ProcessId) $($t.Name): $($_.Exception.Message)"
+        }
 # Stop this repo's FreeToken engine, its worker processes, and the web UI.
 $REPO = Split-Path -Parent $PSScriptRoot
 $repoPattern = [regex]::Escape($REPO)

python/freetoken/engine/engine.py

@@ -3,6 +3,7 @@
 import gc
 import math
 import os
+import sys
 from datetime import timedelta
 from typing import Any, Dict, Iterable, NamedTuple, Tuple
 
@@ -320,8 +321,12 @@ def __init__(self, config: EngineConfig):
         set_rope_device(self.device)
         with torch.device("meta"), torch_dtype(config.dtype):
             self.model = create_model(config.model_config)
+        from freetoken.moe.host_banks import host_mem_summary
+
+        logger.info_rank0(f"Host memory before loading weights: {host_mem_summary()}")
         self.model.load_state_dict(self._load_weight_state_dict(config))
         post_weights_free = self._sync_get_memory()[0]
+        logger.info_rank0(f"Host memory after loading weights:  {host_mem_summary()}")
         self._weights_bytes = self._baseline_free - post_weights_free
         # Pool-budget baseline for the desktop cache sliders: free VRAM after the weights are
         # resident but before ANY runtime cache pool (MoE expert cache below, KV pages, GDN
@@ -977,6 +982,19 @@ def _ensure_expandable_segments() -> None:
     """
     if os.environ.get("PYTORCH_ALLOC_CONF") or os.environ.get("PYTORCH_CUDA_ALLOC_CONF"):
         return
+    # ROCm on Windows backs an expandable segment with HOST memory: creating the segment,
+    # and every later growth of it, mirrors the VRAM heap into the process working set
+    # (measured ~18 GiB per event on a 16 GiB card). On an offload MoE run host RAM is
+    # exactly the resource the design is spending, so this silently eats tens of GiB
+    # before a single weight is read and the expert-bank pin then fails with availPhys
+    # at zero -- a fragmentation guard that costs more than the fragmentation. Opt back
+    # in with PYTORCH_ALLOC_CONF=expandable_segments:True.
+    if sys.platform == "win32" and getattr(torch.version, "hip", None) is not None:
+        logger.info_rank0(
+            "expandable_segments left OFF: on ROCm/Windows the segment is host-backed "
+            "(set PYTORCH_ALLOC_CONF=expandable_segments:True to override)"
+        )
+        return
     try:
         torch.cuda.memory._set_allocator_settings("expandable_segments:True")
     except Exception as exc:  # pragma: no cover - depends on torch build