feat(moe): serve experts of several formats through per-format slot pools

[jason-fxz]

Part of #671. Base for the GGUF experts on the CPU executor, which follows in a separate PR.

What this does

Until now every MoE layer of a model had to use the same expert format; shared_offload_method refused a model whose layers disagree. With this PR each MoE layer keeps its own expert format and its own banks. Nothing in the mechanism is specific to a pair of formats: the layers of one checkpoint can use any of bf16, block FP8, NVFP4 and MXFP4, for example bf16 routed experts in the first and last layers and NVFP4 in the rest. The tests cover bf16 + NVFP4.

The first commit is @andy-yang-1's change. The second commit hands the CPU executor every layer's format together with its options and treats all formats alike; the GGUF follow-up needs the options.

How layers with different banks live together

A bank is one kind of per-expert array. Each format has its own set:

format	banks
bf16	gate_up, down (bf16 weights)
fp8_block	gate_up, down (FP8 weights), gate_up_scale, down_scale (bf16, one per 128x128 block)
nvfp4	gate_up, down (packed 4-bit weights), gate_up_scale, down_scale (FP8, one per 16 values), gate_up_global, down_global (fp16, one per row)
mxfp4_triton (gpt-oss)	gate_up, down (4-bit blocks), gate_up_scale, down_scale (e8m0), gate_up_bias, down_bias (bf16)
ds_fp4	gate_up, down (packed 4-bit weights), gate_up_scale, down_scale (e8m0, one per 32 values)

• Host banks. The model's bank set is the union of its formats' sets. A layer whose format has no such bank gets a zero-byte [E, 0] placeholder there. Each format loads only its own layers, then the parts are merged (merge_expert_banks), and ExpertBanks.layer_formats records which format each layer uses.

• GPU slot cache. Each bank is one GPU buffer, as before. The slot pools of feat(moe): split the offload expert cache into slot pools #689 cut every buffer into one region per pool, and here the layers are grouped by format, so each format gets its own pool with its own LRU and its own row size per bank. The same bank name can hold a different dtype and row shape in each pool (gate_up_scale is bf16 blocks for fp8_block and FP8 for nvfp4), and a pool whose format has no such bank takes zero bytes in it. With bf16 and NVFP4 layers, for example:

               | NVFP4 pool (narrow rows) | bf16 pool (wide rows) |
  gate_up      | ######################## | ##################### |
  down         | ######################## | ##################### |
  *_scale      | ####                     | (0 bytes)             |
  *_global     | #                        | (0 bytes)             |
  

  Each layer's kernel reads the views of its own pool. Staging, the fused copy plan and the hit-D2D batch skip the zero-byte banks.

• FTW. The converter writes each layer's own banks and records layer_quant_formats; the reader fills the missing banks with placeholders.

• CPU executor (hybrid decode). One format object per distinct format; each layer points to one of them. All formats must share H and I. See the next section for how the formats reach the executor.

What cannot be mixed:

• Kernels that keep GPU-resident alphas instead of banks (nvfp4_marlin, nvfp4_b12x): refused when loading the banks.
• --moe-strategy cpu, --moe-cpu-layers and resident experts: refused at startup.

How GGUF fits

GGUF already has a different quant per layer: in Qwen3.8-Flash-Next GSQ-RCO IQ3_XXS, one layer is IQ2_XS gate_up with IQ4_NL down, another IQ3_S gate_up with Q2_0 down. That case is simpler: every layer has the same two banks, gate_up and down, holding raw ggml rows, and only the row size changes with the type. #690 tags each layer with its (gate_up type, down type), so the same slot pools already give each type pair its own pool. No placeholders are needed.

What GGUF still needs is for the CPU executor to know each layer's ggml types, and the format name alone (gguf) cannot carry that. The second commit makes room for it.

Each layer's CPU format and options (second commit)

