From 37e8724ef3b069740da0b4df4a666c34059015f9 Mon Sep 17 00:00:00 2001 From: xlzheng Date: Sat, 15 Nov 2025 22:02:42 +0800 Subject: [PATCH] perf: optimize TypeBasedDispatcher using dict for O(1) lookup (#12001) --- python/sglang/utils.py | 31 ++- test/srt/run_suite.py | 2 + test/srt/test_bench_typebaseddispatcher.py | 261 +++++++++++++++++++++ test/srt/test_type_based_dispatcher.py | 222 ++++++++++++++++++ 4 files changed, 513 insertions(+), 3 deletions(-) create mode 100644 test/srt/test_bench_typebaseddispatcher.py create mode 100644 test/srt/test_type_based_dispatcher.py diff --git a/python/sglang/utils.py b/python/sglang/utils.py index 041630093..d378f22b3 100644 --- a/python/sglang/utils.py +++ b/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}") diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index d516feb43..add4e7e11 100644 --- a/test/srt/run_suite.py +++ b/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), diff --git a/test/srt/test_bench_typebaseddispatcher.py b/test/srt/test_bench_typebaseddispatcher.py new file mode 100644 index 000000000..1c048b89c --- /dev/null +++ b/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() diff --git a/test/srt/test_type_based_dispatcher.py b/test/srt/test_type_based_dispatcher.py new file mode 100644 index 000000000..1049a40ae --- /dev/null +++ b/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()