Explore zero-copy mmap-backed expert paging #77

New Issue

hugo · 2026-08-30T10:34:34Z

hugo commented

2026-08-30 10:34:34 +00:00

Idea

Explore a zero-copy, mmap-backed routed-expert paging mode between the current fully mapped/resident path and explicit SSD streaming.

The model GGUF would remain a single read-only mmap. After each layer's router selects experts, Metal would receive page-aligned newBufferWithBytesNoCopy views for only the selected (layer, expert) gate/up/down ranges. A bounded cache would retain useful views across tokens; evicted ranges could receive POSIX_MADV_DONTNEED after all GPU work using them has completed.

This is an experiment, not yet a commitment to a third production mode. The first goal is to determine whether avoiding pread copies improves GLM-5.2 inference without losing the predictable memory bound of the existing streaming cache.

Why this may work

Gguf::open already mmaps the complete model file.
The Metal backend already wraps page-aligned mmap ranges with newBufferWithBytesNoCopy.
Routed expert offsets and sizes are already available through StreamExpertTable.
The current SSD path already reads selected expert IDs asynchronously, maintains a bounded layer/expert cache, tracks in-flight command buffers, and exposes cache metrics.
File mappings may be larger than physical memory; only the active pages need to be resident. The important limit is the dense weights, KV/cache state, scratch space, and selected routed-expert working set, not the virtual mapping size.

For GLM-5.2 the access unit is a layer-specific expert, not a global expert file: the current shape has 79 layers, 3 leading dense layers, 256 experts, and top-8 routing. Splitting the artifact into one file per expert would either create thousands of layer/expert files or load unrelated layers together. GGUF tensor offsets already provide the required random-access ranges.

Non-goals

Do not split or regenerate GGUF files for the initial spike.
Do not add a new crate, cache abstraction, UI preference, or persistent configuration before the spike passes.
Do not remove or weaken the current resident or pread streaming paths.
Do not attempt CUDA/Mojo support in this issue; first establish the behavior on the existing macOS/Metal backend.
Do not evict every expert after every token. The next token's experts are unknown until routing, and immediate eviction would discard useful temporal locality.

Phase 1: reproducible baseline

Before changing the loader, record resident and SSD-streaming baselines on the same Mac with the same hashed GGUF, prompt tokens, context, sampling settings, prefill chunk, power setting, cache budget, and background load.

Use at least:

GLM-5.2 as the memory-constrained target.
DeepSeek V4 Flash as a resident/streaming regression reference.
A short decode workload and a representative long-prefill workload.

Run sequential repetitions with cooldown. Do not run DS4 and DS4Server simultaneously. Record:

prefill tokens/s and time to first token;
stable decode tokens/s and per-token p50/p95 latency;
peak process footprint, Metal tracked-live bytes, macOS memory pressure, and swap delta;
page faults and SSD bytes read;
streaming cache hits, misses, evictions, waits, pread bytes, and pread time;
memory retained immediately after generation and after a fixed idle interval.

Define the minimum worthwhile improvement before reviewing spike results. It must exceed measured run-to-run noise; a suggested gate is at least 5% improvement in stable decode throughput or p95 token latency versus pread streaming at the same effective working-set budget.

Phase 2: narrow mmap expert-view spike

Add one hidden, explicit experimental switch to the existing SSD-streaming path. Reuse its selected-ID handoff, expert table, budget, cache policy, metrics plumbing, and command-buffer lifetime tracking.

For a cache miss:

Validate the selected expert and gate/up/down ranges exactly as the current loader does.
Page-align each range and create or reuse a no-copy Metal view over the existing model mmap.
Keep the underlying Gguf mapping alive for longer than every Metal view.
Bind only the selected expert views to the routed-MoE operation.
Keep views alive until the last command buffer that references them has completed.
Evict only when the configured expert/byte budget requires it.
After safe eviction, optionally issue POSIX_MADV_DONTNEED for pages not shared with another live view.

Start without speculative prefetch or a new cache algorithm. If synchronous GPU page faults dominate, test the smallest next step: WILLNEED/readahead after routing and before binding. Do not add prediction unless profiling proves that ordinary readahead cannot overlap enough I/O.

Page alignment is a correctness boundary. Adjacent expert ranges may share a VM page, so an eviction must not advise away a page still referenced by another live view. Group or reference-count only those overlapping pages if the existing layout actually requires it.

Phase 3: correctness and safety gates

Resident, current SSD-streaming, and mmap-view modes must produce identical greedy tokens for the existing DeepSeek and GLM comparison fixtures.
Preserve DS4 behavior for routing, quantization, prefill/decode ordering, cache accounting, checkpoints, MTP, and cancellation.
Exercise cache hit, miss, eviction, wrap/reuse, cancellation, generation teardown, and model teardown.
Add the smallest focused test for page-aligned range calculation and shared-page eviction safety.
Run the normal repository commit gates.
A failed view creation, invalid range, or unsupported Metal limit must fail clearly or fall back only through an explicit, observable policy; never silently change numerical behavior.