• Each layer describes its experts to the CPU executor as (name, options): MoEMethod.cpu_spec() returns (cpu_format, kernel.cpu_options(cfg)), and settings that belong to one format travel as its options. ds_fp4 now carries its fp8 activation block as {"act_block": ...} instead of a separate executor argument taken from the first layer; the GGUF follow-up will pass ("gguf", {"gate_up_type": ..., "down_type": ...}).
• The engine passes every layer's spec. The executor builds one format object per distinct (name, options), and the C++ constructor takes that list plus each layer's index into it. No format is treated as the main one, and set_layer_format from the first commit is gone.
• The checks that depend on the format (router weight on the input, the stale-extension probe for activations) look at every format. The ds_fp4 input round-trip on the GPU rewrites the whole input, so it runs only when every layer is ds_fp4. With these per-format, ds_fp4 and mxfp4_triton no longer have to be the only format in the executor.
• CpuMoeExecutor takes the formats only as layer_formats; without it every layer uses cache.quant_format with no options. Callers that build an executor directly (ft bench bw, tests) pass e.g. layer_formats=[("ds_fp4", {"act_block": 128})] * num_layers.
• The resident (fused) path now passes every layer's method when it loads the experts, so experts of several formats stop there with the refusal from the first commit. Before, it built every layer's banks in layer 0's format; with dummy banks that failed later in attach_resident_banks with a bank-set mismatch.

Review fixes (third commit)

• Every expert reader takes layers (bank layer ids) and skips the other layers' tensors before reading them; completeness is checked against layers. This covers the generic NVFP4 reader, the generic stacked bf16 reader, and the deepseek_v4, deepseek_v41, glm5_next, gpt_oss and qwen3_5_moe readers. _mixed_expert_banks hands each format its own layers, so each layer is read once.
• Every format's reader is resolved before any bank is allocated, so a missing parallel reader falls back to the serial build with nothing pinned yet.
• pin_banks takes per-layer dicts keyed by layer id, so the dummy path labels the layers of one format with their real ids.
• OffloadMoeCache refuses any resident bank spec in layer_layouts, instead of comparing role names.
• alloc_layer_banks had no callers left and is removed.

Validation

Tested on: NVIDIA H100 80GB HBM3, driver 580.95.05, CUDA 13.0; Intel Xeon Platinum 8480+ (13 physical cores visible to the process); torch 2.14.1+cu130, Python 3.14.3, gcc 11.4. No checkpoint was loaded; the tests below build their own experts.

python setup.py build_ext --inplace
pytest tests -m "not slow" --ignore=tests/e2e

• 2121 passed, 219 skipped, 1 failed at the head of this PR. The failure is test_scheme_for_agrees_with_the_stored_tensors[apodex/Apodex-1.1-mini-NVFP4[index-only]], a local HF cache entry with an unsupported modelopt_mixed config.
• The new tests/moe/test_mixed_expert_formats.py builds one model with bf16 and NVFP4 layers; its 5 tests run here and pass: the engine binds the most common format and checks every layer's banks; dummy banks merge the roles with placeholders; decode and overlap prefill through the per-format pools match a dequantized reference; the CPU executor computes each layer in its own format, also in hybrid decode; the FTW round trip. A sixth test checks that _mixed_expert_banks gives each format's reader its own layers and resolves every reader before reading any.
• The reader tests in tests/models/test_qwen3_5_moe_weight.py (generic NVFP4 reader on five checkpoint layouts, the qwen3_5_moe block-FP8 reader serial and parallel) and tests/models/test_models_loader.py (stacked bf16) also read one layer through layers.

Not covered

