perf: optimize TypeBasedDispatcher using dict for O(1) lookup (#12001)
This commit is contained in:
xlzheng
@@ -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}")
|
||||
|
||||
@@ -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),
|
||||
|
||||
261
test/srt/test_bench_typebaseddispatcher.py
Normal file
261
test/srt/test_bench_typebaseddispatcher.py
Normal file
@@ -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()
|
||||
222
test/srt/test_type_based_dispatcher.py
Normal file
222
test/srt/test_type_based_dispatcher.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user