cuda: Windows MoE pin budget - layer-subset routing + chunked pinned store

[jasonlnheath]

Complements #76 (bounded pin budget) with a Windows-tested layer-subset router and a chunked pinned-store allocator, plus the WDDM measurements behind them.

Problem

On Windows the expert cache's full-set pin silently falls back to plain CPU buffers when the driver refuses the allocation — all-zero tensors/covered counters, no warning, uncached speed. WDDM caps page-locked host memory far below what a full expert-set pin needs.

Measured on this rig (RTX 5070 Ti 16GB, 64 GB RAM, CUDA 13.1, Windows 11)

• Bare process: aggregate cudaHostRegister ~= 28 GB obtainable
• Inside a loaded server: per-call limit collapses to ~2 GB, aggregate ~2-4 GB total per process
• Full Flash-Next UD-Q3_K_XL pin = 52.2 GiB request, never grantable on 64 GB WDDM (verified with 52.9 GB free)
• Pin pool is sticky: once spent, even 256 MB fails; pool is shared with the CUDA context (~24 GB/process ceiling on bigger boxes)
• Driver probe: PageableMemoryAccess=0, HostRegisterSupported=1

What this patch adds

1. GGML_CUDA_MOE_PIN_LAYERS subset routing — pin a chosen subset of layers when the full expert set cannot fit, instead of nothing. 66/66 tensors claimed on Flash-Next UD-Q3_K_XL. Includes a scope fix: the pattern string must outlive the loader (otherwise a dangling pointer silently drops claims).
2. Adaptive chunked cudaHostRegister allocator — ladder of shrinking chunk sizes. Doubles as a diagnostic of what a given box/driver actually grants, turning silent fallback into measurable data.

Honest scope

