[MoE] Expert pool follow-ups: MTP, Marlin 32 GiB bank guard, prefill staging

docs/features/moe_expert_pool.md

@@ -36,20 +36,36 @@ allocations. Rows can move between layers at run time.
   with logical expert alignment and a physical-row remap. No host-side
   routing readback or cache planning sits between routing and the GEMM, so
   the MoE layers stay inside the captured CUDA graph (`FULL_DECODE_ONLY`).
-- Wider batches (prefill): resident experts are read from the bank and the
-  rest directly from the pinned host source through its accelerator view.
+- Wider batches (prefill): resident experts are read from the bank. Marlin
+  reads its weights once per M-block, so a non-resident expert read through
+  the host view would cross PCIe several times per step. Prefill steps of
+  at least 512 tokens therefore copy each layer's non-resident rows into a
+  device staging buffer by DMA, one layer ahead of the compute, and read
+  them from there; smaller wide batches read their routed rows through the
+  host view. `--no-moe-expert-pool-prefill-staging` turns staging off and
+  saves its two full-layer buffers of GPU memory.
 
 ## Requirements and limits
 
 - ModelOpt NVFP4 checkpoints with the Marlin MoE backend
-  (`--moe-backend marlin`). Other quantization methods or NVFP4 backends
-  are rejected when the layer is built.
+  (`--moe-backend marlin`). NVFP4 layers on other backends are rejected
+  when the layer is built; MoE layers with another quantization method
+  (e.g. an MTP head the checkpoint leaves in BF16) are not pooled and keep
+  their experts on the GPU.
+- Speculative decoding: the decode lanes are sized for
+  `1 + num_speculative_tokens` tokens per request. A BF16 draft MoE needs
+  its own backend, e.g. `--speculative-config '{"method": "mtp",
+  "num_speculative_tokens": 3, "moe_backend": "triton"}'`.
 - No expert, data or sequence parallelism.
 - `N` must be at least `top_k` and every MoE layer must have the same expert
   count and top-k.
-- Prefill of long inputs reads the non-resident experts from host memory on
-  every chunk, which adds host-read transfer cost per chunk; the dominant
-  term of long-input prefill time has not been profiled.
+- Each bank weight tensor must fit in 32 GiB, the range Marlin MoE indexes
+  with 32-bit offsets (e.g. 8192 rows of 4 MiB gate/up weights); a larger
+  bank is rejected at installation with the largest `N` that fits.
+- Every prefill step still moves each layer's non-resident experts over
+  PCIe once, so long inputs prefill faster with larger
+  `--max-num-batched-tokens`. Staging reads the placement back to the host
+  once per prefill step.
 
 ## Example
 

tests/config/test_moe_expert_pool_rows.py

@@ -40,3 +40,11 @@ def test_cli_flag_and_field():
         ).moe_expert_pool_rows
         == 16
     )
+
+
+def test_prefill_staging_is_on_by_default_and_can_be_disabled():
+    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
+    assert parser.parse_args([]).moe_expert_pool_prefill_staging
+    args = parser.parse_args(["--no-moe-expert-pool-prefill-staging"])
+    assert not args.moe_expert_pool_prefill_staging
+    assert OffloadConfig().moe_expert_pool_prefill_staging

tests/kernels/expert_pool/test_install_guards.py

@@ -60,6 +60,17 @@ def test_planner_width_is_bounded_by_the_decode_lane_cap():
         assert inst._next_power_of_two(top_k * tokens) <= 2 * inst.MAX_DECODE_LANES
 
 
+def test_bank_beyond_the_marlin_weight_index_is_rejected():
+    # 4 MiB w13 rows (GLM-5.3-Flash, TP2): 8192 bank rows fill the 32 GiB.
+    source = {
+        "w13_weight": torch.empty(1, 4 << 20, dtype=torch.uint8, device="meta"),
+        "w2_weight": torch.empty(1, 2 << 20, dtype=torch.uint8, device="meta"),
+    }
+    inst._check_bank_size(source, slots=193, num_layers=42, staging=64)
+    with pytest.raises(ValueError, match="w13_weight .* <= 193"):
+        inst._check_bank_size(source, slots=194, num_layers=42, staging=64)
+
+
 def test_top_k_beyond_the_lane_cap_is_rejected():
     assert inst.install_expert_pool(torch.nn.Module(), torch.device("cpu")) is None
     for bad in (0, inst.MAX_DECODE_LANES + 1):