• Mixes other than bf16 + NVFP4: the code paths are the same for every format above, but only bf16 + NVFP4 is tested.
• End to end on a checkpoint whose MoE layers use different expert formats. The readers are tested on synthetic checkpoints only.
• bf16 experts stored one tensor per expert next to quantized layers: the generic bf16 reader reads only the stacked experts.gate_up_proj / experts.down_proj form, and qwen3_5_moe hands its stacked bf16 experts to it only when the model-wide expert_quant is "none".
• Resident (fused) experts of several formats: refused at startup. On unified-memory GPUs, where auto picks fused, such a model needs --moe-strategy offload or hybrid.
• Slot pools for several formats, left to a follow-up pool policy: the prefill staging area takes E rows of the widest format of each bank (twice that with prefill overlap) from the front of the pools and the slot split does not account for it, so a small pool can lose its slots to staging on every prefill, or the cache can fail to stage a layer at startup; a small pool lowers max_running_req for the whole server; the slot limit is the minimum over all kernels (no kernel with a limit can be mixed today); zero-byte banks still launch empty copy blocks.
• The automatic hybrid fetch cap (--moe-hybrid-max-fetch auto) reads the ft bench bw entry of the most common format for every layer.

[andy-yang-1]

The direction looks right to me. Taking the union of each format's banks, with zero-byte placeholders on top of the #689 slot pools, keeps the cache changes small, and the per-layer (name, options) CPU spec is clean (the second commit could land on its own). I found no compute bug in the cache, CPU executor or FTW paths. But real mixed checkpoints cannot load yet, and the default pool layout has some sharp edges.

Tested on H100, torch 2.9.0+cu128: the 5 new tests pass. The full -m "not slow" run (without e2e) gives 2080 passed, 43 failed, 2 errors; all 45 also fail on the merge base b6193e86 (local vLLM marlin signature, import errors), so there are no regressions.

Should be fixed or guarded before merge

1. Loading a raw checkpoint and ft checkpoint fail on a real mixed checkpoint (from reading the code; I had no mixed checkpoint at hand). _mixed_expert_banks calls iter_expert_pieces(kind) once per format and assumes the reader yields only that format's layers. No reader does:

• The generic NVFP4 reader requires exactly num_moe_layers * num_experts * 9 tensors. A ModelOpt checkpoint with bf16 first and last layers stops with found N expert tensors, expected M.
• The bf16 reader (stacked_expert_pieces) walks every expert tensor and raises Unexpected expert weight key on the NVFP4 layers' scales.
• The model hooks (qwen3_5_moe, glm5_next, deepseek_v4, gpt_oss) do not filter by layer either.

The new test passes only because it uses a hand-written _reader. Even with filtering, each format pass reads every expert of the checkpoint from disk. Suggest an explicit iter_expert_pieces(..., layers=) contract: filter keys before reading and count against len(layers).

2. The parallel-to-serial fallback leaks pinned banks. The fallback in load_expert_banks assumes the reader raises NotImplementedError before any bank is allocated. The mixed path builds and pins the largest format first and only then resolves the next reader. If that reader has no parallel version, the first format's banks leak (host banks live for the process) and pinned RAM doubles. Resolve every reader before building any bank.

3. Unified-memory GPUs refuse mixed models with default flags. On GB10 or Jetson, auto picks fused, and nvfp4 is in _RESIDENT_EXPERT_QUANTS, so a bf16 + NVFP4 model stops with "served from the offload slot cache only". Falling back to offload there runs into #369. Resident is the easiest place to support mixed formats: each layer already owns its _expert_banks, and build_expert_banks already takes resident and layers. At least keep auto from picking fused for these models.

Performance (a follow-up is fine, but worth stating in the PR)

4. Shared prefill staging and zero-byte banks. _stage_bytes[k] is E times the widest row of any format. In the NVFP4-only banks (*_scale, *_global), only the NVFP4 pool has space for staging, and split_slots does not take that into account. Measured with E=32, 12 layers, top_k=2:

layers	slots	result
all NVFP4	153	prefill overlap on, 89 slots outside staging
bf16 first/last, NVFP4 rest	153	overlap turned off with only a warning; 0 slots outside staging in both pools, so every prefill wipes the bf16 layers' cache
1 NVFP4 layer of 12	128 or 256	startup fails with ValueError: cannot stage one whole layer of experts

