Merge branch 'main' of https://gh-proxy.org/https://github.com/flagos-ai/vllm-plugin-FL

Author: luoyc123

README.md

@@ -55,7 +55,7 @@ In theory, vllm-plugin-FL can support all models available in vLLM, as long as n
     # or editble install
     pip install --no-build-isolation -e .
     ```
- 
+
     For CUDA-like devices, including CUDA and HIP/ROCm environments that use
     PyTorch's CUDA dispatch key, build the plugin native extension by setting
     `VLLM_VENDOR=cuda` during installation:

tests/unit_tests/ops/test_activation.py

@@ -14,24 +14,24 @@ class TestSiluAndMulFL:
     """Test SiluAndMulFL class behavior."""
 
     @pytest.fixture
-    def mock_call_op(self):
-        with patch("vllm_fl.ops.activation.call_op") as mock:
+    def mock_cached_op(self):
+        with patch("vllm_fl.ops.activation._silu_and_mul") as mock:
             yield mock
 
     @pytest.fixture
     def mock_parent_init(self):
         with patch("vllm_fl.ops.activation.SiluAndMul.__init__", return_value=None):
             yield
 
-    def test_forward_oot_dispatches_correctly(self, mock_parent_init, mock_call_op):
+    def test_forward_oot_dispatches_correctly(self, mock_parent_init, mock_cached_op):
         """Test forward_oot calls dispatch system with correct op name and input."""
         from vllm_fl.ops.activation import SiluAndMulFL
 
-        mock_call_op.return_value = torch.randn(2, 4)
+        mock_cached_op.return_value = torch.randn(2, 4)
         layer = SiluAndMulFL()
         x = torch.randn(2, 8)
 
         result = layer.forward_oot(x)
 
-        mock_call_op.assert_called_once_with("silu_and_mul", layer, x)
+        mock_cached_op.assert_called_once_with(layer, x)
         assert result.shape == (2, 4)

tests/unit_tests/ops/test_layernorm.py

@@ -19,8 +19,8 @@ def __init__(self):
         set_current_vllm_config(VllmConfig())
 
     @pytest.fixture
-    def mock_call_op(self):
-        with patch("vllm_fl.ops.layernorm.call_op") as mock:
+    def mock_cached_op(self):
+        with patch("vllm_fl.ops.layernorm._rms_norm") as mock:
             yield mock
 
     def test_init_creates_weight_parameter(self):
@@ -34,31 +34,30 @@ def test_init_creates_weight_parameter(self):
         assert layer.variance_epsilon == eps
         assert layer.weight.shape == (hidden_size,)
 
-    def test_forward_oot_dispatches_without_residual(self, mock_call_op):
+    def test_forward_oot_dispatches_without_residual(self, mock_cached_op):
         """Test forward_oot calls dispatch system correctly without residual."""
         from vllm_fl.ops.layernorm import RMSNormFL
 
         hidden_size = 128
-        mock_call_op.return_value = torch.randn(2, hidden_size)
+        mock_cached_op.return_value = torch.randn(2, hidden_size)
 
         layer = RMSNormFL(hidden_size=hidden_size)
         x = torch.randn(2, hidden_size)
 
         layer.forward_oot(x)
 
-        mock_call_op.assert_called_once()
-        call_args = mock_call_op.call_args
-        assert call_args[0][0] == "rms_norm"
-        assert call_args[0][1] is layer  # self
-        assert torch.equal(call_args[0][2], x)
-        assert call_args[0][3] is None  # residual should be None
+        mock_cached_op.assert_called_once()
+        call_args = mock_cached_op.call_args
+        assert call_args[0][0] is layer  # self
+        assert torch.equal(call_args[0][1], x)
+        assert call_args[0][2] is None  # residual should be None
 
-    def test_forward_oot_dispatches_with_residual(self, mock_call_op):
+    def test_forward_oot_dispatches_with_residual(self, mock_cached_op):
         """Test forward_oot passes residual to dispatch system."""
         from vllm_fl.ops.layernorm import RMSNormFL
 
         hidden_size = 128
-        mock_call_op.return_value = (
+        mock_cached_op.return_value = (
             torch.randn(2, hidden_size),
             torch.randn(2, hidden_size),
         )
@@ -69,9 +68,8 @@ def test_forward_oot_dispatches_with_residual(self, mock_call_op):
 
         layer.forward_oot(x, residual=residual)
 
-        mock_call_op.assert_called_once()
-        call_args = mock_call_op.call_args
-        assert call_args[0][0] == "rms_norm"
-        assert call_args[0][1] is layer  # self
-        assert torch.equal(call_args[0][2], x)
-        assert torch.equal(call_args[0][3], residual)
+        mock_cached_op.assert_called_once()
+        call_args = mock_cached_op.call_args
+        assert call_args[0][0] is layer  # self
+        assert torch.equal(call_args[0][1], x)
+        assert torch.equal(call_args[0][2], residual)

tests/unit_tests/ops/test_rotary_embedding.py

@@ -14,8 +14,8 @@ class TestRotaryEmbeddingFL:
     """Test RotaryEmbeddingFL class behavior."""
 
     @pytest.fixture
-    def mock_call_op(self):
-        with patch("vllm_fl.ops.rotary_embedding.call_op") as mock:
+    def mock_cached_op(self):
+        with patch("vllm_fl.ops.rotary_embedding._rotary_embedding") as mock:
             yield mock
 
     @pytest.fixture
@@ -25,7 +25,7 @@ def mock_parent_init(self):
         ):
             yield
 
-    def test_forward_oot_dispatches_correctly(self, mock_parent_init, mock_call_op):
+    def test_forward_oot_dispatches_correctly(self, mock_parent_init, mock_cached_op):
         """Test forward_oot calls dispatch system with correct arguments."""
         from vllm_fl.ops.rotary_embedding import RotaryEmbeddingFL
 
@@ -44,14 +44,17 @@ def test_forward_oot_dispatches_correctly(self, mock_parent_init, mock_call_op):
         layer.is_neox_style = True
         layer.cos_sin_cache = torch.randn(2048, 64)
 
-        mock_call_op.return_value = (torch.randn(4, 8, 32), torch.randn(4, 8, 32))
+        mock_cached_op.return_value = (
+            torch.randn(4, 8, 32),
+            torch.randn(4, 8, 32),
+        )
 
         positions = torch.tensor([0, 1, 2, 3])
         query = torch.randn(4, 8, 64)
         key = torch.randn(4, 8, 64)
 
         layer.forward_oot(positions, query, key)
 
-        mock_call_op.assert_called_once()
-        call_args = mock_call_op.call_args
-        assert call_args[0][0] == "rotary_embedding"
+        mock_cached_op.assert_called_once()
+        call_args = mock_cached_op.call_args
+        assert call_args[0][0] is layer

vllm_fl/__init__.py

@@ -128,6 +128,9 @@ def register_router():
     # fused_moe import chain triggers cutlass_scaled_mm_supports_fp8 on MUSA
     if current_platform.device_type == "musa":
         return
+    from vllm_fl.utils import is_oot_enabled
+    if not is_oot_enabled():
+        return
     from vllm_fl.ops.fused_moe.router import replace_router_with_fl
     replace_router_with_fl()
 

vllm_fl/dispatch/__init__.py

@@ -76,12 +76,15 @@
             - reference
 """
 
+import os
+
 from .types import OpImpl, BackendImplKind, BackendPriority, match_token
 from .registry import OpRegistry, OpRegistrySnapshot
 from .policy import (
     SelectionPolicy,
     PolicyManager,
     get_policy,
+    get_policy_epoch,
     set_global_policy,
     reset_global_policy,
     policy_context,
@@ -108,6 +111,7 @@
     enable_io_dump,
     disable_io_dump,
     io_dump_step,
+    is_dump_enabled,
 )
 from .io_common import list_model_layers, register_tensor_stat, tensor_stats
 
@@ -139,6 +143,94 @@ def resolve_op(op_name: str):
     return get_default_manager().resolve(op_name)
 
 
+# Fast-path opt-out: set VLLM_FL_OP_FAST_PATH=0 to disable per-op fn caching
+# in hot OOT layers and route every call back through OpManager.call.
+_OP_FAST_PATH_ENABLED = os.environ.get("VLLM_FL_OP_FAST_PATH", "1") == "1"
+
+
+class CachedOp:
+    """Resolve an op once at the call site and refresh on policy changes.
+
+    OpManager.call preserves fallback and IO-dump hooks, but it also pays the
+    manager/fallback path on every invocation. Hot layer paths can use CachedOp
+    to call the resolved implementation directly after the first lookup.
+
+    The cache is invalidated by both OpManager.policy_epoch and
+    PolicyManager.policy_epoch. The latter matters for policy_context() and
+    set_global_policy(), which can change the effective backend without
+    touching the OpManager instance.
+
+    Cache refresh is best-effort under concurrent calls. If another thread
+    changes policy at the same time, a call may observe the previous impl once
+    before the next epoch check refreshes it.
+    """
+
+    __slots__ = (
+        "_op_name",
+        "_impl",
+        "_use_manager_call",
+        "_manager_id",
+        "_manager_epoch",
+        "_policy_epoch",
+    )
+
+    def __init__(self, op_name: str) -> None:
+        self._op_name = op_name
+        self._impl = None
+        self._use_manager_call = False
+        self._manager_id = -1
+        self._manager_epoch = -1
+        self._policy_epoch = -1
+
+    def __call__(self, *args, **kwargs):
+        mgr = get_default_manager()
+
+        if not _OP_FAST_PATH_ENABLED:
+            return mgr.call(self._op_name, *args, **kwargs)
+
+        if is_dump_enabled():
+            return mgr.call(self._op_name, *args, **kwargs)
+
+        manager_epoch = mgr.policy_epoch
+        manager_id = id(mgr)
+        policy_epoch = get_policy_epoch()
+        if (
+            self._manager_id != manager_id
+            or self._manager_epoch != manager_epoch
+            or self._policy_epoch != policy_epoch
+        ):
+            self._impl = None
+            self._use_manager_call = False
+
+        if self._use_manager_call:
+            return mgr.call(self._op_name, *args, **kwargs)
+
+        impl = self._impl
+        if (
+            impl is None
+            or self._manager_id != manager_id
+            or self._manager_epoch != manager_epoch
+            or self._policy_epoch != policy_epoch
+        ):
+            impl = mgr._resolve_impl(self._op_name)
+            mgr._record_first_use(self._op_name, impl)
+            self._impl = impl
+            # resolve() can initialize the manager and bump its epoch.
+            self._manager_id = manager_id
+            self._manager_epoch = mgr.policy_epoch
+            self._policy_epoch = get_policy_epoch()
+
+        try:
+            return impl.fn(*args, **kwargs)
+        except Exception:
+            self._impl = None
+            if get_policy().strict:
+                raise
+            mgr._mark_failed_impl(self._op_name, impl.impl_id)
+            self._use_manager_call = True
+            return mgr.call(self._op_name, *args, **kwargs)
+
+
 __all__ = [
     # Types
     "OpImpl",
@@ -152,6 +244,7 @@ def resolve_op(op_name: str):
     "SelectionPolicy",
     "PolicyManager",
     "get_policy",
+    "get_policy_epoch",
     "set_global_policy",
     "reset_global_policy",
     "policy_context",
@@ -188,4 +281,5 @@ def resolve_op(op_name: str):
     # Convenience functions
     "call_op",
     "resolve_op",
+    "CachedOp",
 ]

vllm_fl/dispatch/backends/vendor/metax/patches/gdn_linear_attn.py

@@ -112,4 +112,3 @@ def forward(
         core_attn_out = core_attn_out.reshape(z_shape_og)
         core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
         output[:num_tokens], _ = self.out_proj(core_attn_out)
-