This does not reach the full cached path on 64 GB WDDM — the in-process grant is 2-4 GB versus the ~24 GB the cached path wants. Where it helps: bigger-RAM Windows boxes (128 GB: AnimaLoraStudio measured the ~50% cap; vetchahoo's box pins fine) and anyone diagnosing WDDM grant limits. Possible follow-up: tensor data pointers are 64B- rather than 4KB-aligned — cudaHostRegister wants 4K alignment, so alignment fixup may improve grant sizes.

The compute-sanitizer report and full env diagnostics were shared on Discord earlier; happy to attach here too.

🤖 Generated with Claude Code

[GenerelSchwerz]

The Windows measurements are useful, and smaller registration retries could help #76. This patch needs a few fixes before we can integrate it. Reviewing head 343a03be4:

• Non-Windows compilation: the free callback is outside _WIN32 but calls ggml_cuda_moe_cached_win_pinned_free, which is defined only inside it. An extracted C++ compiler check reproduces the undeclared-function error. Guard the callback and its selection consistently.
• Registration ownership: g_moe_win_pin_chunk is shared across allocations, but each allocation can choose a different chunk size. If A registers in 1 GiB chunks and a later B selects 4 GiB, freeing A skips registrations. Concurrent allocations also race on this variable. Store registration information per allocation and use the original registered addresses during cleanup, including failure unwinding. Check unregister/free failures instead of silently discarding them.
• Pinning opt-out: the Windows path calls the new allocator before the helper that checks GGML_CUDA_NO_PINNED, so it can pin memory despite that setting. Check it before either allocation path.
• Layer placement: selecting the first N layers does not itself place the remaining layers on CPU. They follow other overrides or normal placement, so the log's unconditional CPU claim is misleading and GPU placement can consume unexpected VRAM. Make the placement requirement explicit and validate/bound the environment value before constructing the regex. The pattern string's lifetime is correct as written.

For integration with updated #76 (6b0b05807), the chunked allocator changes the legacy allocation path; it does not improve #76's bounded source-registration path when an explicit pin cap is supplied. That path currently attempts one registered prefix and falls back to staging if it fails. A smaller-prefix retry there seems worth testing under the existing budget and ownership. I would keep layer-selection policy separate from that change.

Please attach the native Windows commands, allocation logs and sanitizer results mentioned in the description, with completed inference/output comparisons and repeated allocation/free coverage. We still need evidence that the proposed registration change improves the bounded path. The checks above are source review plus an extracted compiler check, not a completed CUDA build or native Windows runtime validation.

AI disclosure: Codex assisted with this review and comment under repository-owner direction.

[jasonlnheath]

Thanks for the detailed review. All four source issues are fixed in 0aed729, the branch is rebased onto moe-cache@d30efee88, and native Windows evidence is below.

1. Non-Windows build

The free callback and the win_chunked flag now live inside the _WIN32 guard; the iface assignment is a default plus a guarded override. Linux never sees a Windows-only symbol. The ubuntu CUDA workflow needs maintainer approval to run for first-time fork contributors and has not started on this push, so your extracted compiler check is the reference - it should now pass.

2. Registration ownership

g_moe_win_pin_chunk is gone. Each allocation now records its registered segments:

• struct moe_win_pin_segment {addr, bytes}, collected per chunk during registration
• a mutex-guarded registry (base, size, segments), same pattern as the mmap range registry already in this file (moe-cache.cu:334)
• free and failure-unwind walk the segments in reverse, at the original registered addresses
• every cudaHostUnregister and VirtualFree return is checked; on any unregister failure the VA range is deliberately leaked and logged, never freed while registered
• the ladder probe's unregister is checked too - a silent failure there would corrupt the next real registration of the same range

Why a registry instead of per-buffer state: ggml_backend_cpu_buffer_from_ptr uses buffer->context as the data base for the shared CPU iface (ggml-backend.cpp:2527), so per-allocation metadata cannot ride on the buffer without a custom iface. Pin and unpin are serialized by a mutex; the WDDM quota is process-wide.

Test coverage: test-moe-cache gained interleaved alloc/free cycles (512 MiB + 64 MiB allocated together, freed out of order, 3 cycles) with a --win-pin-only mode. Passes on this rig, and runs clean under compute-sanitizer memcheck (ERROR SUMMARY: 0 errors).

3. GGML_CUDA_NO_PINNED

Checked at the top of the Windows allocator, before VirtualAlloc. With the variable set there are zero registration attempts and the model generates on plain CPU buffers, same as the other paths honoring the variable.

4. Layer placement

• strtol + endptr validation; non-positive or garbage logs a warning and falls back to the all-layers pattern
• clamped at 512 with a warning; the cap only bounds the regex since a count above the model depth matches nothing. A true n_layer clamp is not reachable here: overrides are passed to the llama_model_loader constructor before hparams exist
• the log now says remaining layers keep their normal placement (--n-gpu-layers window, or --cpu-moe / --n-cpu-moe where they match)
• the "regardless of --cpu-moe" warning is now scoped to the pin-layers case
• kept the single bounded alternation rather than per-layer overrides: the loader builds a std::regex per tensor per override entry (llama-model-loader.cpp:1235), so N per-layer entries multiply failed compiles across every non-expert tensor. Happy to switch if you prefer the convention
• pattern-string lifetime untouched as you confirmed

Rebase

Replayed onto d30efee. The src/llama.cpp hunk is re-expressed on the proc-address buffer-type acquisition from 12a4d1d; nothing else in that block changed.

Evidence (Windows 11, RTX 5070 Ti 16 GB, 64 GB RAM, CUDA 13.1, sm_120, Flash-Next UD-Q3_K_XL)

run	result
PIN_LAYERS=2, cache + --cpu-moe	pin engages ("Windows chunked pin engaged"), blk.0-1 expert tensors to CUDA_MoE_Cached, blk.2+ to CUDA_Host; first generation hits an illegal memory access surfacing at cudaFuncSetAttribute after a mm_ids_helper launch
PIN_LAYERS=6	same as above
PIN_LAYERS=48	ladder fails at quota, clean unwind, CPU fallback, generates at 23 t/s, exit 0
GGML_CUDA_NO_PINNED=1	zero pin attempts, generates at 8 t/s, exit 0
PIN_LAYERS=abc, =-1	"not a positive integer - ignoring", all-layers pattern, graceful fallback, exit 0
PIN_LAYERS=100000	"clamped to 512", exit 0

Greedy comparison (--cpu-moe reference vs cache-enabled with pin failure fallback): byte-identical outputs on 3 prompts x 120 tokens, temp 0.

The PIN_LAYERS=2 crash is the partial-pin crash class the original commit documented: routing is verified correct in the logs, the subset pins fully, and the fault appears only when the pinned-subset cache path engages at generation. It did not reproduce on the pre-Sep-8 base this patch was developed against, so it most likely interacts with the grouped-decode changes merged since. I have not root-caused it yet; the win-pin allocator itself is memcheck-clean, so I believe the allocator is not the origin.

Full-model compute-sanitizer is not feasible on this box: the sanitizer's own allocations push the WDDM pinned quota over during model load (cudaMallocHost OOM under instrumentation). The Sep 7 targeted-repro sanitizer report in lv5-console-for-gen.zip (previously shared on Discord) remains the memcheck evidence for the split-staging stream-memop issue - and note the Windows test suite currently fails at the split-staging sync test (test-moe-cache.cpp:10982) with the same signature. The suite also needs a 32 MB stack on MSVC (default 1 MB overflows in the certificate tests).

#76 interaction

Kept separate as you suggested. I test-drove a smaller-prefix retry in ggml_backend_cuda_moe_pin_sources on 6b0b058:

• test first: injected failure via a fail_source_register_after counter (same pattern as fail_stage_after); on unpatched code the single-shot path registers nothing and falls to staging (red)
• with a halving retry (floor at the 64 KiB registration granularity), the test lands at the half and quarter prefixes (green)
• real model, --moe-expert-cache-host-pinned-mb 3072: full 55.8 GB registration fails, the ladder registers a 1.93 GiB contiguous prefix and generation completes; unpatched, the same run pins nothing and stages everything

Patch (retry loop + the two test hooks) is at jasonlnheath@868a29314 on top of 6b0b058 - or I can send the diff directly, whichever you prefer. One contiguous registration per source is preserved; your budget and ownership structures are untouched; layer policy stays out.

[jasonlnheath]

Follow-up on the PIN_LAYERS=2 illegal access from my earlier comment - root-caused, and it is not this patch's allocator.

What it is not (all verified on the rig, logs available):

• not cuda: fail closed grouped decode ROUTE to cached mmid #75: reproduces identically on 12a4d1d + only the original pin commit
• not the win-pin allocator: every H2D from the pinned buffer was checked with a per-operation stream sync (hundreds of acquires, main and sibling prefetch paths) - all clean; the alloc/free cycle test is memcheck-clean; registration flags (Mapped|Portable vs Default) make no difference; pool size and L2 (disabled via LLAMA_ARG_MOE_EXPERT_CACHE_L2_PINNED_MB=0) make no difference
• not pointer math: all acquire host pointers verified inside the registered range

What it is: the fault is asynchronous and outside the calling thread's streams. Every main-thread sync point reports clean, then the next CUDA API call (cudaFuncSetAttribute inside launch_mm_ids_helper, serving a CUDA_Host layer's prefill MMQ) returns the sticky illegal access. The crash appears exactly when a pinned expert cache engages during prefill; the same routing with pageable fallback buffers (GGML_CUDA_NO_PINNED=1) generates fine at 8 t/s, and full-set fallback generates fine at 23 t/s.

