Cached MoE Prefill

Cached MoE prefill

Cached prefill runs routed expert matrix work on CUDA while expert weights are managed by the cache. Prefill route density and matrix geometry differ from one-token decode, so it retains the existing shape-selected MMQ or MMVQ consumer instead of forcing one kernel family.

Important

CUDA REQUIRED: The cached expert buffers, mapped dispatch, MMQ/MMVQ consumers, and staging pipeline are CUDA-only.

Execution behavior

Routing IDs are reused across sibling expert projections where valid. Large unique sets use bounded overflow staging. Partial tiles and quantized source padding are initialized and bounded before MMQ reads them. Compact expert shapes remain eligible for MMVQ when that is the normal faster consumer; the cache does not force MMQ merely because prefill is active.

Eligible small contiguous F32 expert biases can remain resident through the auxiliary-bank policy. Quantized block-scale banks remain tied to the expert slot that owns their weight data.

Prefill can touch most experts in a layer, so cache hit rate alone is not a sufficient success metric. Compare prompt tokens per second, transfer bytes, staging wait, selected kernel family, and exact/coherent output at the real ubatch.

Branch implementations

Branch	Implementation
moe-cache	mmid.cu, mmq.cu, mmvq.cu, moe-cache.cu

Review history: PR 41, PR 52, and the promoted grouped architecture in PR 56.

Feature index