Options: a dedicated staging region outside the pools for mixed models (costs 2 x E x the widest row per bank), or a staging-aware slot split that fails with the minimum cache size when it cannot fit.

5. A small bf16 pool caps the decode batch of the whole server. Pools start from _moe_pools(1), which guarantees only top_k slots each, and _clamp_decode_batch_to_pools lowers max_running_req to fit the smallest pool. With E=128, top_k=8 and 60 layers (bf16 first and last), a cache of 10%, 20% or 30% of all experts gives max_running_req 3, 6 or 9. A single-format model with the same slot count is not capped. Mixed formats now always produce several pools, so this happens by default.

Nits

• _slot_limit applies the minimum over all kernels to every pool. Since marlin and b12x are refused, it is always None today.
• CpuMoeExecutor silently ignores fmt when layer_formats is also passed (the new test passes both). The new opts parameter of _resolve_banks is unused.
• In the dummy path, pin_banks({role: [per_layer[l] for l in ids]}) enumerates list positions, not layer ids. If a residency plan is ever active, the layers get the wrong labels; today this is masked by the refusal of mixed formats with residency.
• __post_init__ uses the tuple of role names to detect several formats, but fp8_block and ds_fp4 have the same role names. Refusing any resident spec in layer_layouts would be simpler.
• alloc_layer_banks has no callers left. The fused copy still launches kBlocksPerBank empty blocks for each zero-byte bank.

[jason-fxz]

Thanks for the careful review. b218e5b addresses the loader issues and most nits; the pool layout items are left to a follow-up pool policy, and the PR description now lists them under Not covered.

1. Readers. Fixed. iter_expert_pieces(..., layers=) now limits a reader to the given bank layers: it skips the other layers' tensors before reading them and checks completeness against layers. This covers the generic NVFP4 reader, the generic stacked bf16 reader, and the deepseek_v4, deepseek_v41, glm5_next, gpt_oss and qwen3_5_moe readers, and _mixed_expert_banks passes each format its own layers, so each layer is read once. The NVFP4 and block-FP8 reader tests and the stacked bf16 test now also read a single layer this way. Two gaps remain and are listed in the description: the generic bf16 reader only reads the stacked experts.gate_up_proj / experts.down_proj form, so bf16 experts stored one tensor per expert next to quantized layers are not readable yet, and qwen3_5_moe hands its stacked bf16 experts over only when the model-wide expert_quant is "none". I have no mixed checkpoint to check the real layout against.

2. Fallback leak. Fixed. _mixed_expert_banks resolves every format's reader before it builds any bank; a new test checks that all readers are resolved before the first one is read.

3. Unified-memory GPUs. Not in this PR. Resident experts of several formats stay refused, now with the clear error on the fused path as well (_load_resident_experts passes the layer methods). Support for them on the resident path, and keeping auto away from fused for these models, will come separately.

4, 5. Staging and small pools. Agreed on both. The cache already lays the pools out in whatever order the planner hands it, so these belong to a pool policy (pool order, a minimum slot count per pool that leaves room for staging and the decode batch, filling small pools when they fit) in a follow-up PR.

Nits.

• _slot_limit: left as is for now; with marlin and b12x refused it is always None. The policy PR will make it per pool.
• fmt vs layer_formats: fmt is gone; CpuMoeExecutor takes only layer_formats, and without it every layer uses cache.quant_format. The opts argument of _resolve_banks is for the GGUF follow-up, which passes each layer's ggml types as options.
• pin_banks: now takes per-layer dicts keyed by layer id, and the dummy path passes them.
• __post_init__: now refuses any resident spec in layer_layouts.
• alloc_layer_banks: removed.
• Empty copy blocks for zero-byte banks: left for later.