That converges with two things already on record: the Sep 7 sanitizer report (cuStreamWriteValue32_v2 to a host-mapped stage_ready flag in ggml_cuda_moe_cache_prepare_split_staging - stream memops against host-mapped memory are WDDM-restricted) and the Windows test suite failing today at the split-staging sync test (test-moe-cache.cpp:10982) with the same signature. Pageable sources take plain copies and never enter that path; pinned sources do.

So: GGML_CUDA_MOE_PIN_LAYERS=2 routes layers 0-1 experts to the cached buffer type, the pin engages (first configuration on a 64 GB box that ever engaged the cache), the first multi-row evaluation splits into staged waves, and the split-staging stream memop faults. The subset routing itself is verified correct in the load logs.

Repro: GGML_CUDA_MOE_PIN_LAYERS=2 llama-cli -m Flash-Next-UD-Q3_K_XL -ngl all -fit off --cpu-moe --moe-expert-cache-size 80 --load-mode none --lazy-mode on -c 4096 -st -n 16 -p "..." - faults at first generation on Windows/CUDA 13.1/WDDM, RTX 5070 Ti.

Until the split-staging path gains the device-memory-flag/event-sync treatment discussed for it earlier, the safe Windows story for this patch remains: pin budget + ladder as diagnostics, cache fallback and NO_PINNED paths as the runtime modes. Happy to hand over the instrumented logs (per-op traces, ~490 sync points) if useful for the fix.

[GenerelSchwerz]

Follow-up on this: I introduced fixes based off of this fork in a commit a few days ago. If you'd like to test again and see if this has been removed, that would be great.