@@ -100,11 +111,22 @@ def test_layer_guard_accepts_the_supported_geometry():
         (_routed(dp_size=2), "data parallelism"),
         (_routed(sp=True), "sequence parallelism"),
         (_routed(method=_method(NvFp4MoeBackend.VLLM_CUTLASS)), "Marlin"),
-        (_routed(method=SimpleNamespace(nvfp4_backend=None)), "Marlin"),
     ],
 )
 def test_layer_guard_rejects_before_any_weight_is_allocated(bad, match):
     from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
 
     with pytest.raises(ValueError, match=match):
         RoutedExperts._validate_expert_pool_supported(bad)
+
+
+def test_layer_guard_keeps_unquantized_layers_on_the_gpu():
+    # e.g. the BF16 MTP head of an NVFP4 checkpoint: served outside the pool.
+    from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
+    from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
+        UnquantizedFusedMoEMethod,
+    )
+
+    layer = _routed(method=UnquantizedFusedMoEMethod.__new__(UnquantizedFusedMoEMethod))
+    RoutedExperts._validate_expert_pool_supported(layer)
+    assert layer._moe_expert_pool_rows == 0

tests/kernels/expert_pool/test_pool_marlin_cuda.py

@@ -16,6 +16,7 @@
     TOP_K,
     E,
     K,
+    M,
     dist_env,  # noqa: F401
     make_layer,
     quantized_weights,
@@ -27,6 +28,7 @@
 )
 from vllm.model_executor.layers.fused_moe.expert_pool.pool import verify_bank_rows
 from vllm.model_executor.layers.fused_moe.expert_pool.tables import (
+    TENSORS,
     check_global_tables,
     resident_per_layer,
     set_gate,
@@ -159,3 +161,64 @@ def run(i, x, logits, n):
         run(0, x, _decode(order, device), 1)
     run(1, x, _decode([0, 1], device), 1)
     assert pls[0].decode_steps == 9 and pls[1].decode_steps == 4  # 1 gate-closed
+
+
+def test_prefill_staging_matches_the_host_view(dist_env):  # noqa: F811
+    """Two layers staged at M tokens: each layer reads its non-resident rows
+    from a staging buffer filled while the previous layer ran; the outputs
+    equal the host-view read byte for byte, and a smaller batch is not
+    staged."""
+    device = torch.accelerator.current_accelerator()
+    ref_cfgs, pool_cfgs, refs, layers = [], [], [], []
+    for seed_offset in (0, 1):
+        params = quantized_weights(device, seed_offset=seed_offset)
+        ref_cfgs.append(vllm_config(0))
+        pool_cfgs.append(vllm_config(SLOTS))
+        refs.append(make_layer(ref_cfgs[-1], params))
+        layers.append(make_layer(pool_cfgs[-1], params, host_source=True))
+    for layer in layers:
+        for name in ("w13_weight_scale_2", "w2_weight_scale_2"):
+            p = getattr(layer.routed_experts, name)
+            p.data = p.data.to(device)
+    model = torch.nn.ModuleDict({"a": layers[0], "b": layers[1]})
+    pool = install_expert_pool(model, device, max_decode_tokens=1, stage_min_tokens=M)
+    assert pool is not None
+    pls = [layer.routed_experts.expert_pool_layer for layer in layers]
+    stage = pls[0].stage
+    assert stage is not None and stage is pls[1].stage and stage.layers == pls
+    from vllm.forward_context import set_forward_context
+
+    def forward(i, x, logits):
+        with set_forward_context(None, pool_cfgs[i], num_tokens=x.shape[0]):
+            out = layers[i](x, logits)
+        torch.accelerator.synchronize(device)
+        return out
+
+    # Rows 0..3 of each layer are resident. Rows 4..7 are staged whether
+    # routed (4, 5) or not (6, 7); resident rows are not copied.
+    logits = routing([0, 4, 1, 5, 2, 3, 0, 1], device)
+    x = torch.randn(M, K, dtype=torch.bfloat16, device=device)
+    for buffer in stage.buffers:
+        for t in buffer.values():
+            t.zero_()
+    staged = [forward(i, x, logits) for i in (0, 1)]
+    assert stage.staged == 2 and stage.copied_rows == 8
+    for i, pl in enumerate(pls):
+        for name in TENSORS:
+            buf = stage.buffers[i][name].cpu()
+            assert torch.equal(buf[SLOTS:], pl.sources[name][SLOTS:]), (i, name)
+            assert not buf[:SLOTS].reshape(-1).view(torch.uint8).any(), (i, name)
+    for i, pl in enumerate(pls):
+        pl.stage = None
+        assert torch.equal(staged[i], forward(i, x, logits))
+        pl.stage = stage
+        with set_forward_context(None, ref_cfgs[i], num_tokens=M):
+            want = refs[i](x, logits)
+        torch.testing.assert_close(staged[i], want, rtol=2e-2, atol=2e-2)
+    # Below the threshold (still wider than the decode lanes) the partition
+    # step reads the host view.
+    assert pls[0].width < TOP_K * (M - 1)
+    forward(0, x[: M - 1].contiguous(), logits[: M - 1].contiguous())
+    assert stage.staged == 2 and pls[0].partition_steps == 3
+    check_global_tables(pool.tables)
+    assert int(pool.tables.error[0]) == 0

vllm/config/offload.py

@@ -103,6 +103,12 @@ class OffloadConfig:
     0 (default) keeps every expert on the GPU. Supported for the ModelOpt
     NVFP4 Marlin MoE backend without expert/data parallelism."""
 
+    moe_expert_pool_prefill_staging: bool = True
+    """With the expert pool, large prefill steps copy each layer's
+    non-resident expert rows into a device buffer (prefetched one layer
+    ahead) instead of letting the MoE kernel re-read them from host memory.
+    Costs two layers' worth of expert weights in GPU memory."""
+
     @model_validator(mode="after")
     def validate_offload_config(self) -> "OffloadConfig":
         """Validate offload configuration constraints."""

vllm/engine/arg_utils.py

@@ -552,6 +552,9 @@ class EngineArgs:
     offload_prefetch_step: int = PrefetchOffloadConfig.offload_prefetch_step
     offload_params: set[str] = get_field(PrefetchOffloadConfig, "offload_params")
     moe_expert_pool_rows: int = OffloadConfig.moe_expert_pool_rows
+    moe_expert_pool_prefill_staging: bool = (
+        OffloadConfig.moe_expert_pool_prefill_staging
+    )
     gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization
     kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes
     max_num_batched_tokens: int | None = None
@@ -1360,6 +1363,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         offload_group.add_argument(
             "--moe-expert-pool-rows", **offload_kwargs["moe_expert_pool_rows"]
         )
+        offload_group.add_argument(
+            "--moe-expert-pool-prefill-staging",
+            **offload_kwargs["moe_expert_pool_prefill_staging"],
+        )
 
         # Multimodal related configs
         multimodal_kwargs = get_kwargs(MultiModalConfig)
@@ -2629,6 +2636,7 @@ def create_engine_config(
                 offload_params=self.offload_params,
             ),
             moe_expert_pool_rows=self.moe_expert_pool_rows,
+            moe_expert_pool_prefill_staging=self.moe_expert_pool_prefill_staging,
         )
 
         if self.gdn_prefill_backend is not None:

