Repository navigation
POC: Experts offload - #50717
Draft
tarekziade wants to merge 5 commits into
Draft
POC: Experts offload#50717tarekziade wants to merge 5 commits into
tarekziade wants to merge 5 commits into
Conversation
…evice cache
`from_pretrained(..., experts_offload={"cache_bytes": ...})` keeps the experts
of MoE models on disk and runs them through a shared device cache with the
existing experts kernels.
Optional quantized experts run through gather_qmm from kernels-community/mlx-quantization-metal-kernels.
For now this is MPS-specific
tarekziade
marked this pull request as draft
November 13, 2026 12:26
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
…r dense layers - Quantized experts are read into recycled pinned host buffers, through one descriptor per shard. - The expert -> slot tables live on the host; a layer sends its slots in one small copy. - Sorted prefill runs 64 experts at a time, so a long prompt that routes to every expert stays bounded. - from_pretrained accepts a MetalConfig (dense nn.Linear layers only) with experts_offload.
…guessed reads PyTorch's MPS stream keeps small kernels in an uncommitted command buffer until the next synchronization, so after each layer's routing copy the GPU idled while the host encoded the next layer. Record an event (commits without waiting) after each decoder-layer sub-block and after the experts (mps_commit, on by default). Qwen3.8-Flash-Next decode +10% on an M5; OLMoE unchanged output. install_ahead now skips guessed reads that are still in flight instead of waiting for them before the next layer's launches.
…bytes, host_preload) A pool of expert rows in pinned host memory between storage and the device cache: a device-cache miss costs one host-to-device copy instead of a storage read. For discrete GPUs with more host than device memory; on unified memory the device cache already is host RAM. Rows handed out stay reserved until the copy that reads them has run (next synchronization) or the prefetch is dropped. Qwen3.8-Flash-Next on an A10G (3-bit experts, 8 GiB device cache, experts on a network file system): decode 4.06 -> 5.8 tok/s with every expert preloaded.
This branch has not been deployed
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.
What does this PR do?
See https://tarekziade-moe-expert-streaming.static.hf.space/index.html