Windows/ROCm: make GGUF MoE offload work end-to-end on gfx1201 (24.4 -> 37.9 tok/s) - #5
Open
uploadyours wants to merge 6 commits into
Open
Conversation
Found while bringing the fork up on an RX 9070 XT (gfx1201, Windows 11) against
a real Qwen3.5-35B-A3B Q4_K_M GGUF.
dist/install.ps1 — the one-command install could not run at all:
- `$PIP` was a string ("python.exe -m pip"), so `$PIP install ...` is a
PowerShell parse error ("Unexpected token 'install'") on four lines and the
script died before doing anything. Use `& $PYEXE -m pip install`, the idiom
install-offline.ps1 already uses.
- flashlib==0.3.0 was missing from the helper list. It is a hard runtime dep
(the slot_cache LRU admission kernel behind the MoE expert cache), and
install-offline.ps1 does install it. Installed here with --no-deps, because
it declares torch>=2.0 and resolving that pulls the CUDA torch from PyPI
straight over the ROCm wheel from step 2. Its other deps are in the main
list; numba added.
- The `rocm` metapackage sdist is installed with --no-build-isolation in step 2,
but setuptools only arrives in step 3, so on a fresh 3.12 venv it fails with
"Cannot import 'setuptools.build_meta'". Without that package there is no
rocm_sdk module for torch's _rocm_init to import. setuptools/wheel now go in
right after the venv is created.
- $ErrorActionPreference does not apply to native exit codes, so the failure
above was silent and the script printed "All done!" over an install whose GPU
check had already thrown "Torch not compiled with CUDA enabled". The check now
fails the script via $LASTEXITCODE.
dist/run-server.ps1:
- `ft serve` takes --port (dest server_port); --server-port is not a flag and
argparse rejects the whole command line.
- vcvarsall.bat was searched only under %ProgramFiles%\Microsoft Visual Studio.
Build Tools installs land in the x86 root, so the launcher silently ran with
no MSVC CRT on LIB. Ask vswhere first, then scan both roots, and warn when
it is genuinely absent.
- The previous run's logs were never cleared, and the `>` redirect only
truncates them once cmd reaches the `ft serve` line (after vcvarsall). The
readiness poll fires at 5s and matches the OLD traceback, so every launch
after a failure aborts instantly reporting the previous error as this one's.
Clear them up front; a log still locked by an orphaned worker now throws
something actionable instead.
python/freetoken/models/qwen3_5_moe/gguf.py — first real Qwen3.5 file through
the adapter (it had only been exercised against synthetic random-weight GGUFs):
- llama.cpp writes attention.head_count[_kv] as a per-layer array for hybrid
models, 0 on the GDN linear-attention layers: [0, 0, 0, 2, ...] for
full_attention_interval=4. parse_gguf_config called int() on the list and
raised TypeError. Collapse to the count the full-attention layers agree on.
- The dt_bias tensor is named bare `blk.N.ssm_dt` in real conversions, not
`ssm_dt.bias` — same bare spelling as its sibling `ssm_a`. Confirmed by shape
and dtype: (32,) F32, matching ssm.time_step_rank, one entry per value head.
Accept both spellings.
With these, the engine installs, initialises ROCm on gfx1201, and loads the
model through to expert-bank construction.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015bsqLButwcuTr7NL7wiWnP
…indows A 19.45 GiB Qwen3.5-35B-A3B Q4_K_M never finished loading under `--moe-backend offload`: the expert-bank build died at ~bank 38 with `hipHostRegister ... hipError 1` and availPhys at zero. Four distinct causes, each measured rather than guessed: 1. `expandable_segments` is HOST-backed on ROCm/Windows. The engine enables it by default before the first CUDA allocation. On this build, creating the segment mirrors the VRAM heap into the process working set: procWS went 0.39 -> 18.38 GiB on context creation alone, and a 4 GiB VRAM allocation added another 17.9 GiB. With it off, host WS stayed at 0.46 GiB through the same sequence. That is ~40 GiB gone before a single weight is read, leaving 8 GiB for a 19.45 GiB bank build. Now skipped on ROCm/Windows, with PYTORCH_ALLOC_CONF=expandable_segments:True to opt back in. 2. The GGUF bank fill held the checkpoint twice. Every expert byte is copied out of the reader's mmap, but the mapped source pages stay resident alongside the bank they were copied into -- measured at exactly 2x (0.97 GiB of working set per 0.49 GiB copied). `release_mapped_pages`/`PageReleaser` hand each consumed region back (Windows VirtualUnlock's documented working-set-trim side effect, MADV_DONTNEED on POSIX), byte-budgeted so a fill does not walk the mapping's page tables once per tensor. Same 8 layers now cost 3.73 GiB instead of 7.24. 3. `dequant_any`'s gguf-py fallback materialized whole tensors as fp32 -- a 1.9 GiB spike to produce a 0.9 GiB bf16 vocab head, at the point in the load where host RAM is scarcest. Quant blocks never span a row, so it dequantizes row slabs into the output instead (256 MiB cap). Verified bit-identical to the old path on 40 K-quant tensors. 4. A bare "cudaHostRegister failed" named no budget. The runtime returns hipErrorInvalidValue both for a bad argument and for "cannot lock these pages", which sent an earlier investigation chasing alignment, size, quotas and range overlap. The failure now reports banks pinned so far, bytes pinned, and both host budgets -- physical and commit, which differ sharply on Windows. The load also logs host memory at bank allocation and every 16 pins. Result: all 80 banks pin (19.45 GiB) with 23.68 GiB physical still free, and the model serves. The launcher had three more bugs on the path behind that one, each only reachable once the load got further: - ROCM_HOME was never exported, so tvm_ffi's JIT died at the first kernel build, long after a full model load. Note ROCM_HOME *only*: also setting ROCM_PATH makes clang look for the device bitcode at %ROCM_PATH%\amdgcn\bitcode, which is not the TheRock wheel layout (it lives under lib\llvm\), and every kernel build then fails to find it. - PYTORCH_ROCM_ARCH was unset, so torch compiled every JIT extension for all visible cards -- including the CPU's iGPU (gfx1036 on a 9800X3D), which is not the serving device. Known-not-fixed: the model loads and runs but generates incoherent text, which is a separate numerics bug in the qwen35moe GGUF adapter (the adapter had only ever been exercised against synthetic random-weight fixtures). The norm-scale and expert-bank-type conventions were checked and are correct; the GDN / attention conventions are not yet ruled out. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015bsqLButwcuTr7NL7wiWnP
…g key head A Qwen3.5-35B-A3B Q4_K_M GGUF loaded, served, and answered every prompt with fluent, grammatical, meaning-free text. The weights were not corrupt and no kernel was wrong; the value heads were in the wrong order. GDN has fewer key heads than value heads (16 vs 32 here), so every value-head- indexed axis needs a convention for which key head a value head belongs to. HF groups them -- [G0_v0..v(r-1), G1_v0..v(r-1), ...], value head h under key head h // r. llama.cpp's converter (conversion/qwen.py, _LinearAttentionVReorderBase) rewrites them tiled -- [G0_v0, G1_v0, ..., G0_v1, G1_v1, ...] -- so a plain ggml broadcast can stand in for an interleaved repeat. FreeToken's fla kernels are HF-convention on BOTH paths: chunk prefill (fla/chunk_o.py, i_h // (H // Hg)) and fused decode (fla/fused_sigmoid_gating_recurrent.py, i_hv // (HV // H)). Loading the GGUF order verbatim therefore pairs every value head with the wrong key head on every linear-attention layer -- 30 of this model's 40. The 10 full-attention layers are untouched, which is why output stayed fluent instead of collapsing into noise. _untile_v_heads inverts the converter's permutation, and it applies to all eight affected tensors or to none: attn_qkv's v rows, attn_gate (z), ssm_beta, ssm_alpha, ssm_a (A_log), ssm_dt (dt_bias), ssm_conv1d's v channels, and ssm_out's input columns. ssm_out moves out of _SUFFIX_MAP for that reason -- it needs dim=1. ssm_norm is per-head-dim and stays as it was. Un-tiling is the DEFAULT, not opt-in. The converter reorders unconditionally whenever num_k_heads != num_v_heads but only sometimes writes <arch>.ssm.v_head_reordered: this Qwen3.5 file has the key, and neither Qwen3.6 UD-Q4_K_XL nor Ornith-1.5 does. A file that explicitly sets it false is honoured. The arch prefix is read from general.architecture rather than hardcoded to "qwen35moe" for the same reason. Tests pin the inverse against a verbatim copy of the converter's permutation, for all six tensor shapes, with a negative control asserting the two orders actually differ (a no-op _untile_v_heads cannot pass). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018ZnDQ7zcqLT7omCX7Tk4KA
… accept -ExtraArgs
Five launcher bugs, each of which cost a debugging session on its own.
run-server.ps1
1. It demanded HIP_PATH or -RocmPath and told the user to go install a system
ROCm. There is no system ROCm on this port and there should not be: install.ps1
unpacks the TheRock wheels into the venv at site-packages\_rocm_sdk_core, which
is exactly the lib\llvm\... layout the env block below assumes, and HIP_PATH is
set by nothing. Ask the venv's python where _rocm_sdk_core lives, confirm
clang.exe is under it, and only then fail -- pointing at install.ps1 rather than
at AMD's download page.
2. The readiness poll treated any "Traceback" in the log as fatal. torch LOGS
tracebacks as warnings and carries on -- cpp_extension's "Error checking
compiler version" probe prints a full one on every ROCm start -- so the poll
killed loads that would have served fine. torch stamps its warning lines
([rank0]:W0830 ...) and a real crash does not, so match on the absence of the
stamp.
3. -ExtraArgs silently mangled its own documented form. `powershell -File` does no
PowerShell parsing, so -ExtraArgs "--moe-backend","offload" arrives as the ONE
literal element "--moe-backend,offload" and `ft serve` rejects the lot as a
single unrecognized argument. Split every element on commas and whitespace so
the array form and the space-separated string form both reach ft identically,
and say so in the usage header. Same bug in the -KVPages append.
4. The "log is still locked" guard blamed "a scheduler worker" -- true, but
grepping for `freetoken` finds nothing, because the holder is a python.exe
running multiprocessing.spawn. Name that, and point at stop-server.ps1.
stop-server.ps1
5. It orphaned every worker it was supposed to kill. It matched
CommandLine -match "serve|http.server 1420|freetoken", but the scheduler's
workers are spawned as
python.exe -c "from multiprocessing.spawn import spawn_main; ..."
which matches none of those. They survived, kept ~20 GiB of expert banks pinned,
and held the inherited serve.log/serve_err.log handles -- which then tripped
guard 4 on the next launch, so the visible symptom was never "stop didn't work".
Now: match the roots as before, then walk Win32_Process.ParentProcessId down
from each and kill the tree deepest-first (a parent killed early can spawn a
replacement on its way out). Workers already orphaned by an earlier partial stop
are swept separately, identified by ExecutablePath under this repo -- their
command line is generic, so anything else on the box using multiprocessing must
not be caught. Child links are only followed when the child started after its
claimed parent, since Windows reuses PIDs. Each kill reports its working set, so
"did that actually free the RAM" is answerable from the output.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018ZnDQ7zcqLT7omCX7Tk4KA
…r table
Ornith-1.5-35B-A3B declares `qwen35moe` and will not load: it reports
`block_count = 41` where the model has 40 layers. The 41st block is not a layer.
The header says so outright -- `qwen35moe.nextn_predict_layers = 1`, and
`blk.40.nextn.{eh_proj,enorm,hnorm,shared_head_norm}` sit alongside ordinary
`blk.40.attn_*` / `blk.40.ffn_*_exps` tensors -- so block 40 is a multi-token-
prediction draft head. Qwen3.5's own GGUF names the same thing `mtp.*`, at the top
level, where the `blk.` prefix check already dropped it; Ornith's converter numbers
it into the block sequence instead, so it walked straight into the layer loop.
That mattered three ways, none of which raised:
- `(i+1) % full_attention_interval` typed block 40 as linear attention, so its
`attn_q/k/v` failed the full-attention guard, missed `_IN_PROJ_SLOTS`, missed
`_SUFFIX_MAP`, and were dropped in silence -- the fused-group assert only catches
a group that is half filled, never one that is never started.
- Its routed experts joined the bank-type uniformity vote, which can promote (and
so requantize) both banks for a block that is never used.
- `load_q4_0_expert_sources` allocates `config.num_layers` banks and indexes them by
the tensor's own layer number, so block 40's experts index one past the end.
FreeToken has no speculative decoder anywhere in `engine/` or `scheduler/`, so an MTP
block can only be dropped. `_num_model_layers` subtracts `nextn_predict_layers` from
`block_count`, and the config, the weight iterator and the expert-bank loader all cut
at that.
Layer kinds now come from the tensor table rather than the interval pattern:
a block with `blk.i.ssm_*` is GDN linear attention, one without is full attention.
`full_attention_interval` describes a repeating pattern and mistypes any stack that
is not purely that pattern; the names are a direct statement and cost one header
read, which `_expert_types` was already paying. Files with no ssm tensors at all
(metadata-only shims, synthetic fixtures) keep the interval behaviour, and a
disagreement between the two is logged rather than hidden.
Finally, a `blk.` tensor inside the model stack that matches no rule is now an
error instead of a silent drop. That silence is what turned a wrongly-typed block
into "loads fine, answers nonsense" -- the same failure shape as the value-head
bug -- and it is the only reason this took a session to find.
Verified against all four qwen35moe GGUFs on disk: Ornith resolves to 40 layers
with experts 0..39, the three Qwen files are unchanged (40 layers, full attention
at 3/7/.../39, same bank types), and no file has an unhandled `blk.` tensor.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018ZnDQ7zcqLT7omCX7Tk4KA
…4.4 -> 37.9 tok/s Decode on this port ran the staged safe-copy fallback for every miss on every layer, and CUDA-graph capture aborted outright. Both trace back to one capability probe asking the wrong question. `OffloadMoeCache.should_use_safe_offload_copy` calls `inspect_host_mapping` to decide whether the GPU may dereference the host expert banks directly. That check demanded the packaged `_pinned_tensor` C++ extension and accepted no other evidence -- but the Windows/ROCm port never builds it (`ext_modules = []`, the same gap that dropped `_cpu_moe`). So the answer was structurally always "unmapped, unsafe", on a box where the mapping demonstrably works: hipHostRegister + hipHostGetDevicePointer translate a bank at 0x31356010000 to device 0x205c80000, and a D2D copy from that device address returns the correct bytes. `pinned.device_ptr` has resolved through the HIP runtime by ctypes since 83d2417, and `OffloadMoeCache` already builds `_copy_src_ptrs` from it -- the pointers the gather kernel dereferences and the pointers the safety check inspected had drifted apart. The probe now uses the same resolution, reports which mechanism answered (`mapping_source=extension|hip_runtime|host_ptr_identity`), and still fails closed when none can translate: an unregistered bank raises from hipHostGetDevicePointer and is called unsafe, exactly as before. Consequences, measured on Gemma-4-26B-A4B QAT q4_0, `--moe-backend offload --moe-cache-size 2048`, median of 3 x 256 tokens: eager + staged safe copy (the old behaviour) 24.4 tok/s eager + fused zero-copy gather 37.9 tok/s graphs + fused zero-copy gather 10.4 tok/s Answers stay coherent on the fast path (capital of France, the all-but-9 sheep trick, first 8 primes, a translation), so the gather is reading the right expert rows. The second consequence was the capture abort. `_safe_copy_plan` reads the miss count with `num_indices.item()`; under `torch.cuda.graph` that is hipErrorStreamCaptureUnsupported, which invalidates the capture and kills the backend worker during startup. So the probe -- not any gfx1201 driver bug -- is what made `--cuda-graph-max-bs 0` mandatory. Graphs now capture. They are also a 3.6x LOSS here, which inverts the standing assumption that the graph crash was where the missing speed lived; the launcher header now says so. Two build fixes were needed to get the fused kernel to compile at all: - `DEFAULT_CUDA_CFLAGS` decided "is this a HIP build?" from `HIP_PATH`, i.e. from whether a launcher had exported an environment variable, not from the toolchain. Ask `torch.version.hip`. Otherwise clang gets `--expt-relaxed-constexpr` and refuses it. - Nothing outside `dist/run-server.ps1` set `HIP_PATH` at all, and the patched `tvm_ffi` decides the ENTIRE Windows toolchain from it -- clang++ vs cl for the host compile, and whether `-lamdhip64` reaches the link. Without it the build failed with unresolved HIP symbols and no hint about the cause. `ensure_rocm_env` fills in `HIP_PATH`/`ROCM_HOME` from the venv's `_rocm_sdk_core` (never `ROCM_PATH` -- that sends clang to a bitcode path the wheel does not have), and names the GPU family from `gcnArchName` rather than the hardcoded gfx1201, since `rocm_agent_enumerator` is not in the wheel and torch's default is gfx906. `fast_index_copy_multi` now builds from a plain shell with nothing exported. Also: `run-server.ps1` resolves a relative `-Model`. The generated runner does `cd /d %TEMP%` first, so `models\foo.gguf` missed, and transformers then read it as a Hugging Face repo id -- the error talked about repo-id character rules and never mentioned the path. Test suite over tests/moe is unchanged: identical failure set before and after (all pre-existing, needing hardware/flashinfer this box lacks), plus four new tests covering a HIP-runtime mapping being accepted, an untranslatable one still being refused, and the no-mechanism-at-all case. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018ZnDQ7zcqLT7omCX7Tk4KA
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Six fixes that take GGUF MoE offload on Windows/ROCm from "does not load" to working and roughly 1.55x faster. All measured on a single box — Radeon RX 9070 XT (gfx1201), Windows 11, Python 3.12, ROCm from the
_rocm_sdk_corewheel (no system ROCm install). I have no other AMD hardware, so the Linux and non-gfx1201 paths are untested by me; the changes are written to be no-ops off Windows, but that claim deserves a second pair of eyes.The two correctness bugs are the important ones
Value heads were paired with the wrong key heads on 30 of 40 layers (
f48edc5). llama.cpp's converter rewrites every value-head-indexed axis from HF's grouped order into a tiled one. FreeToken'sflakernels are HF-convention on both paths, so without un-tiling at load, every value head reads the wrong key head. The symptom is the nasty kind: fluent, grammatical, completely meaningless output — nothing crashes, so it reads as a bad model rather than a bug. Un-tiling has to be the default, because upstream reorders unconditionally but only sometimes writes the<arch>.ssm.v_head_reorderedkey that would let you detect it.block_countincludes MTP draft blocks (9132d3c).<arch>.nextn_predict_layerstrailing blocks are next-token-prediction heads, not model layers. Qwen3.5 names themmtp.*and they were already dropped; Ornith numbers them into the block sequence asblk.40.*, so they walked straight into the layer loop. This was the real blocker on Ornith-1.5-35B-A3B — the symptom looked like "an extra full-attention layer", which sent me the wrong way for a while. Layer types now come from the tensor table rather than being inferred.The performance fix
The offload copy-path probe failed closed on the wrong evidence (
c95ca83) — 24.4 → 37.9 tok/s, a permanent ~1.55x tax.should_use_safe_offload_copyasked whether the_pinned_tensorC++ extension had been built. This port never builds it, so the answer was structurally always "unmapped, unsafe", and every decode miss took a staged host-side copy — on a machine wherehipHostRegister+hipHostGetDevicePointerdemonstrably work. Meanwhile the fast path's own pointer table was already being built from the ctypes HIP fallback, so the safety check and the pointers the kernel actually dereferenced had drifted apart. Both now come from one resolution.Worth noting for anyone who has hit the gfx1201 CUDA-graph capture crash: it was never a driver bug. The safe-copy path calls
num_indices.item(), which is illegal under stream capture. Fixing the above makes graphs capturable — and they are then a 3.6x loss (10.4 vs 37.9 tok/s eager), so--cuda-graph-max-bs 0stays the right default, but it is now a tuning choice rather than a workaround.Windows plumbing
Host memory blew up on HIP context creation (
abc6ebd) — working set jumped 0.39 → 18.38 GiB before any weights were loaded, which made GGUF MoE offload unloadable. On ROCm/Windows PyTorch'sexpandable_segmentsis host-backed. Worth suspecting in any ROCm/Windows torch process that mysteriously eats host RAM.Toolchain detection was reading the launcher, not the toolchain (
116543e,55a27ac). The patchedtvm_ffidecides the entire Windows toolchain fromHIP_PATH— clang++ vs cl, and whether-lamdhip64reaches the link — butHIP_PATHwas only ever set bydist\run-server.ps1. Every other entry point failed to build kernels with unresolved-HIP-symbol errors that named nothing about the cause.kernel/_toolchain.py::ensure_rocm_envnow derives it from the venv's_rocm_sdk_coreand takes the arch fromgcnArchName(there is norocm_agent_enumeratorin the wheel, and torch's default of gfx906 silently produces dead kernels). The launcher also now stops the whole process tree — it was reporting success while leaving a worker holding the log handle, which made the next launch report the previous run's error.Measured
--moe-backend offload --moe-cache-size 2048 --cuda-graph-max-bs 0, median of 3 x 256 tokens:Tests
tests/moe/test_safe_offload_copy.pyis new and covers the probe/pointer-source agreement — the drift above was invisible precisely because nothing asserted the two came from the same place.tests/moe/test_decode_trace.pyis extended.One thing I could not fix
--moe-backend cpu|hybridstill raisesImportError: cannot import name '_cpu_moe'; this port's6d7d6c4setext_modules = []. Reviving it needs HIP substitutes forcudart/cudaLaunchHostFunc(amdhip64,hipLaunchHostFunc). Left alone deliberately rather than half-done.