vllm/model_executor/layers/fused_moe/expert_pool/install.py

@@ -20,6 +20,7 @@
 from vllm.model_executor.layers.fused_moe.expert_pool.copy import configure_copy
 from vllm.model_executor.layers.fused_moe.expert_pool.layer import PoolLayer
 from vllm.model_executor.layers.fused_moe.expert_pool.pool import GlobalPool
+from vllm.model_executor.layers.fused_moe.expert_pool.stage import PrefillStage
 from vllm.model_executor.layers.fused_moe.expert_pool.tables import (
     TENSORS,
     allocate_step_buffers,
@@ -34,6 +35,13 @@
 # single-program planner loops over WIDTH lanes and scans the pool per miss.
 MAX_DECODE_LANES = 64
 
+# Marlin MoE indexes its weight operand in 16-byte units with a 32-bit int.
+MARLIN_MAX_WEIGHT_BYTES = 16 * 2**31
+
+# Partition steps with at least this many tokens stage their non-resident
+# rows (see stage.py); smaller ones read only their routed rows from the host.
+PREFILL_STAGE_MIN_TOKENS = 512
+
 
 def _next_power_of_two(value: int) -> int:
     return 1 << max(int(value) - 1, 0).bit_length()
@@ -46,6 +54,19 @@ def _check_top_k(top_k: int) -> None:
         )
 
 