Phase 4: performance decision

Repeat the Phase 1 matrix with the mmap-view mode at exactly the same effective cache budget.

Promote the experiment only if all of the following hold:

token-level correctness and existing tests remain green;
peak memory stays within the configured working-set target plus the previously measured non-routed runtime overhead;
no sustained critical memory pressure or new swap growth appears;
mmap-view performance beats the current streaming path by more than run-to-run noise on GLM-5.2, or provides an otherwise clearly measured latency/memory advantage;
short and long prefill do not show pathological fault thrashing;
generation cancellation and teardown release all Metal views safely.

If it does not pass, keep the current pread cache and close this idea with the benchmark evidence. A zero-copy path that is merely different but not measurably better is not worth maintaining.

Phase 5: production implementation, only after a passing spike

Replace the hidden switch with a small explicit expert-loading mode in engine settings while preserving resident and pread streaming as supported choices.
Reuse the existing cache budget and diagnostics rather than creating parallel settings.
Report mmap-view faults/readahead/advice separately from pread bytes so the UI and metrics do not mislabel zero-copy paging as copied streaming.
Keep the current explicit streaming path as the fallback for machines/workloads where GPU page faults are slower or less predictable.
Document which hardware/model combinations passed the gate and keep comparisons hardware-specific.
Add user-facing configuration only after the final default/auto-selection policy is backed by measurements.

Separate expert files or a new sharded artifact format should be a later issue only if profiling proves that page alignment, Metal VM-region limits, or GGUF tensor layout prevents the single-file approach from meeting the gate.

Likely code areas

src/engine/gguf.rs
src/engine/metal.rs
src/engine/metal/gpu.rs
src/engine/metal/glm.rs
native/metal/ds4_metal.m
src/metrics.rs

Completion

This issue is complete when the spike has produced reproducible A/B evidence and either:

the mmap-view mode passes the gates and the production phases above are implemented and documented; or
the idea is rejected with the measured reason and no experimental path remains enabled by default.

