perf: optimize TypeBasedDispatcher using dict for O(1) lookup (#12001)

python/sglang/utils.py

@@ -13,6 +13,7 @@ import time
 import traceback
 import urllib.request
 import weakref
+from collections import OrderedDict
 from concurrent.futures import ThreadPoolExecutor
 from functools import wraps
 from io import BytesIO
@@ -480,21 +481,45 @@ def wait_for_server(base_url: str, timeout: int = None) -> None:
 
 class TypeBasedDispatcher:
     def __init__(self, mapping: List[Tuple[Type, Callable]]):
-        self._mapping = mapping
+        # Use dictionary for fast exact type matching, using OrderedDict(mapping)
+        # to maintains registration order
+        self._mapping = OrderedDict(mapping)
+        # MRO cache for inheritance-based matching
+        self._mro_cache = {}
         self._fallback_fn = None
 
     def add_fallback_fn(self, fallback_fn: Callable):
         self._fallback_fn = fallback_fn
 
     def __iadd__(self, other: "TypeBasedDispatcher"):
-        self._mapping.extend(other._mapping)
+        for ty, fn in other._mapping.items():
+            if ty not in self._mapping:
+                self._mapping[ty] = fn
+
+        self._mro_cache.clear()
         return self
 
     def __call__(self, obj: Any):
-        for ty, fn in self._mapping:
+        obj_type = type(obj)
+        # 1. First try exact match(o(1))
+        fn = self._mapping.get(obj_type)
+        if fn is not None:
+            return fn(obj)
+
+        # 2. If exact match fails, check MRO cache
+        cached_fn = self._mro_cache.get(obj_type)
+        if cached_fn is not None:
+            return cached_fn(obj)
+
+        # 3.search in registration order for compatible type(maintains origin behavior)
+        for ty, fn in self._mapping.items():
             if isinstance(obj, ty):
+                self._mro_cache[obj_type] = fn
                 return fn(obj)
 
+        # 4. if no matching type found, cache this result
+        self._mro_cache[obj_type] = None
+
         if self._fallback_fn is not None:
             return self._fallback_fn(obj)
         raise ValueError(f"Invalid object: {obj}")

test/srt/run_suite.py

@@ -422,6 +422,7 @@ suite_amd = {
         TestFile("quant/test_fused_rms_fp8_group_quant.py", 10),
         TestFile("rl/test_update_weights_from_disk.py", 210),
         TestFile("test_abort.py", 51),
+        TestFile("test_bench_typebaseddispatcher.py", 10),
         TestFile("test_chunked_prefill.py", 410),
         TestFile("test_create_kvindices.py", 2),
         TestFile("test_eval_fp8_accuracy.py", 303),
@@ -454,6 +455,7 @@ suite_amd = {
         # TestFile("test_triton_attention_kernels.py", 4),
         TestFile("test_triton_attention_backend.py", 150),
         TestFile("test_triton_sliding_window.py", 250),
+        TestFile("test_type_based_dispatcher.py", 10),
         TestFile("test_wave_attention_kernels.py", 2),
         # Disabled temporarily
         # TestFile("test_vlm_input_format.py", 300),

test/srt/test_bench_typebaseddispatcher.py

@@ -0,0 +1,261 @@
+import timeit
+from typing import Any, Callable, List, Tuple, Type
+
+from sglang.utils import TypeBasedDispatcher
+
+
+class TypeBasedDispatcherList:
+    def __init__(self, mapping: List[Tuple[Type, Callable]]):
+        self._mapping = mapping
+        self._fallback_fn = None
+
+    def add_fallback_fn(self, fallback_fn: Callable):
+        self._fallback_fn = fallback_fn
+
+    def __iadd__(self, other: "TypeBasedDispatcher"):
+        self._mapping.extend(other._mapping)
+        return self
+
+    def __call__(self, obj: Any):
+        for ty, fn in self._mapping:
+            if isinstance(obj, ty):
+                return fn(obj)
+
+        if self._fallback_fn is not None:
+            return self._fallback_fn(obj)
+        raise ValueError(f"Invalid object: {obj}")
+
+
+def create_test_mapping(num_types=30):
+    types = [type(f"RequestType{i}", (), {}) for i in range(num_types)]
+
+    def create_handler(i):
+        def handler(req):
+            return f"handler{i}"
+
+        return handler
+
+    handlers = [create_handler(i) for i in range(num_types)]
+
+    return list(zip(types, handlers))
+
+
+def test_inheritance():
+    print("\n" + "=" * 60)
+    print("test for inheritance")
+    print("=" * 60)
+
+    class BaseRequest:
+        pass
+
+    def base_handler(req):
+        return "base_handler"
+
+    class DerivedRequest(BaseRequest):
+        pass
+
+    mapping = [(BaseRequest, base_handler)]
+    dict_dispatcher = TypeBasedDispatcher(mapping)
+
+    derived_obj = DerivedRequest()
+    expected = "base_handler"
+
+    # This test will fail with the current implementation, but pass with the suggested MRO-based fix
+    result_dict = dict_dispatcher(derived_obj)
+    assert result_dict == expected, f"Expected '{expected}', but got '{result_dict}'"
+    print("Pass: dict dispatcher handles inheritance.")
+
+
+def benchmark_with_inheritance():
+    """Performance test with inheritance scenarios"""
+    print("\nBenchmarking with inheritance scenarios...")
+
+    # Create type hierarchy with inheritance relationships
+    class BaseType:
+        pass
+
+    class ChildType1(BaseType):
+        pass
+
+    class ChildType2(BaseType):
+        pass
+
+    class GrandChildType(ChildType1):
+        pass
+
+    class UnrelatedType:
+        pass
+
+    def base_handler(obj):
+        return "handled"
+
+    mapping = [(BaseType, base_handler)]
+    dispatcher = TypeBasedDispatcher(mapping)
+
+    test_cases = [
+        BaseType(),
+        ChildType1(),
+        ChildType2(),
+        GrandChildType(),
+        UnrelatedType(),
+    ]
+
+    # Test first call (includes MRO lookup)
+    first_call_times = []
+    for case in test_cases:
+        if not isinstance(case, UnrelatedType):
+            time_taken = timeit.timeit(lambda: dispatcher(case), number=1000)
+            first_call_times.append(time_taken)
+
+    # Test subsequent calls (using cache)
+    cached_call_times = []
+    for case in test_cases:
+        if not isinstance(case, UnrelatedType):
+            time_taken = timeit.timeit(lambda: dispatcher(case), number=1000)
+            cached_call_times.append(time_taken)
+
+    print(
+        f"First call (with MRO lookup): {sum(first_call_times)/len(first_call_times):.6f}s avg"
+    )
+    print(f"Cached call: {sum(cached_call_times)/len(cached_call_times):.6f}s avg")
+    print(f"Caching improvement: {sum(first_call_times)/sum(cached_call_times):.2f}x")
+
+
+def benchmark_dispatchers():
+    mapping = create_test_mapping(30)
+    list_dispatcher = TypeBasedDispatcherList(mapping)
+    dist_dispatcher = TypeBasedDispatcher(mapping)
+
+    test_cases = []
+    for _, (ty, _) in enumerate(mapping):
+        test_cases.append(ty())
+
+    test_scenarios = [
+        ("the first", [test_cases[0]] * 1000),
+        ("the middle", [test_cases[len(test_cases) // 2]] * 1000),
+        ("the last", [test_cases[-1]] * 1000),
+        ("the random", test_cases * 1000),
+    ]
+
+    print("=" * 60)
+    print("TypeBasedDispatcher benchmark test")
+    print("=" * 60)
+
+    for scenario_name, cases in test_scenarios:
+        print(f"\ntest scenario: {scenario_name}")
+        print(f"\ntest numbers: {len(cases)}")
+
+        list_time = timeit.timeit(
+            lambda: [list_dispatcher(case) for case in cases], number=10
+        )
+
+        dict_time = timeit.timeit(
+            lambda: [dist_dispatcher(case) for case in cases], number=10
+        )
+
+        print(f"for list: {list_time:.4f} s")
+        print(f"for dict: {dict_time:.4f} s")
+        print(f"improvement: {list_time/dict_time:.2f} x")
+        print(f"time reduce: {(1-dict_time/list_time) * 100:.1f} %")
+
+
+def test_memory_usage():
+    import sys
+
+    mapping = create_test_mapping(30)
+    list_dispatcher = TypeBasedDispatcherList(mapping)
+    dict_dispatcher = TypeBasedDispatcher(mapping)
+
+    print("\n" + "=" * 60)
+    print("compare memory used:")
+    print("=" * 60)
+
+    list_size = sys.getsizeof(list_dispatcher._mapping)
+    dict_size = sys.getsizeof(dict_dispatcher._mapping)
+
+    print(f"memory used by list version: {list_size} bytes")
+    print(f"memory used by dict version: {dict_size} bytes")
+    print(f"compare memory used by the two version: {dict_size - list_size} bytes")
+
+
+def test_edge_case():
+    """test for edge case"""
+    print("\n" + "=" * 60)
+    print("test for edge case")
+    print("=" * 60)
+
+    mapping = create_test_mapping(30)
+    list_dispatcher = TypeBasedDispatcherList(mapping)
+    dict_dispatcher = TypeBasedDispatcher(mapping)
+
+    test_obj = mapping[0][0]()
+    result1 = list_dispatcher(test_obj)
+    result2 = dict_dispatcher(test_obj)
+
+    assert result1 == result2
+    print("Pass for normal test")
+
+    class UnkownType:
+        pass
+
+    try:
+        list_dispatcher(UnkownType())
+        print("exception was thrown from list version as expected")
+    except ValueError:
+        print("exception thrown from list version was processed...")
+
+    try:
+        dict_dispatcher(UnkownType())
+        print("exception was thrown from dict version as expected")
+    except ValueError:
+        print("exception thrown from dict version was processed...")
+
+
+def simulate_real_workload():
+    """simulate real workload"""
+
+    print("\n" + "=" * 60)
+    print("simulate real workload")
+    print("=" * 60)
+
+    mapping = create_test_mapping(30)
+
+    request_distribution = {
+        0: 0.2,
+        5: 0.3,
+        10: 0.1,
+        15: 0.15,
+    }
+
+    list_dispatcher = TypeBasedDispatcherList(mapping)
+    dict_dispatcher = TypeBasedDispatcher(mapping)
+
+    test_requests = []
+    for idx, prob in request_distribution.items():
+        count = int(1000 * prob)
+        test_requests.extend([mapping[idx][0]()] * count)
+
+    remaining = 1000 - len(test_requests)
+    for i in range(remaining):
+        test_requests.append(mapping[i % len(mapping)][0]())
+
+    list_time = timeit.timeit(
+        lambda: [list_dispatcher(req) for req in test_requests], number=100
+    )
+
+    dict_time = timeit.timeit(
+        lambda: [dict_dispatcher(req) for req in test_requests], number=100
+    )
+
+    print(f"list version: {list_time:.4f} s")
+    print(f"dict version: {dict_time:.4f} s")
+    print(f"improvement: {list_time/dict_time:.2f} x")
+
+
+if __name__ == "__main__":
+    benchmark_dispatchers()
+    test_memory_usage()
+    test_edge_case()
+    simulate_real_workload()
+    test_inheritance()
+    benchmark_with_inheritance()

test/srt/test_type_based_dispatcher.py

@@ -0,0 +1,222 @@
+# tests/benchmarks/test_type_dispatcher_e2e.py
+"""
+E2E test for TypeBasedDispatcher optimization.
+Tests real-world scenarios with actual request types.
+"""
+
+import timeit
+import unittest
+
+from sglang.srt.managers.io_struct import SamplingParams
+from sglang.utils import TypeBasedDispatcher
+
+
+class TestTypeBasedDispatcher(unittest.TestCase):
+    """Unit tests for TypeBasedDispatcher e2e performance."""
+
+    def test_type_dispatcher_e2e_performance(self):
+        """End-to-end performance test with real request types"""
+        print("E2E Performance Test for TypeBasedDispatcher")
+        print("=" * 50)
+
+        from sglang.srt.managers.io_struct import (
+            AbortReq,
+            BatchTokenizedEmbeddingReqInput,
+            BatchTokenizedGenerateReqInput,
+            ClearHiCacheReqInput,
+            CloseSessionReqInput,
+            DestroyWeightsUpdateGroupReqInput,
+            ExpertDistributionReq,
+            FlushCacheReqInput,
+            FreezeGCReq,
+            GetInternalStateReq,
+            GetLoadReqInput,
+            GetWeightsByNameReqInput,
+            InitWeightsSendGroupForRemoteInstanceReqInput,
+            InitWeightsUpdateGroupReqInput,
+            LoadLoRAAdapterReqInput,
+            OpenSessionReqInput,
+            ProfileReq,
+            ReleaseMemoryOccupationReqInput,
+            ResumeMemoryOccupationReqInput,
+            RpcReqInput,
+            SendWeightsToRemoteInstanceReqInput,
+            SetInternalStateReq,
+            SlowDownReqInput,
+            TokenizedEmbeddingReqInput,
+            TokenizedGenerateReqInput,
+            UnloadLoRAAdapterReqInput,
+            UpdateWeightFromDiskReqInput,
+            UpdateWeightsFromIPCReqInput,
+            UpdateWeightsFromTensorReqInput,
+        )
+
+        mapping = [
+            (TokenizedGenerateReqInput, lambda req: "generate_handled"),
+            (TokenizedEmbeddingReqInput, lambda req: "embedding_handled"),
+            (BatchTokenizedGenerateReqInput, lambda req: "batch_generate_handled"),
+            (
+                BatchTokenizedEmbeddingReqInput,
+                lambda req: "batch_generate_embedding_handled",
+            ),
+            (FlushCacheReqInput, lambda req: "flush_cache_handled"),
+            (ClearHiCacheReqInput, lambda req: "clear_hicache_handled"),
+            (AbortReq, lambda req: "abort_handled"),
+            (OpenSessionReqInput, lambda req: "open_session_handled"),
+            (CloseSessionReqInput, lambda req: "close_session_handled"),
+            (
+                UpdateWeightFromDiskReqInput,
+                lambda req: "update_weights_from_disk_handled",
+            ),
+            (
+                InitWeightsUpdateGroupReqInput,
+                lambda req: "init_weights_update_group_handled",
+            ),
+            (
+                DestroyWeightsUpdateGroupReqInput,
+                lambda req: "destroy_weights_update_group_handled",
+            ),
+            (
+                InitWeightsSendGroupForRemoteInstanceReqInput,
+                lambda req: "init_weights_send_group_for_remote_instance_handled",
+            ),
+            (
+                SendWeightsToRemoteInstanceReqInput,
+                lambda req: "send_weights_to_remote_instance_handled",
+            ),
+            (
+                UpdateWeightsFromTensorReqInput,
+                lambda req: "update_weights_from_tensor_handled",
+            ),
+            (
+                UpdateWeightsFromIPCReqInput,
+                lambda req: "update_weights_from_ipc_handled",
+            ),
+            (GetWeightsByNameReqInput, lambda req: "get_weights_by_name_handled"),
+            (
+                ReleaseMemoryOccupationReqInput,
+                lambda req: "release_memory_occupation_handled",
+            ),
+            (
+                ResumeMemoryOccupationReqInput,
+                lambda req: "resume_memory_occupation_handled",
+            ),
+            (SlowDownReqInput, lambda req: "slow_down_handled"),
+            (ProfileReq, lambda req: "profile_handled"),
+            (FreezeGCReq, lambda req: "freeze_gc_handled"),
+            (GetInternalStateReq, lambda req: "get_internal_state_handled"),
+            (SetInternalStateReq, lambda req: "set_internal_state_handled"),
+            (RpcReqInput, lambda req: "rpc_request_handled"),
+            (ExpertDistributionReq, lambda req: "expert_distribution_handled"),
+            (LoadLoRAAdapterReqInput, lambda req: "load_lora_adapter_handled"),
+            (UnloadLoRAAdapterReqInput, lambda req: "unload_lora_adapter_handled"),
+            (GetLoadReqInput, lambda req: "get_load_handled"),
+        ]
+
+        # Create requests that conforms to the real distribution
+        test_requests = []
+
+        test_requests.append(
+            TokenizedGenerateReqInput(
+                input_text="",
+                input_ids=[1, 2],
+                mm_inputs=dict(),
+                sampling_params=SamplingParams(),
+                return_logprob=False,
+                logprob_start_len=0,
+                top_logprobs_num=0,
+                token_ids_logprob=[1, 2],
+                stream=False,
+            )
+        )
+
+        test_requests.append(
+            TokenizedEmbeddingReqInput(
+                input_text="",
+                input_ids=[1, 2],
+                image_inputs=dict(),
+                token_type_ids=[1, 2],
+                sampling_params=SamplingParams(),
+            )
+        )
+
+        test_requests.append(
+            BatchTokenizedGenerateReqInput(
+                batch=[
+                    TokenizedGenerateReqInput(
+                        input_text="",
+                        input_ids=[1, 2],
+                        mm_inputs=dict(),
+                        sampling_params=SamplingParams(),
+                        return_logprob=False,
+                        logprob_start_len=0,
+                        top_logprobs_num=0,
+                        token_ids_logprob=[1, 2],
+                        stream=False,
+                    )
+                ]
+            )
+        )
+        test_requests.append(
+            BatchTokenizedEmbeddingReqInput(
+                batch=[
+                    TokenizedEmbeddingReqInput(
+                        input_text="",
+                        input_ids=[1, 2],
+                        image_inputs=dict(),
+                        token_type_ids=[1, 2],
+                        sampling_params=SamplingParams(),
+                    )
+                ]
+            )
+        )
+
+        test_requests.append(FlushCacheReqInput())
+        test_requests.append(ClearHiCacheReqInput())
+        test_requests.append(AbortReq())
+        test_requests.append(OpenSessionReqInput(capacity_of_str_len=0))
+        test_requests.append(CloseSessionReqInput(session_id=""))
+        test_requests.append(UpdateWeightFromDiskReqInput(model_path=""))
+        test_requests.append(
+            InitWeightsUpdateGroupReqInput(
+                master_address="",
+                master_port=0,
+                rank_offset=0,
+                world_size=0,
+                group_name="",
+            )
+        )
+        test_requests.append(DestroyWeightsUpdateGroupReqInput())
+        test_requests.append(
+            InitWeightsSendGroupForRemoteInstanceReqInput(
+                master_address="", ports="", group_name="", world_size=0, group_rank=0
+            )
+        )
+        test_requests.append(
+            SendWeightsToRemoteInstanceReqInput(master_address="", ports="")
+        )
+        test_requests.append(
+            UpdateWeightsFromTensorReqInput(serialized_named_tensors=[])
+        )
+        test_requests.append(GetWeightsByNameReqInput(name=""))
+        test_requests.append(ReleaseMemoryOccupationReqInput())
+        test_requests.append(RpcReqInput(method=""))
+        test_requests.append(GetLoadReqInput())
+
+        dispatcher = TypeBasedDispatcher(mapping)
+
+        # test
+        time_taken = timeit.timeit(
+            lambda: [dispatcher(req) for req in test_requests],
+            number=100,  # Average of 100 runs
+        )
+
+        print(f"Total requests: {len(test_requests)}")
+        print(f"Time taken: {time_taken:.4f}s")
+        print(f"Requests per second: {len(test_requests) * 100 / time_taken:.0f}")
+
+        return time_taken
+
+
+if __name__ == "__main__":
+    unittest.main()