+def _check_bank_size(
+    source: dict[str, torch.Tensor], slots: int, num_layers: int, staging: int
+) -> None:
+    for name in ("w13_weight", "w2_weight"):
+        row_bytes = source[name][0].nbytes
+        if (slots * num_layers + staging) * row_bytes > MARLIN_MAX_WEIGHT_BYTES:
+            max_slots = (MARLIN_MAX_WEIGHT_BYTES // row_bytes - staging) // num_layers
+            raise ValueError(
+                f"expert pool bank: {name} would exceed the 32 GiB that Marlin "
+                f"MoE can address; use --moe-expert-pool-rows <= {max_slots}"
+            )
+
+
 def check_pool_layers(layers: list[tuple[str, torch.nn.Module]]) -> None:
     """Reject geometries the pool cannot serve, before anything is allocated:
     every layer must share expert count, top-k, resident rows and backend,
@@ -115,9 +136,14 @@ def _sources(name: str, layer: torch.nn.Module) -> dict[str, torch.Tensor]:
 
 
 def install_expert_pool(
-    model: torch.nn.Module, device: torch.device, max_decode_tokens: int = 1
+    model: torch.nn.Module,
+    device: torch.device,
+    max_decode_tokens: int = 1,
+    stage_min_tokens: int = 0,
 ) -> GlobalPool | None:
-    """Allocate the shared bank and bind every pending pool layer to it."""
+    """Allocate the shared bank and bind every pending pool layer to it.
+    With stage_min_tokens > 0, partition steps of at least that many tokens
+    read their non-resident rows from prefill staging buffers."""
     from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
         make_nvfp4_moe_kernel,
     )
@@ -154,8 +180,16 @@ def install_expert_pool(
                 raise RuntimeError(
                     f"{name}.{tensor_name}: layer rows differ from {first_name}"
                 )
+    _check_bank_size(sources[0], slots, len(layers), staging)
     pool = GlobalPool(device, sources[0], [slots] * len(layers), staging)
     configure_copy("chunks")
+    # Allocated here, before memory profiling, so the KV cache is sized
+    # around the staging buffers.
+    stage = (
+        PrefillStage(device, sources[0], num_experts, stage_min_tokens)
+        if stage_min_tokens > 0
+        else None
+    )
     for index, ((name, layer), src) in enumerate(zip(layers, sources)):
         method = layer.quant_method
         start = pool.offset(index)
@@ -201,6 +235,7 @@ def install_expert_pool(
             top_k=top_k,
             activation=layer.activation,
             apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            stage=stage,
         )
         layer.expert_pool_pending = False
     # Placement policy (the lab run's values): promotions on every forward,
@@ -231,6 +266,14 @@ def install_expert_pool(
         (pool.pool_bytes + pool.staging_bytes) / 2**30,
         type(layers[0][1].quant_method).__name__,
     )
+    if stage is not None:
+        stage.layers = [layer.expert_pool_layer for _, layer in layers]
+        logger.info(
+            "Expert pool prefill staging: partition steps with >= %d tokens "
+            "read non-resident rows from two prefetched device buffers (%.2f GiB)",
+            stage_min_tokens,
+            stage.bytes / 2**30,
+        )
     return pool
 
 

vllm/model_executor/layers/fused_moe/expert_pool/layer.py

@@ -21,6 +21,7 @@
 import torch
 
 from vllm.model_executor.layers.fused_moe.expert_pool.pool import GlobalPool, copy_in
+from vllm.model_executor.layers.fused_moe.expert_pool.stage import PrefillStage
 from vllm.model_executor.layers.fused_moe.expert_pool.tables import (
     TENSORS,
     StepBuffers,
@@ -86,6 +87,7 @@ def __init__(
         top_k: int,
         activation: Any,
         apply_router_weight_on_input: bool,
+        stage: PrefillStage | None = None,
     ) -> None:
         self.index = index
         self.pool = pool
@@ -106,6 +108,7 @@ def __init__(
         self.width = buffers.gather_src.shape[0]
         self.hot_map = pool.tables.layer_slice(pool.tables.hot_phys, index)
         self.cold_map = pool.tables.layer_slice(pool.tables.cold_phys, index)
+        self.stage = stage  # shared by every layer, or None
         self.decode_steps = 0
         self.partition_steps = 0
 
@@ -124,15 +127,19 @@ def apply(
                 x, weights, ids, ((self.bank, self.buffers.step_map, self.bank_rows),)
             )
         # Wider batches (prefill): resident rows from the bank, the rest read
-        # straight from the pinned host source through its accelerator view.
+        # straight from the pinned host source through its accelerator view,
+        # or, for large uncaptured steps, from the prefill staging buffer.
         self.partition_steps += 1
+        cold = self.host
+        if self.stage is not None and self.stage.applies(x):
+            cold = self.stage.acquire(self)
         return self._run_marlin(
             x,
             weights,
             ids,
             (
                 (self.bank, self.hot_map, self.bank_rows),
-                (self.host, self.cold_map, self.num_experts),
+                (cold, self.cold_map, self.num_experts),
             ),
         )