perf(moe): avoid max-stride padding for heterogeneous GGUF expert banks

[pedro-moser]

Summary

The mixed-quantization GGUF expert-bank layout introduced by #102 pads every layer in a bank to the largest per-layer row stride. That is safe for kernel addressing, but it makes Laguna-S Q4_K_M impractical on 128 GiB hosts and wastes most of the GPU expert-cache budget.

This issue proposes keeping the host banks compact per layer and carving the existing GPU bank arenas into independent exact-geometry decode pools. The external moe_cache_size byte budget remains unchanged.

Reproduction geometry

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

The 47 MoE layers contain:

• 39 Q4_K layers: 3,538,944-byte gate/up row + 1,769,472-byte down row = 5.0625 MiB per expert
• 8 BF16 layers: 12,582,912-byte gate/up row + 6,291,456-byte down row = 18 MiB per expert
• 256 experts per layer, top-10 routing

Padding every host row to 18 MiB makes the expert banks 211.5 GiB. Compact per-layer rows require 85.359375 GiB.

On the GPU, a unified 18 MiB slot cache with 298 slots is smaller than Laguna's 470 expert selections per token. With a cold cache, the exact routed payload is 3.33435 GiB/token. At the measured ~48 GB/s PCIe gather bandwidth, the transfer-only roofline is ~13.41 tok/s; 11.24 tok/s was observed in the original 298-slot bring-up.

Proposed design

• Preserve one OffloadMoeCache facade and the existing byte-denominated bank arenas.
• Store GGUF host rows compactly per layer.
• Group GGUF layers by the tuple of bank row sizes.
• Carve each bank arena into disjoint exact-width views, with independent tags and LRU state per geometry.
• Route decode views and LRU updates by layer_id.
• Keep full-layer prefill on a shared max-stride overlay; invalidate geometry pools before the overlay is written.
• Copy compact pinned-host rows directly into wider CUDA rows with a strided JIT kernel, without a payload-sized CUDA temporary.
• Restrict heterogeneous rows to contiguous 2-D uint8 GGUF sources; uniform formats retain their current path.
• Treat raw GGUF BF16 expert rows as BF16 views and dispatch them through the dense expert kernel instead of MMVQ.
• If the saved bandwidth profile recommends hybrid, keep auto-selection on offload for GGUF because the CPU MoE executor has no GGUF expert implementation. Explicit backend choices remain unchanged.
• Add opt-in cumulative, rank-local decode miss/H2D telemetry so mixed geometry and mixed CPU/GPU transfer behavior can be validated without attributing overlapping look-ahead work to one request.

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.

Both arms used the same Laguna-S Q4_K_M checkpoint, TP=1, Triton attention, Q8_0 KV, 8,192 KV tokens, one running request, --moe-cache-size 480, CUDA graphs disabled, prefill overlap disabled, identical prompt/sampling, and a 247-token warm response:

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

The geometry layout improved warm decode by 23.13% at the same cache-byte budget. A cold 127-token request measured 18.686 tok/s with geometry pools.

Validation already completed

• Real Laguna-S Q4_K_M load and generation with compact host banks
• CUDA JIT strided-copy correctness on SM 12.0
• Peak-memory regression proving no payload-sized CUDA staging allocation
• Synthetic mixed Q8_0/BF16 GGUF loader, cache, rebuild, prefill, decode, and dispatch coverage
• Geometry planner, independent LRU, mixed CPU/GPU stats, and non-GGUF compatibility tests
• 77 focused cache/loader/kernel/scheduler tests and 9 isolated Laguna config/module tests pass on the rebased stack

Limitations and risk

• This builds on feat(laguna): native GGUF support for poolside Laguna (S/XS) #102 and applies to direct GGUF serving; FTW conversion still does not preserve per-layer GGUF types.
• The optional cudaMemcpyBatchAsync prefill-hit test fails identically on the rebased feat(laguna): native GGUF support for poolside Laguna (S/XS) #102 baseline on this driver. Geometry mode intentionally disables that optimization because the prefill overlay can overwrite source rows.
• A 65,536-token Q8_0 configuration does not fit the tested 16 GiB card; the validated load was 32K, and the controlled cache-layout A/B was 8K.
• The strong performance claim is only the Laguna 480-vs-480 A/B above. Cross-model llama.cpp measurements are not equivalent and are intentionally excluded.

Related: #102