Repository navigation
Host-weight prefetch in the scheduler, faster MoE cache prompt and decode, cache seeding under llama-server - #268
Open
danielhanchen wants to merge 88 commits into
Open
danielhanchen wants to merge 88 commits into
danielhanchen wants to merge 88 commits into
Conversation
…for large ubatches, seed the cache from the end of a prompt
…sched Weights in host buffers that are used by GPU splits were uploaded on the compute stream after a full synchronization, so the transfers never overlapped compute. The scheduler now binds the weight copies of each such split to one of `depth` staging slots of a device arena it owns (outside of the graph allocator, so the copies never alias other tensors) and uploads them on a second backend instance of the same device, up to depth - 1 splits ahead. A copy into a slot waits for the event recorded after the last compute that read the slot, and the compute waits for the event recorded after the copy. Slot addresses are fixed per graph, so graph reuse and CUDA graphs are unaffected. - enabled for CUDA, ROCm and MUSA devices with async copies and events, GGML_SCHED_PREFETCH=0 disables it, =2 allows any backend with the caps - GGML_SCHED_PREFETCH_DEPTH (default 2) staging slots per backend - with a copy callback, MUL_MAT_ID weights of batches smaller than GGML_SCHED_PREFETCH_MIN_BATCH (default 256) keep the callback, which uploads only the routed experts (into the slot) - copies read by more than one split stay with the graph allocator - disabled with pipeline parallelism (n_copies > 1) - the staging memory is included in ggml_backend_sched_reserve_size and ggml_backend_sched_get_buffer_size - new ggml_backend_sched_set_prefetch_callback to exclude host tensors that change between graph computations; llama uses it to exclude the MoE cache slot maps
With the MoE cache on, the scheduler prefetch uploaded every host expert tensor whole from host memory for large ubatches, so the experts already held in the cache were sent over PCIe again. The prefetch callback now returns a mode per weight (NONE, SYNC, ASYNC), and ASYNC copies go through the copy callback on the copy queue backend. For experts, llama copies the cached ones device to device from the cache banks and uploads only the others, issued ahead of the split like any prefetch. Its ids-free predicate (sched_copy_all_experts) decides both the mode and the copy, so an ASYNC copy never reads the routing ids. - the ubatch that seeds the MoE cache (outputs, > 32 tokens) needs the routing counts, so its cached experts are SYNC: they stay in the staging slot and are filled at the split after the ids readback, as before - GGML_SCHED_PREFETCH_MIN_BATCH moved to llama (it is a property of the copy callback); GGML_SCHED_PREFETCH=0 restores the ids-based copy - the copy queue waits for an event recorded on the compute backend at the start of each graph, so copies that read device memory (cache banks) are ordered after the writes of earlier graphs
…r large batches, issue them before the uploads
The MoE cache resolved each layer of a decode step on the host: the slot map was a host weight, so every MoE layer started a new split (a full sync in ggml_backend_sched_copy_input), then the copy callback read the routing ids back (a second sync), planned the LRU and uploaded the misses from pageable memory. Decode ran at ~11 ms/token whatever the hit rate. Add GGML_OP_MOE_CACHE_LOOKUP (CUDA only): one warp looks up the slots of the selected experts in a device slot map, gives the misses a slot with a CLOCK sweep (hits and this lookup's slots are never evicted) and a second kernel copies the missing experts from mapped pinned host memory into their slots. The matmuls read the banks with the returned slots as before, so the results are bit-identical to the host lookup. No host round trip, no split per layer: one CUDA graph per token. - llama_moe_cache::lookup builds the op for ubatches of <= 32 tokens (decided by the ubatch size, so the output-only last layer agrees with the others); prepare() moves the state between the host LRU (still used by large batches) and the device when the path changes. - host experts are mapped with cudaHostRegister(Mapped|ReadOnly) or used as is if already pinned; pageable (HMM) access is refused, it is ~10x slower. mmap'ed files usually cannot be pinned, so --moe-cache-mib loads the model without mmap when the load mode is auto. - LLAMA_MOE_CACHE_DEVICE=0 restores the host lookup; backends without the op or failed mapping fall back to it.
With --cpu-moe the per-expert scales (gemma4 ffn_down_exps.scale) and biases of the cached layers stay in host memory, so their get_rows ran on the CPU and every MoE layer of a decode step went back to the host. The cache now keeps a device copy of these small tensors and the lookup path of small batches reads them, bit-identical values, no host round trip.
- the device lookup is decided per cache group (one per device and expert layout): a group whose backend lacks the op or the ggml_backend_host_map proc keeps the host lookups, the others use the device, so mixed and multi-GPU setups work - the host mapping is exposed as the generic ggml_backend_host_map / ggml_backend_host_unmap procs; the CUDA implementation builds for CUDA, HIP and MUSA (vendor aliases added); pinned memory is detected with the pointer attributes, pageable memory is registered, never read through HMM - device copies of per-expert scales and biases have no size limit - --moe-cache-mib with load mode auto loads without mmap only if the model files take at most 80% of system memory, larger models keep mmap and the host lookups
…che D2D and uploads) in the stats With the scheduler prefetch the expert fill runs on the copy queue, so the number of copy calls per ubatch is the cost to watch: the 4 MiB minimum cache run keeps it near one upload per gap instead of one per scattered expert.
… expert tensor After seeding, the cached experts sit in scattered slots, so using them for a large ubatch needs one device copy per cached run plus one upload per gap (thousands of copies per ubatch). The 4 MiB minimum run avoids that by uploading almost everything from host again, so the cache did not cut PCIe bytes on the prompt path. New ggml_backend_tensor_copy_regions_async: queues many (src, offset, dst, offset, size) copies as one call. CUDA implements it with cudaMemcpyBatchAsync (CUDA >= 12.8) for device sources of the same device and pinned host sources, in stream order; it returns false (nothing copied) otherwise (HIP, MUSA, older CUDA, pageable host memory, GGML_CUDA_COPY_BATCH=0), and the callers keep the previous per-copy path. llama issues the fill of each expert tensor (every cached run from the cache banks plus the uploads of the rest) as one batch, on the backend the scheduler passes (the copy queue when prefetched), and the seed fills of the cache slots as one batch per tensor. The stats report copies and batched calls per ubatch.
Conflict in src/llama-moe-cache.cpp: keep the large-batch helpers of the prompt path (count_batch, count_upload, seed, seed_end) next to the device lookup, and log_stats becomes non-const (it reads the device counters). seed() admits experts through moe_cache_lru::plan, which bumps lru.gen, so the next device lookup uploads the seeded host state; the seed_end rollback now bumps lru.gen too.
- user inputs reached through views (view_src chain with GGML_TENSOR_FLAG_INPUT) are never prefetched: they may be overwritten as soon as the computation is queued (B200 probe: stale result 99 vs 1) - host weights written by the graph itself (a node or a view of one) are never prefetched: the copy must follow the writer (probe: stale 720) - a computation that stops on an error drains the copy queues before returning, so a copy issued ahead cannot write a slot after it is reused - if the graph allocator fails after the staging arena was allocated, the arena is freed, prefetch is disabled on that backend and the allocation is retried without it (reserve and alloc_graph) - only splits whose copies are all issued ahead get a slot; splits with SYNC copies use the graph allocator, as without prefetch, so no staging memory or copy queue is created for them - a copy or prefetch callback set after the graph was allocated: its bound copies use the regular path (current callbacks) until the next allocation binds again - SYCL always keeps the regular path (pinned staging workaround for mmap, per-tensor extras that the arena would accumulate)
- ggml_backend_host_map (CUDA) reference counts the buffers it registers, so a second context mapping the same buffer holds its own reference and the first context's unmap no longer unregisters memory still in use - registration tries read-only first and falls back to a read-write registration (B200 reports HostRegisterReadOnlySupported=0, so pageable buffers were never mapped) - memory measurement (no_alloc) enables the device lookup only for groups whose device has the host_map procs and supports the op, like the real initialization - the slot sizing counts the device copies of the per-expert scales and biases and the lookup map/state inside the cache budget - get_slot_map takes the ubatch size: only ubatches that use the device lookup skip the host slot map, so the output-only last layer of a large ubatch reads the cache again as on master (first-token logits now match master byte for byte) - copies from the cache and seeding skip a group whose device state is newer than the host LRU (e.g. a warmup ubatch of 32 tokens that did not use the device lookup), instead of reading stale slots - lookup kernel: 32-bit epoch wrap clears the pins, out-of-range ids read expert 0; ggml_moe_cache_lookup asserts that the distinct experts fit in the slots - ggml-rpc.h: GGML_OP_COUNT 102, RPC_PROTO_PATCH_VERSION 1 - load mode auto: the RAM check uses the available system memory (MemAvailable on Linux, free + inactive + purgeable pages on macOS), the CPU backend reports free = total there
This reverts commit 521ca75.
…ough views or host weights written by the graph A view of a GGML_TENSOR_FLAG_INPUT tensor was prefetched on the copy queue, so an overwrite by the caller after graph_compute_async returned changed the result. A host weight written by a node of the graph (e.g. a ggml_cpy into a weight view) was uploaded before the split that writes it ran. The candidate check now walks the view_src chain for the input flag and skips host weights that any non-empty node writes, directly or through a view, so these copies keep the regular path after their producer.
…th SYNC copies A split whose copies need results of earlier splits (SYNC, e.g. the routed experts of a small MoE batch) still took a staging slot: depth slots sized for it and a copy queue that never issued anything. Such splits are no longer bound, their copies are allocated by the graph allocator as without prefetch, so the slots are sized only for the splits whose copies are issued ahead and the copy queue is created only if there is one.
…ging memory If the staging arena was allocated but the graph allocator then failed, reserve and alloc_graph returned false while still holding the arena, unlike a failed arena allocation, which falls back to the regular path. The copies are now unbound, the arenas freed and prefetch disabled, and the graph allocation is retried once, on both the reserve and the alloc_graph paths.
…on stops early compute_splits returned on a compute error or abort with copies already issued up to depth - 1 splits ahead. They kept writing their slots after the scheduler was reset and could land after a later write to the same slot, or read host memory after the caller regained control. The copy queues are now synchronized before returning the error.
…after the graph was allocated The copies to prefetch are selected when the graph is allocated, so a copy or prefetch callback set afterwards was ignored for that graph (e.g. a prefetch callback returning NONE for a weight the caller overwrites). When a callback changes, the bound copies of the allocated graph are made by the regular path with the current callbacks, into their slots, until the next allocation selects them again.
…he arena on each bind SYCL could be opted in with GGML_SCHED_PREFETCH=2, but the async copies skip the pinned staging that its set_tensor uses for weights in mmap'ed memory, so SYCL now always keeps the regular path. The staging arena is also reset before the copies of a new graph are bound, as the graph allocator does with its buffers, so per-tensor state set by init_tensor does not accumulate.
…tive splits With depth 2 the arena held two slots of the largest prefetched split, but only two consecutive splits are live at once (the copy of split k + 1 overlaps the compute of split k). Even ordinals are now placed at the start of the arena and odd ordinals end at its end, so the arena only needs the largest sum of two consecutive splits. E.g. a MoE layer with 425 MiB gate/up and an 850 MiB Q8_0 down tensor needs 1275 MiB instead of 1700 MiB, with the same overlap. Other depths keep uniform slots. The log line reports the staging size and the largest split.
…nce per split Without pipeline parallelism each user input of a split was copied with a backend synchronize plus a synchronous tensor_set (2 host syncs per input). Queue the host-resident inputs as async uploads on the split backend and synchronize once after all of them, before host weights are copied. For a hybrid MoE model with ~30 graph inputs this drops decode from 82 to 9 stream synchronizations per token.
… batch ggml_backend_tensor_copy_regions_async put the device-to-device copies of a MoE cache fill (runs of cached experts) in the same cudaMemcpyBatchAsync as the uploads, so they took copy-engine time from the uploads: nsys of a warm 9.3k prompt (gemma-4-26B-A4B, 4000 MiB cache) shows 65.8 GB D2D using the copy engines for 658 ms next to 2.16 s of uploads. The device copies of a call now run as one kernel (variable-size regions, pointers in the kernel parameters, 128 regions per launch, vector width from the common alignment, at most one 256-thread CTA per SM), issued before the batched uploads on the same stream; the regions are disjoint, so their order does not matter. GGML_CUDA_COPY_REGIONS_KERNEL=0 puts them back in the batch. Host memory registered through ggml_backend_host_map (reference counted) also counts as pinned for the uploads, not only CUDA host buffers.
…uted experts, small ubatches are not prefetched - the copy of all experts without reading the ids is only for the copies the scheduler issues ahead on its copy queue; on the regular path (GGML_SCHED_PREFETCH=0, a staging arena that did not fit, or a split that is not prefetched) the ids are read and only the routed experts are copied, as on master. Copying every expert of a 2048-token ubatch made prompt processing ~7% slower with prefetch off (since the moe-cache prompt path). - ubatches smaller than GGML_SCHED_PREFETCH_MIN_BATCH (default 256) bind no host weights to the staging arena, so generation does no staging, event or copy queue work. - the copy queues are synchronized only when work was queued on them since the last synchronization.
The prefetch copy queue of the scheduler is a second CUDA backend instance that never computes a graph. Its batched copies now keep the device-to-device regions on the copy engines (cudaMemcpyBatchAsync) instead of the SM kernel. With the kernel there, generation after a prefetched prompt was ~5% slower (Qwen3.8-Flash-Next IQ4_XS, MoE cache 42688 MiB, ub 2048, B200: 71.9 -> 76.1 tok/s, equal to GGML_SCHED_PREFETCH=0), at the same prompt speed. Compute instances (the regular copy path) keep the kernel.
…its it The prompt seed only ran on a large ubatch with outputs. llama-server runs the last tokens of a prompt as a separate small batch for recurrent and SWA models (context checkpoints), so the ubatch with outputs had <= 32 tokens and the cache was never seeded there. The last large ubatch of each batch now seeds the cache, with or without outputs. The seed still fills the slots from the experts staged for that ubatch (no extra uploads). To keep multi-batch prompts from waiting for the ids of every batch, the routing of each large ubatch of a cached layer is kept as a graph output and read before the next graph: a later batch of the same prompt seeds from it while its copies of all of the experts stay asynchronous. Only a batch with outputs, or the first large batch of a prompt, reads its own ids first. The stored routing is dropped at any ubatch with outputs, a multi-sequence ubatch, a position jump (context switch, checkpoint restore) or a failed ubatch, and at every graph build. The graph is rebuilt when the seed mode changes, as the scheduler plans the expert copies with it. LLAMA_MOE_CACHE_SEED=0 disables the seed.
The seed of a prompt end copies slots from the experts staged on the device for its ubatch. With the scheduler prefetch fix, the regular path (prefetch off, a ubatch below GGML_SCHED_PREFETCH_MIN_BATCH, a split that is not prefetched) stages only the routed experts, so: - a seed uses the stored routing of an earlier ubatch only when the copies staged every expert (prefetched), otherwise the ids of the ubatch itself; - seed() takes the staged experts: the plan admits only those, and an expert tensor of the layer that did not stage a planned expert does not fill (the layer is dropped by seed_end); - a last large ubatch below the prefetch threshold seeds from its own ids, the routing is not recorded with GGML_SCHED_PREFETCH=0 (no extra sync); - LLAMA_MOE_CACHE_SEED_CHECK=1 (debug) compares the seeded slots with the host experts after each seeding ubatch.
…unless the device has HBM 90a34bc moved the device copies of the prefetch copy queue to the copy engines, which helped generation on a B200 (~5%) at the same prompt speed. On an RTX PRO 6000 (GDDR7, sm_120) the copy engines made a prefetched 8k prompt with the MoE cache 16% (ub 2048) / 9% (ub 4096) slower, with no generation change (2k: 90.3 vs 90.2 tok/s). The copy queue now uses the kernel unless the memory bus is at least 2048 bits wide (HBM); GGML_CUDA_COPY_QUEUE_KERNEL=0/1 overrides.
Member
Author
|
@codex security review |
🛡️ Codex Security ReviewSecurity review completed. No security issues were found in this pull request. Reviewed commit: Only the user who started this review can view the report in Codex. ℹ️ About Codex security reviews in GitHubThis is an experimental Codex feature. Security reviews are triggered when:
Once complete, Codex will leave suggestions, or a comment if no findings are found. |
|
Codex Review: Didn't find any major issues. Swish! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
…nts ungated experts
When ggml_gallocr_reserve_n cannot allocate a buffer, it has already freed the old one but keeps the node and leaf assignments of the new layout. A later ggml_gallocr_alloc_graph of a graph with the same nodes then finds nothing to reallocate and places the tensors in the freed buffer (NULL), which crashed the scheduler retry after a failed graph reallocation with the host weight prefetch disabled. The assignments are now dropped, so the next allocation reserves again.
…refetch staging The staging arena of the host weight prefetch (depth 2 x the largest split, about 117 MiB for Qwen3.8-27B and Qwen3-14B at -ub 2048) was part of the compute buffer a --fit probe measures. Near a layer boundary this moved one dense layer to the host (Qwen3.8-27B at 16 GB: 54/66 layers instead of 55, Qwen3-14B at 8 GB: 23/41 instead of 24, decode -4%). Without the MoE cache a measurement now leaves the staging out, so --fit places the same layers as without the prefetch, and a reserve that leaves less than 512 MiB free on a device with a staging arena drops the prefetch (GGML_SCHED_PREFETCH_HEADROOM_MIB, ggml_backend_sched_set_prefetch_headroom), like a reserve that does not fit with it, so the arena only takes the memory that is left. With the MoE cache the measurement keeps the staging and the cache sizing leaves room for it.
…synchronously With GGML_SCHED_PREFETCH=2 (any backend with async copies and events) a host split input of a --split-mode tensor graph was queued as an async upload into a meta buffer, whose async set only handles some split states (fatal error in ggml_backend_meta_set_tensor_async). A meta buffer now takes the copy it took before (cpy_tensor_async, else the synchronous copy). The default mode (CUDA family only) already skipped the meta backend, which has no registry.
With more than one GPU the cache was shrunk by the total deficit only: a total surplus could hide one device short of memory, the per-device check then fell back to the old placement, which kept the shrunk cache and moved whole layers to the CPU (Qwen3.6-35B-A3B on 2 x T4 at 9 + 9 GB, --fit on --moe-cache-mib 16384: 28/41 layers, 12 GB of cache, decode 9.8 tok/s vs 22 tok/s for --fit on without the cache). The shrink now also removes, for a device that misses its target, its deficit scaled by its share of the cache (measured against the probe without the cache), and disables the cache when a device misses its target without any cache memory to give back, so the fallback places the layers without the cache. (cherry picked from commit 1e4df4cf99a23b78b1714b3af753da45fb4c6981)
…e, not the defaults
…es with the prefetch off Subtracting the staging arena after the prefetch binding had removed the copies from the graph allocator left neither counted, so a run that drops the prefetch could lack the memory its copies need. The measurement now runs with the prefetch off, which is what a fit without the prefetch measures; ggml_backend_sched_get_staging_size is no longer needed.
…not only in reserve ggml_backend_sched_alloc_graph allocated the graph next to a staging arena without checking the headroom that ggml_backend_sched_reserve keeps for the memory the backends allocate while computing (e.g. the CUDA pool), so the compute could fail with GGML_STATUS_ALLOC_FAILED. Check it when an arena was allocated for the graph and after the graph buffers are reserved again, and drop the prefetch like reserve does. Nothing changes with the prefetch off. test-alloc: test_prefetch_headroom with a mock device of limited memory.
…cannot map The device lookup maps the host experts into the device address space with ggml_backend_host_map (cudaHostRegister / hipHostRegister). A model loaded with mmap keeps them in a read-only file mapping, which often cannot be registered: on ROCm hipHostRegister becomes a KFD userptr that is always faulted in for writing (the read-only flag is the coarse-grained flag on HIP), and on CUDA a read-only mapping needs cudaHostRegisterReadOnly, which not every device or driver supports. The cache then fell back to host lookups for every layer, losing most of the gain of the cache under mmap. When mapping a group fails, its memory-mapped expert tensors are copied once into one anonymous host buffer owned by the model (what --load-mode none does for every weight) and the mapping is retried, if they fit in the available system memory (MemAvailable within the cgroup limit, or the available physical memory on Windows) less max(2 GiB, 1/16). The other tensors stay memory-mapped. LLAMA_MOE_CACHE_MATERIALIZE=0 keeps the host lookups. test-moe-cache-materialize checks the copy and repointing on CPU.
The automatic load mode drops mmap for the MoE cache only when the model files fit in 0.8 of the available system memory; the decision is now logged, so a load that keeps mmap (and relies on the device lookup copying the experts it cannot map) can be told apart.
…erts The memory check before copying memory-mapped experts into anonymous host memory only read the memory.max of the cgroup v2 leaf of the process. It missed the limits of the ancestor cgroups and of cgroup v1, counted all of memory.stat:file as reclaimable although it includes tmpfs and shared memory, and returned 0 (unknown, so the copy went ahead) when the cgroup was exhausted. The copy could then exceed a cgroup limit and be OOM-killed. Walk every cgroup v2 and v1 memory hierarchy of the process up to its root, as common_available_memory does for the automatic load mode, count only the non-shmem page cache as reclaimable, and return 1 for exhausted memory so the copy is skipped with a warning and the host lookups are kept. The Linux check takes a root prefix for /proc and /sys/fs/cgroup so that test-moe-cache-available-memory can feed it fake hierarchies.
…the mapped expert copy memory.stat counts mlock'd file pages in file (v2) and total_cache (v1) and again in unevictable / total_unevictable; with mmap and mlock the source pages stay locked while the copy is made, so they are no longer credited as reclaimable.
…KEND_DL builds link
…model memory breakdown
…apped expert copy common_available_memory() kept its own Linux cgroup helpers, which count reclaimable page cache (a warm GGUF) as used and give up without /proc/meminfo, so the load mode could pick mmap while the mapped expert copy found the memory. Both now use llama_moe_cache_available_memory(), which also reads the free, inactive and purgeable pages on macOS instead of the total memory the CPU backend reports.
- user-fixed layers: the cache is shrunk to the margins too, else the context was reduced to make room for it - the cache shrunk for the user placement at the full context is measured again at its requested size for the placements with the experts in system memory - the handoff probe splits the layers over the devices and checks every device against its own target instead of the total surplus - if steps 3 and 4 cannot place every layer in device memory, the placement with all experts in system memory and the cache is used instead
Splits invalidated by a new copy or prefetch callback are copied on the compute queue, but n_async kept the copy queues waiting on the compute queues of every later computation of the allocated graph. test-alloc checks that the recomputed graph makes no copy queue wait.
…backend_cuda_host_map ggml_cuda_host_map maps memory pinned by someone else without registering it, but the batched copies only accepted registered memory and fell back to one copy per region. Memory pinned for the device of the copy or portable is accepted too, pageable memory is still rejected.
…e the larger batch The small-batch reserve case reserved the same graph that alloc_graph then allocated with a 256 batch, which GGML_SCHED_DEBUG_REALLOC (set by CI) counts as an unexpected reallocation and aborts on, also with the prefetch off. llama.cpp reserves the largest batch of each graph, so the reallocation next to a staging arena happens when the graph size changes: the reserve graph gets one more node. The test now also counts unexpected reallocations through the realloc callback and asserts none, with or without GGML_SCHED_DEBUG_REALLOC.
- the probe of the placement with all layers in device memory and the experts of one layer read the tensor_split buffer after it was restored, so it measured the saved split instead of its own - measuring its own split (proportional to the memory each device can use) is not better: the output layer, the compute buffers and the shares of the cache are per-device costs, a split by layers misjudges them and a failed probe returned all experts in system memory where steps 3 and 4 keep experts in device memory - several devices: steps 3 and 4 always run, their placement is replaced by the one with all experts in system memory and the cache only if it cannot keep every layer in device memory - one device: unchanged, the probe is the first placement of step 3
This was referenced Oct 19, 2026
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
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.
Faster MoE and dense inference when weights live in host memory (partial
-ngl,-ot/--cpu-moe,--fitspill), and a faster--moe-cache-mibcache. Everything is generic across architectures and on by default, with kill switches.What changes
1. Host-weight prefetch in
ggml_backend_sched(ggml/src/ggml-backend.cpp)wait_streampattern as Unsloth's offloaded gradient checkpointing.--fitand the memory breakdown count it.GGML_SCHED_PREFETCH_MIN_BATCH, default 256) copy whole expert tensors ahead without the routing-ids sync. Smaller ubatches keep the existing routed-only copy.GGML_SCHED_PREFETCH=0disables it. One log line reports depth, staging size and eligible splits.2. MoE cache, prompt path (
src/llama-moe-cache.cpp, newggml-cuda/copy-regions.cu)cudaMemcpyBatchAsyncplus a copy-regions kernel per expert tensor, instead of a full-layer copy on the compute stream.GGML_CUDA_COPY_QUEUE_KERNEL=0/1overrides this.3. MoE cache, decode path (new
GGML_OP_MOE_CACHE_LOOKUP,ggml-cuda/moe-cache.cu)LLAMA_MOE_CACHE_DEVICE=0falls back to the host lookup.4. Cache seeding at the end of a prompt
LLAMA_MOE_CACHE_SEED=0disables seeding.LLAMA_MOE_CACHE_SEED_CHECK=1verifies every seeded slot byte for byte (debug).5. Host overhead: host split inputs are uploaded asynchronously, with one synchronize per split.
The four new files carry the Unsloth AGPL-3.0 SPDX header. Upstream files keep their license.
Results
RTX PRO 6000 Blackwell (Colab G4)
--fit, load-mode noneGGML_SCHED_PREFETCH=0GGML_SCHED_PREFETCH=0B200
Correctness
-ngl 12and Qwen3.6-35B-A3B--cpu-moewith the cache off and on, ub 512 and 2048, mmap and-lm none. Perplexity logits,llama-debuglogits and 128-token greedy text are byte-identical to master in all 12, with the prefetch line logged in every run.Not covered
Base is upstream b11577
cbe28bb3c(base/upstream-cbe28bb3c). The branch merges b11577 to absorb upstream ggml-org#30329 (scheduler creation moved intollama_context::sched_new()). Like every pin inscripts/unsloth/pr-set.json, this is merged onto the aged upstreamb####tag for the prebuilt and dropped once upstream has the equivalent.