## Idea Explore a zero-copy, mmap-backed routed-expert paging mode between the current fully mapped/resident path and explicit SSD streaming. The model GGUF would remain a single read-only mmap. After each layer's router selects experts, Metal would receive page-aligned `newBufferWithBytesNoCopy` views for only the selected `(layer, expert)` gate/up/down ranges. A bounded cache would retain useful views across tokens; evicted ranges could receive `POSIX_MADV_DONTNEED` after all GPU work using them has completed. This is an experiment, not yet a commitment to a third production mode. The first goal is to determine whether avoiding `pread` copies improves GLM-5.2 inference without losing the predictable memory bound of the existing streaming cache. ## Why this may work - `Gguf::open` already mmaps the complete model file. - The Metal backend already wraps page-aligned mmap ranges with `newBufferWithBytesNoCopy`. - Routed expert offsets and sizes are already available through `StreamExpertTable`. - The current SSD path already reads selected expert IDs asynchronously, maintains a bounded layer/expert cache, tracks in-flight command buffers, and exposes cache metrics. - File mappings may be larger than physical memory; only the active pages need to be resident. The important limit is the dense weights, KV/cache state, scratch space, and selected routed-expert working set, not the virtual mapping size. For GLM-5.2 the access unit is a layer-specific expert, not a global expert file: the current shape has 79 layers, 3 leading dense layers, 256 experts, and top-8 routing. Splitting the artifact into one file per expert would either create thousands of layer/expert files or load unrelated layers together. GGUF tensor offsets already provide the required random-access ranges. ## Non-goals - Do not split or regenerate GGUF files for the initial spike. - Do not add a new crate, cache abstraction, UI preference, or persistent configuration before the spike passes. - Do not remove or weaken the current resident or `pread` streaming paths. - Do not attempt CUDA/Mojo support in this issue; first establish the behavior on the existing macOS/Metal backend. - Do not evict every expert after every token. The next token's experts are unknown until routing, and immediate eviction would discard useful temporal locality. ## Phase 1: reproducible baseline Before changing the loader, record resident and SSD-streaming baselines on the same Mac with the same hashed GGUF, prompt tokens, context, sampling settings, prefill chunk, power setting, cache budget, and background load. Use at least: 1. GLM-5.2 as the memory-constrained target. 2. DeepSeek V4 Flash as a resident/streaming regression reference. 3. A short decode workload and a representative long-prefill workload. Run sequential repetitions with cooldown. Do not run DS4 and DS4Server simultaneously. Record: - prefill tokens/s and time to first token; - stable decode tokens/s and per-token p50/p95 latency; - peak process footprint, Metal tracked-live bytes, macOS memory pressure, and swap delta; - page faults and SSD bytes read; - streaming cache hits, misses, evictions, waits, `pread` bytes, and `pread` time; - memory retained immediately after generation and after a fixed idle interval. Define the minimum worthwhile improvement before reviewing spike results. It must exceed measured run-to-run noise; a suggested gate is at least 5% improvement in stable decode throughput or p95 token latency versus `pread` streaming at the same effective working-set budget. ## Phase 2: narrow mmap expert-view spike Add one hidden, explicit experimental switch to the existing SSD-streaming path. Reuse its selected-ID handoff, expert table, budget, cache policy, metrics plumbing, and command-buffer lifetime tracking. For a cache miss: 1. Validate the selected expert and gate/up/down ranges exactly as the current loader does. 2. Page-align each range and create or reuse a no-copy Metal view over the existing model mmap. 3. Keep the underlying `Gguf` mapping alive for longer than every Metal view. 4. Bind only the selected expert views to the routed-MoE operation. 5. Keep views alive until the last command buffer that references them has completed. 6. Evict only when the configured expert/byte budget requires it. 7. After safe eviction, optionally issue `POSIX_MADV_DONTNEED` for pages not shared with another live view. Start without speculative prefetch or a new cache algorithm. If synchronous GPU page faults dominate, test the smallest next step: `WILLNEED`/readahead after routing and before binding. Do not add prediction unless profiling proves that ordinary readahead cannot overlap enough I/O. Page alignment is a correctness boundary. Adjacent expert ranges may share a VM page, so an eviction must not advise away a page still referenced by another live view. Group or reference-count only those overlapping pages if the existing layout actually requires it. ## Phase 3: correctness and safety gates - Resident, current SSD-streaming, and mmap-view modes must produce identical greedy tokens for the existing DeepSeek and GLM comparison fixtures. - Preserve DS4 behavior for routing, quantization, prefill/decode ordering, cache accounting, checkpoints, MTP, and cancellation. - Exercise cache hit, miss, eviction, wrap/reuse, cancellation, generation teardown, and model teardown. - Add the smallest focused test for page-aligned range calculation and shared-page eviction safety. - Run the normal repository commit gates. - A failed view creation, invalid range, or unsupported Metal limit must fail clearly or fall back only through an explicit, observable policy; never silently change numerical behavior. ## Phase 4: performance decision Repeat the Phase 1 matrix with the mmap-view mode at exactly the same effective cache budget. Promote the experiment only if all of the following hold: - token-level correctness and existing tests remain green; - peak memory stays within the configured working-set target plus the previously measured non-routed runtime overhead; - no sustained critical memory pressure or new swap growth appears; - mmap-view performance beats the current streaming path by more than run-to-run noise on GLM-5.2, or provides an otherwise clearly measured latency/memory advantage; - short and long prefill do not show pathological fault thrashing; - generation cancellation and teardown release all Metal views safely. If it does not pass, keep the current `pread` cache and close this idea with the benchmark evidence. A zero-copy path that is merely different but not measurably better is not worth maintaining. ## Phase 5: production implementation, only after a passing spike 1. Replace the hidden switch with a small explicit expert-loading mode in engine settings while preserving resident and `pread` streaming as supported choices. 2. Reuse the existing cache budget and diagnostics rather than creating parallel settings. 3. Report mmap-view faults/readahead/advice separately from `pread` bytes so the UI and metrics do not mislabel zero-copy paging as copied streaming. 4. Keep the current explicit streaming path as the fallback for machines/workloads where GPU page faults are slower or less predictable. 5. Document which hardware/model combinations passed the gate and keep comparisons hardware-specific. 6. Add user-facing configuration only after the final default/auto-selection policy is backed by measurements. Separate expert files or a new sharded artifact format should be a later issue only if profiling proves that page alignment, Metal VM-region limits, or GGUF tensor layout prevents the single-file approach from meeting the gate. ## Likely code areas - `src/engine/gguf.rs` - `src/engine/metal.rs` - `src/engine/metal/gpu.rs` - `src/engine/metal/glm.rs` - `native/metal/ds4_metal.m` - `src/metrics.rs` ## Completion This issue is complete when the spike has produced reproducible A/B evidence and either: - the mmap-view mode passes the gates and the production phases above are implemented and documented; or - the idea is rejected with the measured reason and no experimental path remains enabled by default.

hugo added the idea label 2026-08-30 10:34:34 +00:00

hugo referenced this issue

2026-08-30 10:56:57 +00:00

Reach current DS4 performance parity on Apple Metal from the Rust execution path #78

Sign in to join this conversation.

1 Participants

Notifications

Due Date

No due date set.

Dependencies

No dependencies set.

Reference: hugo/DS4Server#77