feat(moe): per-layer host-bank residency - #112
Merged
Merged
Conversation
…yers, auto selection under capped pin quota
Collaborator
|
@jason-fxz LGTM |
avlp12
pushed a commit
to avlp12/FreeToken
that referenced
this pull request
Aug 24, 2026
Brings in from FlashML-org/FreeToken main (5 commits since our merge base): bd372b6 fix(checkpoint): FTW truncated shard raises OSError instead of silent garbage load (FlashML-org#25) e0a3bbc Merge pull request FlashML-org#112 from FlashML-org/feat/split-residency 831d38a fix(moe): report the residency banks actually settle at eebb3f5 feat(moe): estimate expert-bank bytes from the model config c41833b feat(moe): per-layer host-bank residency -- split lock-CPU/pin-GPU layers, auto selection under capped pin quota 184a4f1 build(kernel-cache): add sm_80 (A100/A800) to default arches (FlashML-org#75) Conflicts (all in files where upstream's split-residency work and our own MoE/offload patches touched the same lines) were resolved to keep both intents, not to pick a side: * moe/host_banks.py: upstream's split-residency HostBank (backing="cuda" born-pinned option, LOCKED/PinPipeline/ResidencyPlan machinery) merged with our explicit 4096-byte alignment invariant (_off/_reg_len, aligned memoryview()/pin()). The alignment carve now applies to BOTH the mmap and the cuda backing, not just mmap. Also fixed a latent bug the combination would otherwise have introduced: HostBank.lock() mlock'd len(self._buf) (the over-allocated padded buffer) starting at the aligned self.addr, which overruns the mmap by up to one block; now mlocks self._reg_len like pin() already did. * moe/expert_banks.py: our bank-alignment startup diagnostic (_log_bank_alignment/_ALIGN) and upstream's parallel-reader support flag (_PARALLEL_READER_SUPPORTED) are independent additions next to each other -- kept both. In load_expert_banks, kept upstream's requested_residency()/_echo_residency() wrapping and added our _log_bank_alignment(banks) call right after the build, before echoing residency back to the caller. * models/deepseek_v4/moe.py: kept both prefill_routed guards -- our speculative-verify-wearing-a-prefill's-clothes check (route to _decode_routed under hybrid decode) runs first, then upstream's crossover-size-or-unpinned-layer check (route to super()._prefill_routed) -- both gate independently to their own fallback path. * moe/offload_cache.py: in copy_missing(), reordered so the unpinned (LOCKED) whole-layer pageable materialize check (upstream) runs before the --moe-copy-engine DMA doorbell dispatch (ours): an unpinned layer has no device alias, so it must never reach the doorbell path, which assumes a pinned bank layout. * engine/engine.py: took upstream's version of the cpu_layer_ids resolution -- it's a strict superset of ours (adds _auto_cpu_layers selection under a capped pin budget, which is the feature this sync exists to pull in); our `n_moe` local was just an unused-elsewhere alias, nothing of ours was lost. Also (per the task): rewrote sync-upstream.sh to `git merge` instead of `git rebase` -- uran-custom already carries its own merge commit (pr69/pr70/pr71), and rebasing across an existing merge commit invites history corruption. Verified in worktree /root/ft-sync (PYTHONPATH override, not the live editable install at /root/FreeToken): - python3 -m compileall python/freetoken: clean - import freetoken + the 7 touched modules: clean - /root/test_dma_doorbell.py (eager + graph phases): PASS - /root/test_bank_alignment.py: PASS - /root/test_moe_copy_engine_gate.py: PASS - /root/test_dma_dead_flag.py: PASS - ad hoc HostBank backing="cuda" (born-pinned) + lock() mlock-length fix smoke test: PASS Our uran-specific features (host_kv_tier.py, DmaCopyService + gating, alignment invariant, --moe-collect-stats, scheduler idle hooks) are unaffected outside the files above and behave identically. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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
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.
Why
#55: on WSL/Windows, WDDM caps pinned (page-locked) memory to ~50% of RAM -- MoE models whose
host bank sources exceed the cap cannot be served through the offload backend.
Workaround
Per-layer host-bank residency: keep only the GPU layers' banks pinned; the remaining layers'
banks are OS-locked (
mlock, no pin quota) and decode on the CPU MoE executor.Changes
--moe-cpu-layers(id list / count / fraction)--moe-backend hybridtoo (was offload-only)0opts outFREETOKEN_PIN_BUDGET_GB: override the pin budget (default: WSL -> 40% RAM, else uncapped)--moe-backend cpuover budget: degrade to all-locked instead of dying in cudaHostRegisterhooks) honor them ambiently via
requested_residency()-- no loader signature changescopy plan skips their rows, decode routes to the CPU executor
_cpu_moeABI / expert format); otherwise stay fully pinned on the GPU path
Testing
VmLckin every configuration; FTW vs raw identical greedy outputs for the same split