[Draft][MoE][Offload] Row-level expert weight cache provider (moe_expert_cache_provider=row)

tests/config/test_moe_expert_cache_provider.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""moe_expert_cache_provider: config field, CLI -> EngineArgs, hash."""
+
+import pytest
+from pydantic import ValidationError
+
+from vllm.config.offload import OffloadConfig
+from vllm.engine.arg_utils import EngineArgs
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+
+def test_default_is_the_upstream_cached_provider():
+    assert OffloadConfig().moe_expert_cache_provider == "cached"
+
+
+def test_hash_distinguishes_the_provider():
+    cached = OffloadConfig(moe_expert_cache_size=8).compute_hash()
+    row = OffloadConfig(
+        moe_expert_cache_size=8, moe_expert_cache_provider="row"
+    ).compute_hash()
+    assert cached != row
+
+
+def test_invalid_provider_is_rejected():
+    with pytest.raises(ValidationError):
+        OffloadConfig(moe_expert_cache_provider="lru")  # type: ignore[arg-type]
+
+
+def test_cli_reaches_engine_args_and_offload_config():
+    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
+    args = parser.parse_args(
+        ["--moe-expert-cache-size", "16", "--moe-expert-cache-provider", "row"]
+    )
+    engine_args = EngineArgs.from_cli_args(args)
+    assert engine_args.moe_expert_cache_provider == "row"
+    # The same kwargs EngineArgs.create_engine_config passes on.
+    offload = OffloadConfig(
+        moe_expert_cache_size=engine_args.moe_expert_cache_size,
+        moe_expert_cache_split=engine_args.moe_expert_cache_split,
+        moe_expert_cache_provider=engine_args.moe_expert_cache_provider,
+    )
+    assert (offload.moe_expert_cache_size, offload.moe_expert_cache_provider) == (
+        16,
+        "row",
+    )

tests/kernels/moe/test_expert_lru_cache.py

@@ -261,11 +261,54 @@ def test_scale_lifecycle():
     torch.testing.assert_close(provider.buf_w2_scale[slot_7].cpu(), w2_s[7])
 
 
+def test_scale_2_lifecycle():
+    """Per-expert global scales follow the slot through fill, eviction, reuse."""
+    provider, _, _, _ = _make_provider(capacity=2, with_scales=True)
+    assert provider.buf_w13_scale_2 is None and provider.buf_w2_scale_2 is None
+    set_random_seed(42)
+    w13, w2 = _make_weights(8, torch.bfloat16)
+    w13_s, w2_s = _make_scales(8)
+    g13 = torch.rand(8, 2, dtype=torch.float32)  # gate/up globals per expert
+    g2 = torch.rand(8, 1, dtype=torch.float32)
+    provider = CachedWeightProvider(
+        2, w13, w2, w13_s, w2_s, w13_scale_2=g13, w2_scale_2=g2
+    )
+    assert provider.buf_w13_scale_2 is not None
+    assert provider.buf_w2_scale_2 is not None
+    assert provider.buf_w13_scale_2.shape == (2, 2)
+    assert provider.buf_w2_scale_2.shape == (2, 1)
+    result = provider.prepare(_topk([0, 1]))
+    for eid in (0, 1):
+        slot = int(result.expert_map[eid])
+        torch.testing.assert_close(provider.buf_w13_scale_2[slot].cpu(), g13[eid])
+        torch.testing.assert_close(provider.buf_w2_scale_2[slot].cpu(), g2[eid])
+        torch.testing.assert_close(result.w1_scale[slot].cpu(), w13_s[eid])
+    slot_0 = int(result.expert_map[0])
+    result = provider.prepare(_topk([1, 5]))  # evicts 0; 5 reuses its slot
+    assert int(result.expert_map[5]) == slot_0
+    torch.testing.assert_close(provider.buf_w13_scale_2[slot_0].cpu(), g13[5])
+    torch.testing.assert_close(provider.buf_w2_scale_2[slot_0].cpu(), g2[5])
+    assert result.w1_scale is provider.buf_w13_scale  # result surface unchanged
+    assert not hasattr(result, "w1_scale_2")
+
+
+def test_scale_2_validation():
+    w13, w2 = _make_weights(8, torch.bfloat16)
+    with pytest.raises(ValueError, match="together"):
+        CachedWeightProvider(2, w13, w2, w13_scale_2=torch.rand(8, 2))
+    with pytest.raises(ValueError, match="rows"):
+        CachedWeightProvider(
+            2, w13, w2, w13_scale_2=torch.rand(7, 2), w2_scale_2=torch.rand(8, 1)
+        )
+
+
 def test_no_scales_when_not_provided():
     """Without scale inputs, scale buffers remain None."""
     provider, *_ = _make_provider()
     assert provider.buf_w13_scale is None
     assert provider.buf_w2_scale is None
+    assert provider.buf_w13_scale_2 is None
+    assert provider.buf_w2_scale_2 is None
     result = provider.prepare(_topk([0]))
     assert result.w1_scale is None
     assert result.w2_scale is None

tests/kernels/moe/test_moe_forward_breakable.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The MoE custom op is a breakable-CUDA-graph break point only for layers
+that hold an expert cache provider; uncached layers stay in the segment."""
+
+from types import SimpleNamespace
+from unittest import mock
+
+import torch
+
+from vllm.compilation import breakable_cudagraph as bcg
+from vllm.model_executor.layers.fused_moe.runner import moe_runner
+from vllm.model_executor.layers.fused_moe.runner.moe_runner_interface import (
+    MoERunnerInterface,
+)
+
+
+def _fake_runner(tag: str, cached: bool) -> MoERunnerInterface:
+    """A MoERunnerInterface instance for get_layer_from_name's isinstance
+    check: every abstract member is stubbed (none is called here) and
+    __init__ is bypassed."""
+    stubs = {
+        name: (lambda *a, **k: None) for name in MoERunnerInterface.__abstractmethods__
+    }
+    cls: type = type("_FakeRunner", (MoERunnerInterface,), stubs)
+    runner: MoERunnerInterface = object.__new__(cls)
+    runner.tag = tag
+    runner.routed_experts = SimpleNamespace(
+        expert_weight_provider=object() if cached else None
+    )
+    return runner
+
+
+def _run(monkeypatch, cached: bool):
+    monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "1")
+    calls: list[str] = []
+
+    def fn(hidden_states, router_logits, shared_experts_input, input_ids, name, pad):
+        calls.append("fn")
+        return hidden_states
+
+    class Capture:
+        _capturing = True
+
+        def add_eager(self, thunk):
+            calls.append("eager")
+            return thunk()
+
+    layer = SimpleNamespace(
+        routed_experts=SimpleNamespace(
+            expert_weight_provider=object() if cached else None
+        )
+    )
+    with (
+        mock.patch.object(moe_runner, "get_layer_from_name", lambda n: layer),
+        mock.patch.object(
+            bcg.BreakableCUDAGraphCapture, "current", classmethod(lambda cls: Capture())
+        ),
+        mock.patch.object(bcg, "is_forward_context_available", lambda: False),
+        # weak_ref_tensor is CUDA-only; tensor lifetime is a GPU acceptance item.
+        mock.patch.object(bcg, "weak_ref_tensor", lambda t: t),
+    ):
+        wrapped = moe_runner._eager_break_when_cached(fn)
+        assert wrapped is not fn
+        x = torch.zeros(2, 4)
+        out = wrapped(x, x, None, None, "layer", 0)
+    assert out is x
+    return calls
+
+
+def test_cached_layer_breaks_the_capture(monkeypatch):
+    assert _run(monkeypatch, cached=True) == ["eager", "fn"]
+
+
+def test_uncached_layer_stays_in_the_segment(monkeypatch):
+    assert _run(monkeypatch, cached=False) == ["fn"]
+
+
+def test_legacy_placeholder_consumes_one_layer_per_op(monkeypatch):
+    """With the legacy "from_forward_context" name the lookup is stateful;
+    the wrapper must resolve it once per op, for cached and uncached layers
+    alike (the shared variant uses the same wrapper)."""
+    monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "1")
+    seen: list[str] = []
+
+    def fn(hidden_states, router_logits, shared_experts_input, input_ids, name, pad):
+        seen.append(moe_runner.get_layer_from_name(name).tag)  # plain fetch
+        return hidden_states
+
+    class Capture:
+        _capturing = True
+        thunks: list = []
+
+        def add_eager(self, thunk):
+            Capture.thunks.append(thunk)  # kept, as the real capture does
+            return thunk()
+
+    Capture.thunks = []
+    layers = {"a": _fake_runner("a", cached=True), "b": _fake_runner("b", cached=False)}
+    ctx = SimpleNamespace(
+        all_moe_layers=["a", "b"], moe_layer_index=0, no_compile_layers=layers
+    )
+    with (
+        mock.patch.object(moe_runner, "_USE_LAYERNAME", False),
+        mock.patch.object(moe_runner, "get_forward_context", lambda: ctx),
+        mock.patch.object(moe_runner, "_resolve_layer_name", lambda n: n),
+        mock.patch.object(
+            bcg.BreakableCUDAGraphCapture, "current", classmethod(lambda cls: Capture())
+        ),
+        mock.patch.object(bcg, "is_forward_context_available", lambda: False),
+        # weak_ref_tensor is CUDA-only; tensor lifetime is a GPU acceptance item.
+        mock.patch.object(bcg, "weak_ref_tensor", lambda t: t),
+    ):
+        wrapped = moe_runner._eager_break_when_cached(fn)
+        x = torch.zeros(2, 4)
+        wrapped(x, x, None, None, "from_forward_context", 0)
+        wrapped(x, x, None, None, "from_forward_context", 0)
+        assert seen == ["a", "b"]
+        assert ctx.moe_layer_index == 2
+        # Replay of the recorded (cached) thunk fetches the fixed concrete
+        # name again and does not touch the forward-context index.
+        assert len(Capture.thunks) == 1
+        Capture.thunks[0]()
+    assert seen == ["a", "b", "a"]
+    assert ctx.moe_layer_index == 2
+
+
+def test_identity_when_breakable_graphs_are_off(monkeypatch):
+    monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "0")
+
+    def fn(*a):
+        return a
+
+    assert moe_runner._eager_break_when_cached(fn) is fn