perf(moe): pool heterogeneous GGUF cache rows

[pedro-moser]

Important

This is a draft stacked on #102. Until #102 lands, the GitHub diff against main includes its Laguna GGUF commits. The heterogeneous-cache changes to review are the final two commits (2c6591c and a9f8824, 15 files). I will rebase onto main after #102 merges.

Summary

• store mixed-type GGUF expert banks at each layer's native row width instead of padding every host row to the largest layer
• partition the existing max-stride GPU bank arenas into disjoint exact-geometry decode pools with independent tags and LRU state
• route GGUF decode views by layer_id while preserving a full-layer max-stride prefill overlay
• add direct pinned-host compact-row to padded-CUDA-row JIT copies without payload-sized CUDA temporaries
• dispatch raw GGUF BF16 expert rows through the dense BF16 path instead of MMVQ
• prevent bandwidth-profile auto-selection from choosing hybrid for GGUF, whose experts have no CPU executor
• add opt-in cumulative, rank-local decode miss/H2D telemetry

Closes #194 after #102 lands.

Motivation

Laguna-S Q4_K_M has 47 MoE layers but two very different expert geometries:

layers	format	gate/up row	down row	total/expert
39	Q4_K	3,538,944 B	1,769,472 B	5.0625 MiB
8	BF16	12,582,912 B	6,291,456 B	18 MiB

#102's safe uniform layout pads every layer to the 18 MiB maximum. On this checkpoint that means:

• host expert banks: 211.5 GiB, which prevented Laguna-S validation on 128 GiB RAM
• a nominal 298-slot GPU cache covers only 298 of 470 expert selections per token
• worst-case routed H2D traffic: 3.33435 GiB/token
• transfer roofline at the measured ~48 GB/s gather bandwidth: ~13.41 tok/s
• observed original decode: 11.24 tok/s

Compact host rows reduce expert-bank RAM to 85.359375 GiB. Geometry pools spend the same GPU arena bytes on many more Q4_K rows rather than 18 MiB padding.

Design

OffloadMoeCache remains the single owner and public facade:

1. Validate heterogeneous sources only for contiguous 2-D uint8 GGUF rows. Uniform formats retain their existing shape/dtype/contiguity contract.
2. Price each legacy bank arena from that bank's largest layer row, preserving the external moe_cache_size denomination.
3. Group layers by the tuple of per-bank row bytes and carve disjoint exact-width views from those arenas.
4. Maintain independent slot maps, reverse maps, ages, LRU stats, descriptors, and copy plans per geometry.
5. Route decode bank_views(layer_id=...), ensure_experts, and copies through the selected pool.
6. Preserve full-layer prefill via a shared max-stride overlay. Geometry pools are invalidated before the overlay is written, and the optional prefill hit-D2D path is disabled in geometry mode because the overlay may overwrite rows still needed as sources.
7. Rebuild tears down aliases before reallocating arenas and recreates all views and metadata.

Backend auto-selection also stays on offload when a saved bandwidth profile recommends hybrid for GGUF. Explicit backend choices are unchanged; this only prevents an automatic post-load failure in a CPU executor that has no GGUF expert implementation.

The row-strided CUDA path copies compact registered-host rows directly into padded destinations. A peak-memory regression verifies it does not allocate a CUDA temporary proportional to the payload.

Controlled A/B

Hardware: RTX 5080 16,303 MiB (SM 12.0), driver 610.57.04, CUDA toolkit 13.3, Ryzen 7 9800X3D, 128 GiB RAM. PCIe expert gather measured at 47-48 GB/s.

Checkpoint: unsloth/Laguna-S-2.1-GGUF, revision 750f92f90cf54159c4d7a610cb7b3e74498e75c6, Laguna-S Q4_K_M GGUF (96,031,829,760-byte local file).

Both benchmark arms ran on the same #102 + #103 integration branch and differed only in cache layout. #103 supplied Q8_0 KV for the benchmark but is not a code dependency of this PR.

ft serve --model <Laguna-S-2.1-Q4_K_M.gguf> \
  --tp-size 1 --attn triton --kv-cache-dtype q8_0 \
  --num-tokens 8192 --kv-reserve-tokens 8192 \
  --max-running-requests 1 --cuda-graph-max-bs 0 \
  --memory-ratio 0.95 --moe-backend offload \
  --moe-cache-size 480 --disable-moe-prefill-overlap

Identical prompt, sampling, one warm request, and 247 completion tokens. The client used the public streaming API and measured decode as (completion_tokens - 1) / (last_token_time - first_token_time) so the first token is not counted before the timing window:

cache layout	effective slots	client decode	TTFT
unified max-stride	480 × 18 MiB	15.595 tok/s	2.15 s
geometry pools	988 Q4_K + 202 BF16	19.203 tok/s	2.07 s

Warm decode improved 23.13% at the same cache-byte budget. The geometry branch also measured 18.686 tok/s on a cold 127-token request.

Tests

On current main merged with #102, without #103:

77 passed

across cache budgeting, Laguna mixed-weight loading, MoE/offload behavior, the strided CUDA kernels, CLI stats, and scheduler reporting.

9 passed

for isolated Laguna config/module tests.

Additional verified coverage includes:

• direct real Laguna-S Q4_K_M load and API generation
• Q4_K/BF16 pool carving, disjoint storage, independent LRU, rebuild, and prefill invalidation
• strided copy payload/stride correctness with FREETOKEN_FUSED_COPY enabled and disabled
• no payload-sized CUDA temporary
• GGUF BF16 reinterpret/dispatch
• non-GGUF uniform-shape and contiguity regressions
• exact H2D accounting and fetched-vs-CPU split for uniform GPU, geometry GPU, hybrid, pure CPU, and mixed CPU/GPU modes

tests/moe/test_prefill_hit_d2d.py is 4 passed / 1 failed on both this branch and the rebased #102 baseline: the optional cudaMemcpyBatchAsync probe copies wrong bytes on this driver. Geometry mode intentionally does not use that path.

Compatibility, limitations, and risks

• Direct GGUF only. FTW conversion still does not persist per-layer GGUF types.
• Uniform expert formats keep the legacy cache path.
• A 65,536-token Q8_0 configuration does not fit the tested 16 GiB card. Real bring-up was validated at 32K; the controlled layout A/B was 8K.
• Decode telemetry is cumulative and rank-local. Under overlap scheduling it may include look-ahead work from the next batch; it is not attributed exclusively to the just-finished request.
• The performance claim is only the Laguna 480-vs-480 A/B above. Cross-model llama.cpp numbers are intentionally excluded.