diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 8f39cad01..778c7e598 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -613,7 +613,6 @@ class Scheduler( self.require_mlp_sync = require_mlp_sync(self.server_args) def init_tp_model_worker(self): - worker_kwargs = dict( server_args=self.server_args, gpu_id=self.gpu_id, @@ -826,7 +825,6 @@ class Scheduler( self.tree_cache = SWAChunkCache(params) else: - if envs.SGLANG_EXPERIMENTAL_CPP_RADIX_TREE.get(): # lazy import to avoid JIT overhead from sglang.srt.mem_cache.radix_cache_cpp import RadixCacheCpp @@ -1622,7 +1620,6 @@ class Scheduler( ): recv_reqs, abort_reqs = self.mm_receiver.process_waiting_requests(recv_reqs) for req, error_msg, error_code in abort_reqs: - status_code = ( HTTPStatus.BAD_REQUEST if error_code == 400 @@ -1878,6 +1875,7 @@ class Scheduler( self.metrics_collector if self.enable_metrics else None ), routing_key=recv_req.routing_key, + extra_key=recv_req.extra_key, http_worker_ipc=recv_req.http_worker_ipc, dllm_config=self.dllm_config, time_stats=recv_req.time_stats, diff --git a/python/sglang/srt/session/session_controller.py b/python/sglang/srt/session/session_controller.py index befab0c1e..ce98514b8 100644 --- a/python/sglang/srt/session/session_controller.py +++ b/python/sglang/srt/session/session_controller.py @@ -235,6 +235,7 @@ class Session: return_routed_experts=req.return_routed_experts, priority=req.priority, routing_key=req.routing_key, + extra_key=req.extra_key, http_worker_ipc=req.http_worker_ipc, time_stats=req.time_stats, ) diff --git a/test/manual/openai_server/features/test_cache_report.py b/test/manual/openai_server/features/test_cache_report.py index 6a5f7bd8a..ec45eb119 100644 --- a/test/manual/openai_server/features/test_cache_report.py +++ b/test/manual/openai_server/features/test_cache_report.py @@ -24,6 +24,7 @@ class TestCacheReport(CustomTestCase): timeout=300, other_args=[ "--chunked-prefill-size=40", + "--attention-backend=triton", "--enable-cache-report", ], ) @@ -206,6 +207,14 @@ class TestCacheReport(CustomTestCase): # asyncio.run(run_test()) + @staticmethod + def _get_cached_tokens(response) -> int: + """Extract cached_tokens from response, returning 0 if prompt_tokens_details is None.""" + details = response.usage.prompt_tokens_details + if details is None: + return 0 + return int(details.cached_tokens) + def test_cache_salt_effectiveness(self): print("=" * 100) print("Testing cache_salt effectiveness") @@ -221,7 +230,7 @@ class TestCacheReport(CustomTestCase): max_tokens=10, extra_body={"cache_salt": "salt1"}, ) - cached_tokens_1_first = int(response1.usage.prompt_tokens_details.cached_tokens) + cached_tokens_1_first = self._get_cached_tokens(response1) prompt_tokens_1 = int(response1.usage.prompt_tokens) print( f"First request with salt1 - cached_tokens: {cached_tokens_1_first}, prompt_tokens: {prompt_tokens_1}" @@ -235,9 +244,7 @@ class TestCacheReport(CustomTestCase): max_tokens=10, extra_body={"cache_salt": "salt1"}, ) - cached_tokens_1_second = int( - response2.usage.prompt_tokens_details.cached_tokens - ) + cached_tokens_1_second = self._get_cached_tokens(response2) print( f"Second request with salt1 - cached_tokens: {cached_tokens_1_second}, prompt_tokens: {prompt_tokens_1}" ) @@ -258,7 +265,7 @@ class TestCacheReport(CustomTestCase): max_tokens=10, extra_body={"cache_salt": "salt2"}, ) - cached_tokens_2_first = int(response3.usage.prompt_tokens_details.cached_tokens) + cached_tokens_2_first = self._get_cached_tokens(response3) print(f"First request with salt2 - cached_tokens: {cached_tokens_2_first}") # Verify no cache hit for different salt (should be similar to first request with salt1) @@ -274,14 +281,12 @@ class TestCacheReport(CustomTestCase): max_tokens=10, extra_body={"cache_salt": "salt2"}, ) - cached_tokens_2_second = int( - response4.usage.prompt_tokens_details.cached_tokens - ) + cached_tokens_2_second = self._get_cached_tokens(response4) print(f"Second request with salt2 - cached_tokens: {cached_tokens_2_second}") # Verify cache hit for salt2 assert ( - cached_tokens_2_second == cached_tokens_2_first + cached_tokens_2_second > cached_tokens_2_first ), "Should have cache hit with same cache_salt for salt2"