From f393964a31edea0b4ca702d7c3c58bc8e6bc3cee Mon Sep 17 00:00:00 2001 From: Max de Bayser Date: Thu, 21 May 2026 14:23:19 -0400 Subject: [PATCH 001/106] Add stub of function to calculate required blocks Signed-off-by: Max de Bayser --- sendnn_inference/v1/core/scheduler.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 9dc742d4c..3fd412a0f 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -253,6 +253,16 @@ def adjust_computed_tokens( # Otherwise just account for the left padding return computed_tokens - left_padding + + def get_block_required_for_decode_batch(self) -> int: + """ + Returns the number of blocks that the current decode batch needs to + finish all requests. Reducing the number of available blocks below + this number will cause deadlocks. + """ + raise NotImplementedError() + + def schedule(self) -> "SchedulerOutput": """ The chunked prefill scheduling policy is enforced in this method, then From 526fac3d6032d84939338dc252f26a63022c1864 Mon Sep 17 00:00:00 2001 From: Max de Bayser Date: Thu, 21 May 2026 14:33:13 -0400 Subject: [PATCH 002/106] fix typo Signed-off-by: Max de Bayser --- sendnn_inference/v1/core/scheduler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 3fd412a0f..59f43844a 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -254,7 +254,7 @@ def adjust_computed_tokens( return computed_tokens - left_padding - def get_block_required_for_decode_batch(self) -> int: + def get_blocks_required_for_decode_batch(self) -> int: """ Returns the number of blocks that the current decode batch needs to finish all requests. Reducing the number of available blocks below From 0c696970425af80d64e760d807693c8fb9b58ecb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 20 May 2026 16:48:19 +0000 Subject: [PATCH 003/106] input_batch adaptation for preemption MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../v1/worker/spyre_model_runner.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index 94cac9eee..2e2601958 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -1447,6 +1447,29 @@ def _update_batch(self, scheduler_output: SchedulerOutput): - Refreshes metadata for logits processors """ req_data = scheduler_output.scheduled_cached_reqs + + # Synchronize input_batch with scheduler output: remove requests that are not in scheduler output + # This handles preemption cases where scheduler temporarily removes requests from running queue + scheduled_req_ids = set(req_data.req_ids) + current_batch_req_ids = set(self.input_batch.req_id_to_index.keys()) + + # Find requests that are in input_batch but not in scheduler output (preempted) + preempted_req_ids = current_batch_req_ids - scheduled_req_ids + for req_id in preempted_req_ids: + # Only remove if it's not a finished request (finished requests are handled separately) + if req_id not in (scheduler_output.finished_req_ids or []): + logger.info(f"Removing preempted request {req_id} from input_batch") + self.input_batch.remove_request(req_id) + + # Find requests that are in scheduler output but not in input_batch (restored from preemption) + restored_req_ids = scheduled_req_ids - current_batch_req_ids + for req_id in restored_req_ids: + # Add back the request that was preempted + if req_id in self.requests: + logger.info(f"Restoring preempted request {req_id} to input_batch") + req_state = self.requests[req_id] + self.input_batch.add_request(req_state) + for i, req_id in enumerate(req_data.req_ids): req_state: SamplingRequestState = self.requests[req_id] From 9cc43104d96eca68f359a6d2fe6ae4c67902901f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Thu, 21 May 2026 19:07:40 +0000 Subject: [PATCH 004/106] run pre-commit hook MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 2 -- .../v1/worker/spyre_model_runner.py | 28 ++++++++++--------- 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 59f43844a..26634bf54 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -253,7 +253,6 @@ def adjust_computed_tokens( # Otherwise just account for the left padding return computed_tokens - left_padding - def get_blocks_required_for_decode_batch(self) -> int: """ Returns the number of blocks that the current decode batch needs to @@ -261,7 +260,6 @@ def get_blocks_required_for_decode_batch(self) -> int: this number will cause deadlocks. """ raise NotImplementedError() - def schedule(self) -> "SchedulerOutput": """ diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index 2e2601958..8da51623d 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -1447,29 +1447,31 @@ def _update_batch(self, scheduler_output: SchedulerOutput): - Refreshes metadata for logits processors """ req_data = scheduler_output.scheduled_cached_reqs - - # Synchronize input_batch with scheduler output: remove requests that are not in scheduler output - # This handles preemption cases where scheduler temporarily removes requests from running queue + + # Synchronize input_batch with scheduler output: remove requests + # that are not in scheduler output. This handles hold back cases + # where scheduler temporarily removes requests from running queue scheduled_req_ids = set(req_data.req_ids) current_batch_req_ids = set(self.input_batch.req_id_to_index.keys()) - - # Find requests that are in input_batch but not in scheduler output (preempted) - preempted_req_ids = current_batch_req_ids - scheduled_req_ids - for req_id in preempted_req_ids: + + # Find requests that are in input_batch but not in scheduler output (held back) + heldback_req_ids = current_batch_req_ids - scheduled_req_ids + for req_id in heldback_req_ids: # Only remove if it's not a finished request (finished requests are handled separately) if req_id not in (scheduler_output.finished_req_ids or []): - logger.info(f"Removing preempted request {req_id} from input_batch") + logger.info("Removing held back request %s from input_batch", req_id) self.input_batch.remove_request(req_id) - - # Find requests that are in scheduler output but not in input_batch (restored from preemption) + + # Find requests that are in scheduler output but not in input_batch + # (restored from holding back) restored_req_ids = scheduled_req_ids - current_batch_req_ids for req_id in restored_req_ids: - # Add back the request that was preempted + # Add back the request that was held back if req_id in self.requests: - logger.info(f"Restoring preempted request {req_id} to input_batch") + logger.info("Restoring held back request %s to input_batch", req_id) req_state = self.requests[req_id] self.input_batch.add_request(req_state) - + for i, req_id in enumerate(req_data.req_ids): req_state: SamplingRequestState = self.requests[req_id] From 7e2c5168c579582a09543bbcb5f887f9e95f9dc3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Thu, 21 May 2026 22:17:20 +0000 Subject: [PATCH 005/106] make constraint greedy when scheduling new requests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 84 ++++++--------------------- 1 file changed, 18 insertions(+), 66 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 26634bf54..afb728a7a 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -475,18 +475,13 @@ def _satisfies_last_chunk_constraints(self, request: Request) -> bool: n_blocks = math.floor(max(self.tkv, prompt_len) / self.block_size) new_req_tkv = n_blocks * self.block_size + prompt_len % self.block_size - # check that the number of requested tokens can be served for the - # new sequence (optimal condition) - # note that the -1 comes from the token we generate during prefill - cond2 = request.max_tokens - 1 <= (max_context_len - new_req_tkv) + # check that no tkv exceeds max_context_len (immediate constraint check) + cond2 = new_req_tkv <= max_context_len # check cond2 for all other sequences in the current decode batch for req in decoding_requests: # current tkv of the (left aligned) decode sequence dec_req_tkv = n_blocks * self.block_size + req.num_computed_tokens % self.block_size - n_generated_output_tokens = req.num_computed_tokens - req.num_prompt_tokens - max_tokens_remaining = req.max_tokens - n_generated_output_tokens - # note that the -1 comes from the token we generate during prefill - cond2_current = max_tokens_remaining - 1 <= (max_context_len - dec_req_tkv) + cond2_current = dec_req_tkv <= max_context_len cond2 = cond2 and cond2_current # early exiting loop if violated 2nd condition if not cond2: @@ -495,7 +490,7 @@ def _satisfies_last_chunk_constraints(self, request: Request) -> bool: # check that batch size x tkv is smaller than the max supported number # Note: using max_tkv is a conservative upper bound here. For the # optimal check we need model runner to return per sequence tkvs - cond3 = lambda: self.check_batch_tkv_limit_cp( + cond3 = lambda: self.check_batch_tkv_limit( request=request, new_req_tkv=new_req_tkv, running=decoding_requests, @@ -521,77 +516,34 @@ def _has_scheduling_priority(self, request): num_prefills = len(self.waiting) + len(self.ongoing_prefills) return num_prefills < max_concurrent_prefills - def check_batch_tkv_limit_cp(self, request: Request, new_req_tkv: int, running) -> bool: + def check_batch_tkv_limit(self, request: Request, new_req_tkv: int, running) -> bool: """ Check whether adding a new sequence to the decode batch would violate Spyre's maximum batch volume constraint for chunked prefill. In Spyre, the product of `batch_size` and the current `tkv` (tokens-per-sequence) must not exceed the limit defined by - `VLLM_DT_MAX_BATCH_TKV_LIMIT`. Before scheduling a new sequence, - we must ensure that this constraint will hold for all decoding - steps that result from combining the new sequence with the currently - running decode batch. - - This implementation: - 1. Computes the maximum possible `tkv` for each sequence in the - decode batch. - 2. Sorts these values in ascending order. - 3. Iterates through them, stopping once the `tkv` of the new sequence. - is reached. Remaining sequences do not need to be checked explicitly, - since they were validated when they were added (by inductive reasoning). - - Note: drawing explaining the algorithm in more detail uploaded here: - https://github.com/torch-spyre/sendnn-inference/pull/363#issuecomment-3173605517 + `VLLM_DT_MAX_BATCH_TKV_LIMIT`. This checks the immediate constraint + only, not future states. """ + # Calculate the current max tkv across all sequences + # new_req_tkv is already the current tkv for the new request + current_max_tkv = new_req_tkv - # Compute the effective token length of the new request - # Rounded up to the nearest block size to account for potential padding - new_req_max_tkv = round_up_to_block_size(new_req_tkv + request.max_tokens - 1) - # Extra block of slack: left-padding can push a sequence's runtime tkv up to - # one block past the scheduler's estimate when the batch re-aligns on admission. - new_req_max_tkv += self.block_size - - # Compute token lengths for all running requests (decode batch) - decode_req_max_tkvs = [] - # Decide new tkv based on max of current tkv or new request prompt tokens - dec_req_tkv = max(self.tkv, request.num_prompt_tokens) + # Check current tkv for all running requests + n_blocks = math.floor(max(self.tkv, request.num_prompt_tokens) / self.block_size) for req in running: - n_generated_output_tokens = req.num_computed_tokens - req.num_prompt_tokens - # Rounded up to the nearest block size to account for potential padding - dec_req_max_tkv = round_up_to_block_size( - dec_req_tkv + (req.max_tokens - n_generated_output_tokens) - 1 - ) - # Extra block of slack: left-padding can push a sequence's runtime tkv up to - # one block past the scheduler's estimate when the batch re-aligns on admission. - dec_req_max_tkv += self.block_size - - decode_req_max_tkvs.append(dec_req_max_tkv) - - # Sort decode requests token lengths in ascending order - decode_req_max_tkvs.sort() + dec_req_tkv = n_blocks * self.block_size + req.num_computed_tokens % self.block_size + current_max_tkv = max(current_max_tkv, dec_req_tkv) - # Initialize values - # The request is already in the running queue if it has done a first - # chunked prefill + # Calculate batch size (including the new request if not already running) batch_size = len(running) if request not in running: batch_size += 1 - max_batch_tkv = 0 - - # Try adding the new request to the batch and check the max volume - for decode_req_max_tkv in decode_req_max_tkvs: - if new_req_max_tkv <= decode_req_max_tkv: - # If the new request is shorter, it limits the batch volume - max_batch_tkv = max(max_batch_tkv, batch_size * new_req_max_tkv) - break - else: - # Otherwise, use the current (longer) request's volume - max_batch_tkv = max(max_batch_tkv, batch_size * decode_req_max_tkv) - # decrease batch_size by 1 as the current request finished - batch_size -= 1 - return max_batch_tkv <= self.max_batch_tkv_limit + # Check immediate volume constraint + current_batch_tkv = batch_size * current_max_tkv + return current_batch_tkv <= self.max_batch_tkv_limit def finish_requests( self, From baa4de5b00239b696602112bdd43a0bfc66ce0c8 Mon Sep 17 00:00:00 2001 From: Max de Bayser Date: Thu, 21 May 2026 17:07:00 -0400 Subject: [PATCH 006/106] Fill get_blocks_required_for_decode_batch Signed-off-by: Max de Bayser --- sendnn_inference/v1/core/scheduler.py | 33 ++++++++++++++++++++++++--- 1 file changed, 30 insertions(+), 3 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 59f43844a..a3a7761be 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -253,15 +253,42 @@ def adjust_computed_tokens( # Otherwise just account for the left padding return computed_tokens - left_padding + def get_required_blocks(self, request: Request) -> tuple[int, int, bool]: + assert request.prompt_token_ids is not None + assert ( + request.sampling_params is not None and request.sampling_params.max_tokens is not None + ) + max_tokens = len(request.prompt_token_ids) + request.sampling_params.max_tokens + max_tokens = min(self.max_model_len, max_tokens) + + total_blocks = math.ceil(max_tokens / self.block_size) + + block_ids_per_kv_cache_group = self.kv_cache_manager.get_block_ids(request.request_id) + assert len(block_ids_per_kv_cache_group) == 1 + used_blocks = len(block_ids_per_kv_cache_group[0]) + + total_tokens = request.num_tokens + # the request will get a new block in the next iteration + needs_new_block_now = total_tokens < max_tokens and total_tokens % 64 == 0 - def get_blocks_required_for_decode_batch(self) -> int: + return total_blocks, used_blocks, needs_new_block_now + + def get_blocks_required_for_decode_batch(self, before_allocation: bool = True) -> int: """ Returns the number of blocks that the current decode batch needs to finish all requests. Reducing the number of available blocks below this number will cause deadlocks. """ - raise NotImplementedError() - + # Warning, depending on when this function is called, self.running + # might be out of sync with the worker's input_batch + required_blocks = 0 + for request in self.running: + total_blocks, used_blocks, needs_new_block_now = self.get_required_blocks(request) + required_blocks += ( + total_blocks - used_blocks + int(before_allocation and needs_new_block_now) + ) + + return required_blocks def schedule(self) -> "SchedulerOutput": """ From cf457b97a67387498b46e3c5f75a4aaa1fd5b509 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 22 May 2026 15:07:25 +0000 Subject: [PATCH 007/106] temporary tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../test_spyre_holdback_scheduler_steps.py | 673 ++++++++++++++++++ 1 file changed, 673 insertions(+) create mode 100644 tests/e2e/test_spyre_holdback_scheduler_steps.py diff --git a/tests/e2e/test_spyre_holdback_scheduler_steps.py b/tests/e2e/test_spyre_holdback_scheduler_steps.py new file mode 100644 index 000000000..7c7838916 --- /dev/null +++ b/tests/e2e/test_spyre_holdback_scheduler_steps.py @@ -0,0 +1,673 @@ +"""Verification of the holdback feature in the chunked prefill scheduler. + +This tests the relaxed constraint checking where requests are scheduled if +prefill constraints are satisfied (not future constraints). Requests that +would violate constraints during decode will be held back at that time. + +The two main constraints checked at prefill time are: +1. Max-context constraint: current tkv <= max_context_len +2. Volumetric constraint: current_max_tkv * batch_size <= max_batch_tkv_limit + +Run `python -m pytest tests/e2e/test_spyre_holdback_scheduler_steps.py`. +""" + +import pytest +from scheduling_utils import ( + validate_scheduler_steps, + create_request_for_scheduler_test, + random_prompt, +) +from spyre_util import ModelInfo + + +@pytest.mark.chunked_prefill +@pytest.mark.full_model +@pytest.mark.parametrize("max_num_seqs", [2]) +@pytest.mark.parametrize("max_model_len", [128]) +@pytest.mark.parametrize("max_num_batched_tokens", [128]) +@pytest.mark.parametrize("available_blocks", [None]) +def test_holdback_prefill_max_context_ok( + model: ModelInfo, + backend: str, + monkeypatch: pytest.MonkeyPatch, + set_random_seed, + max_num_seqs: int, + max_model_len: int, + max_num_batched_tokens: int, + available_blocks: int, +): + """Test that requests are scheduled when prefill max-context constraint is satisfied. + + With holdback feature, we only check if current tkv <= max_context_len, + not if future max_tokens would fit. This request would have been blocked + in the old scheduler but can now be scheduled. + + Configuration: + * max_num_seqs: 2 + * number of prompts: 2 + * 0: len = 49, max tokens = 20, step joining = 0 + * 1: len = 70, max tokens = 10, step joining = 0 + """ + + requests = [ + create_request_for_scheduler_test( + model=model, + request_id=0, + add_step=0, + max_tokens=20, + prompt=random_prompt(model, seed=0, length=49), + use_golden_token_injection=False, + generate_hf_results=True, + ), + create_request_for_scheduler_test( + model=model, + request_id=1, + add_step=0, + max_tokens=10, + prompt=random_prompt(model, seed=1, length=70), + use_golden_token_injection=False, + generate_hf_results=True, + ), + ] + + checked_steps = [ + { + "step": 0, + "tkv": 0, + "waiting": ["0", "1"], + "running": [], + "request_outputs": [], + "n_used_blocks": 0, + }, + { + # Prefill sequence 0 + "step": 1, + "tkv": 49, + "waiting": ["1"], + "running": ["0"], + "request_outputs": ["0"], + "n_used_blocks": 1, + }, + { + # Decode sequence 0 + "step": 2, + "tkv": 50, + "waiting": ["1"], + "running": ["0"], + "request_outputs": ["0"], + "n_used_blocks": 1, + }, + { + # With holdback: sequence 1 CAN be scheduled now because + # prefill tkv constraint is satisfied (70 <= 128) + # Old scheduler would block this because future max_tokens + # would exceed max_context_len (70 + 98 > 168) + "step": 3, + "tkv": 70, + "waiting": [], + "running": ["1", "0"], + "request_outputs": ["1"], + "n_used_blocks": 3, + }, + { + # Both sequences decode + "step": 4, + "tkv": 115, + "waiting": [], + "running": ["1", "0"], + "request_outputs": ["1", "0"], + "n_used_blocks": 3, + }, + { + # Sequence 1 finishes + "step": 12, + "tkv": 123, + "waiting": [], + "running": ["0"], + "request_outputs": ["1", "0"], + "finished_requests": ["1"], + "n_used_blocks": 1, + }, + { + # Decode sequence 0 + # We removed the padding block induced by sequence 1 + "step": 13, + "tkv": 60, + "waiting": [], + "running": ["0"], + "request_outputs": ["0"], + "n_used_blocks": 1, + }, + { + # Decode sequence 0 + # tkv is expanding to new block + "step": 18, + "tkv": 65, + "waiting": [], + "running": ["0"], + "request_outputs": ["0"], + "n_used_blocks": 2, + }, + { + # Sequence 0 finishes + "step": 21, + "tkv": 68, + "waiting": [], + "running": [], + "request_outputs": ["0"], + "finished_requests": ["0"], + "n_used_blocks": 0, + }, + { + # tkv should be cleared one step later + "step": 22, + "tkv": 0, + "waiting": [], + "running": [], + "request_outputs": [], + "finished_requests": [], + "n_used_blocks": 0, + }, + ] + + validate_scheduler_steps( + model=model, + backend=backend, + monkeypatch=monkeypatch, + requests=requests, + checked_steps=checked_steps, + max_num_seqs=max_num_seqs, + max_model_len=max_model_len, + available_blocks=available_blocks, + max_num_batched_tokens=max_num_batched_tokens, + ) + + +# NOTE Keeping the test draft just in case, but probably we won't need it at all +# and can remove the constraint along with the tests +# @pytest.mark.chunked_prefill +# @pytest.mark.full_model +# @pytest.mark.parametrize("max_num_seqs", [2]) +# @pytest.mark.parametrize("max_model_len", [128]) +# @pytest.mark.parametrize("max_num_batched_tokens", [128]) +# @pytest.mark.parametrize("available_blocks", [None]) +# def test_holdback_prefill_max_context_violated( +# model: ModelInfo, +# backend: str, +# monkeypatch: pytest.MonkeyPatch, +# set_random_seed, +# max_num_seqs: int, +# max_model_len: int, +# max_num_batched_tokens: int, +# available_blocks: int, +# ): +# """Test that requests are blocked when prefill max-context constraint is violated. + +# Even with holdback, if the prefill tkv would exceed max_context_len, +# the request cannot be scheduled. + +# Configuration: +# * max_num_seqs: 2 +# * number of prompts: 2 +# * 0: len = 60, max tokens = 10, step joining = 0 +# * 1: len = 75, max tokens = 5, step joining = 0 (exceeds max_model_len) +# """ + +# requests = [ +# create_request_for_scheduler_test( +# model=model, +# request_id=0, +# add_step=0, +# max_tokens=10, +# prompt=random_prompt(model, seed=0, length=60), +# use_golden_token_injection=False, +# generate_hf_results=False, +# ), +# create_request_for_scheduler_test( +# model=model, +# request_id=1, +# add_step=6, +# max_tokens=10, +# prompt=random_prompt(model, seed=1, length=75), +# use_golden_token_injection=False, +# generate_hf_results=False, +# ), +# ] + +# checked_steps = [ +# { +# "step": 0, +# "tkv": 0, +# "waiting": ["0"], +# "running": [], +# "request_outputs": [], +# "n_used_blocks": 0, +# }, +# { +# # Prefill sequence 0 +# "step": 1, +# "tkv": 60, +# "waiting": [], +# "running": ["0"], +# "request_outputs": ["0"], +# "n_used_blocks": 1, +# }, +# { +# # Decode 1 sequence 0 +# "step": 2, +# "tkv": 61, +# "waiting": [], +# "running": ["0"], +# "request_outputs": ["0"], +# "n_used_blocks": 1, +# }, +# { +# # Decode 5 sequence 0 +# # Request 1 joins the waiting queue +# "step": 6, +# "tkv": 65, +# "waiting": ["1"], +# "running": ["0"], +# "request_outputs": ["0"], +# "n_used_blocks": 2, +# }, +# { +# # Decode 6 sequence 0 +# # Sequence 1 CANNOT be scheduled because the padding-induced to +# # request 0 would shift its tkv beyond max_context_len +# # (64 + 66 = 130 > 128) +# "step": 7, +# "tkv": 75, +# "waiting": [], +# "running": ["1", "0"], +# "request_outputs": ["0"], +# "n_used_blocks": 2, +# }, +# # { +# # # Decode 6 sequence 0 +# # # Sequence 1 CANNOT be scheduled because the padding-induced to +# # # request 0 would shift its tkv beyond max_context_len +# # # (64 + 66 = 130 > 128) +# # "step": 7, +# # "tkv": 66, +# # "waiting": ["1"], +# # "running": ["0"], +# # "request_outputs": ["0"], +# # "n_used_blocks": 2, +# # }, +# { +# # Sequence 0 finishes +# "step": 10, +# "tkv": 70, +# "waiting": ["1"], +# "running": [], +# "request_outputs": ["0"], +# "finished_requests": ["0"], +# "n_used_blocks": 2, +# }, +# { +# # Prefill sequence 1 +# "step": 11, +# "tkv": 70, +# "waiting": ["1"], +# "running": [], +# "request_outputs": ["0"], +# "finished_requests": ["0"], +# "n_used_blocks": 2, +# }, +# # { +# # # Continue decoding sequence 0 +# # "step": 3, +# # "tkv": 66, +# # "waiting": ["1"], +# # "running": ["0"], +# # "request_outputs": ["0"], +# # "n_used_blocks": 1, +# # }, +# ] + +# validate_scheduler_steps( +# model=model, +# backend=backend, +# monkeypatch=monkeypatch, +# requests=requests, +# checked_steps=checked_steps, +# max_num_seqs=max_num_seqs, +# max_model_len=max_model_len, +# available_blocks=available_blocks, +# max_num_batched_tokens=max_num_batched_tokens, +# ) + + +@pytest.mark.chunked_prefill +@pytest.mark.full_model +@pytest.mark.parametrize("max_num_seqs", [4]) +@pytest.mark.parametrize("max_model_len", [128]) +@pytest.mark.parametrize("max_num_batched_tokens", [128]) +@pytest.mark.parametrize("available_blocks", [None]) +def test_holdback_prefill_volumetric_ok( + model: ModelInfo, + backend: str, + monkeypatch: pytest.MonkeyPatch, + set_random_seed, + max_num_seqs: int, + max_model_len: int, + max_num_batched_tokens: int, + available_blocks: int, +): + """Test that requests are scheduled when prefill volumetric constraint is satisfied. + + With holdback, we only check current_max_tkv * batch_size <= limit, + not future volumetric constraints. + + Configuration: + * max_num_seqs: 2 + * number of prompts: 2 + * 0: len = 64, max tokens = 100, step joining = 0 + * 1: len = 65, max tokens = 100, step joining = 0 + """ + + # Prefill volume: 2 * 65 = 130 (should pass) + # Old scheduler would check future: 2 * (65 + 100) = 330 (would fail with limit 200) + max_batch_tkv_limit = 256 + + requests = [ + create_request_for_scheduler_test( + model=model, + request_id=0, + add_step=0, + max_tokens=100, + prompt=random_prompt(model, seed=0, length=64), + use_golden_token_injection=False, + generate_hf_results=False, + ), + create_request_for_scheduler_test( + model=model, + request_id=1, + add_step=0, + max_tokens=100, + prompt=random_prompt(model, seed=1, length=65), + use_golden_token_injection=False, + generate_hf_results=False, + ), + ] + + checked_steps = [ + { + "step": 0, + "tkv": 0, + "waiting": ["0", "1"], + "running": [], + "request_outputs": [], + "n_used_blocks": 0, + }, + { + # Prefill sequence 0 + "step": 1, + "tkv": 64, + "waiting": ["1"], + "running": ["0"], + "request_outputs": ["0"], + "n_used_blocks": 1, + }, + { + # With holdback: sequence 1 CAN be scheduled + # Prefill volume: 2 * 65 = 130 <= 200 (passes) + # Old scheduler would block: future volume 2 * 165 = 330 > 200 + "step": 2, + "tkv": 65, + "waiting": [], + "running": ["1", "0"], + "request_outputs": ["1"], + "n_used_blocks": 3, + }, + { + # Both sequences decode + "step": 3, + "tkv": 66, + "waiting": [], + "running": ["1", "0"], + "request_outputs": ["1", "0"], + "n_used_blocks": 3, + }, + ] + + validate_scheduler_steps( + model=model, + backend=backend, + monkeypatch=monkeypatch, + requests=requests, + checked_steps=checked_steps, + max_num_seqs=max_num_seqs, + max_model_len=max_model_len, + available_blocks=available_blocks, + max_batch_tkv_limit=max_batch_tkv_limit, + max_num_batched_tokens=max_num_batched_tokens, + ) + + +@pytest.mark.chunked_prefill +@pytest.mark.full_model +@pytest.mark.parametrize("max_num_seqs", [2]) +@pytest.mark.parametrize("max_model_len", [2048]) +@pytest.mark.parametrize("max_num_batched_tokens", [128]) +@pytest.mark.parametrize("available_blocks", [None]) +def test_holdback_prefill_volumetric_violated( + model: ModelInfo, + backend: str, + monkeypatch: pytest.MonkeyPatch, + set_random_seed, + max_num_seqs: int, + max_model_len: int, + max_num_batched_tokens: int, + available_blocks: int, +): + """Test that requests are blocked when prefill volumetric constraint is violated. + + Even with holdback, if prefill volume exceeds limit, request cannot be scheduled. + + Configuration: + * max_num_seqs: 2 + * number of prompts: 2 + * 0: len = 64, max tokens = 10, step joining = 0 + * 1: len = 65, max tokens = 10, step joining = 0 + """ + + requests = [ + create_request_for_scheduler_test( + model=model, + request_id=0, + add_step=0, + max_tokens=10, + prompt=random_prompt(model, seed=0, length=64), + use_golden_token_injection=False, + generate_hf_results=False, + ), + create_request_for_scheduler_test( + model=model, + request_id=1, + add_step=0, + max_tokens=10, + prompt=random_prompt(model, seed=1, length=65), + use_golden_token_injection=False, + generate_hf_results=False, + ), + ] + + # Prefill volume would be: 2 * 65 = 130 + max_batch_tkv_limit = 129 # Just below the prefill requirement + + checked_steps = [ + { + "step": 0, + "tkv": 0, + "waiting": ["0", "1"], + "running": [], + "request_outputs": [], + "n_used_blocks": 0, + }, + { + # Prefill sequence 0 + "step": 1, + "tkv": 64, + "waiting": ["1"], + "running": ["0"], + "request_outputs": ["0"], + "n_used_blocks": 1, + }, + { + # Sequence 1 CANNOT be scheduled + # Prefill volume: 2 * 65 = 130 > 129 (limit) + "step": 2, + "tkv": 65, + "waiting": ["1"], + "running": ["0"], + "request_outputs": ["0"], + "n_used_blocks": 1, + }, + { + # Sequence 0 continues decoding + "step": 3, + "tkv": 66, + "waiting": ["1"], + "running": ["0"], + "request_outputs": ["0"], + "n_used_blocks": 1, + }, + ] + + validate_scheduler_steps( + model=model, + backend=backend, + monkeypatch=monkeypatch, + requests=requests, + checked_steps=checked_steps, + max_num_seqs=max_num_seqs, + max_model_len=max_model_len, + available_blocks=available_blocks, + max_batch_tkv_limit=max_batch_tkv_limit, + max_num_batched_tokens=max_num_batched_tokens, + ) + + +@pytest.mark.chunked_prefill +@pytest.mark.full_model +@pytest.mark.parametrize("max_num_seqs", [3]) +@pytest.mark.parametrize("max_model_len", [2048]) +@pytest.mark.parametrize("max_num_batched_tokens", [128]) +@pytest.mark.parametrize("available_blocks", [None]) +def test_holdback_prefill_multiple_requests_ok( + model: ModelInfo, + backend: str, + monkeypatch: pytest.MonkeyPatch, + set_random_seed, + max_num_seqs: int, + max_model_len: int, + max_num_batched_tokens: int, + available_blocks: int, +): + """Test scheduling multiple requests with prefill constraint checking. + + All requests can be scheduled as long as prefill constraints are satisfied, + regardless of future max_tokens requirements. + + Configuration: + * max_num_seqs: 3 + * number of prompts: 3 + * 0: len = 50, max tokens = 200, step joining = 0 + * 1: len = 60, max tokens = 200, step joining = 0 + * 2: len = 70, max tokens = 200, step joining = 0 + """ + + monkeypatch.setenv("SENDNN_INFERENCE_CP_INTERLEAVE_STEPS", "0") + + requests = [ + create_request_for_scheduler_test( + model=model, + request_id=0, + add_step=0, + max_tokens=200, + prompt=random_prompt(model, seed=0, length=50), + use_golden_token_injection=False, + generate_hf_results=False, + ), + create_request_for_scheduler_test( + model=model, + request_id=1, + add_step=0, + max_tokens=200, + prompt=random_prompt(model, seed=1, length=60), + use_golden_token_injection=False, + generate_hf_results=False, + ), + create_request_for_scheduler_test( + model=model, + request_id=2, + add_step=0, + max_tokens=200, + prompt=random_prompt(model, seed=2, length=70), + use_golden_token_injection=False, + generate_hf_results=False, + ), + ] + + checked_steps = [ + { + "step": 0, + "tkv": 0, + "waiting": ["0", "1", "2"], + "running": [], + "request_outputs": [], + "n_used_blocks": 0, + }, + { + # Prefill sequence 0 + "step": 1, + "tkv": 50, + "waiting": ["1", "2"], + "running": ["0"], + "request_outputs": ["0"], + "n_used_blocks": 1, + }, + { + # Prefill sequence 1 + # With holdback: can schedule even though future tokens would be large + "step": 2, + "tkv": 60, + "waiting": ["2"], + "running": ["1", "0"], + "request_outputs": ["1"], + "n_used_blocks": 2, + }, + { + # Prefill sequence 2 + # With holdback: can schedule even though future tokens would be large + "step": 3, + "tkv": 70, + "waiting": [], + "running": ["2", "1", "0"], + "request_outputs": ["2"], + "n_used_blocks": 3, + }, + { + # All three sequences decode together + "step": 4, + "tkv": 71, + "waiting": [], + "running": ["2", "1", "0"], + "request_outputs": ["2", "1", "0"], + "n_used_blocks": 3, + }, + ] + + validate_scheduler_steps( + model=model, + backend=backend, + monkeypatch=monkeypatch, + requests=requests, + checked_steps=checked_steps, + max_num_seqs=max_num_seqs, + max_model_len=max_model_len, + available_blocks=available_blocks, + max_num_batched_tokens=max_num_batched_tokens, + ) From 0baa5ecac9a71ed931ff54671234815d341c5aa9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 26 May 2026 22:45:32 +0000 Subject: [PATCH 008/106] correct volumetric violated prefill holdback test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../test_spyre_holdback_scheduler_steps.py | 636 +++++++++--------- tests/hf_cache.json | 80 +++ 2 files changed, 404 insertions(+), 312 deletions(-) diff --git a/tests/e2e/test_spyre_holdback_scheduler_steps.py b/tests/e2e/test_spyre_holdback_scheduler_steps.py index 7c7838916..5fe1de541 100644 --- a/tests/e2e/test_spyre_holdback_scheduler_steps.py +++ b/tests/e2e/test_spyre_holdback_scheduler_steps.py @@ -20,167 +20,167 @@ from spyre_util import ModelInfo -@pytest.mark.chunked_prefill -@pytest.mark.full_model -@pytest.mark.parametrize("max_num_seqs", [2]) -@pytest.mark.parametrize("max_model_len", [128]) -@pytest.mark.parametrize("max_num_batched_tokens", [128]) -@pytest.mark.parametrize("available_blocks", [None]) -def test_holdback_prefill_max_context_ok( - model: ModelInfo, - backend: str, - monkeypatch: pytest.MonkeyPatch, - set_random_seed, - max_num_seqs: int, - max_model_len: int, - max_num_batched_tokens: int, - available_blocks: int, -): - """Test that requests are scheduled when prefill max-context constraint is satisfied. +# @pytest.mark.chunked_prefill +# @pytest.mark.full_model +# @pytest.mark.parametrize("max_num_seqs", [2]) +# @pytest.mark.parametrize("max_model_len", [128]) +# @pytest.mark.parametrize("max_num_batched_tokens", [128]) +# @pytest.mark.parametrize("available_blocks", [None]) +# def test_holdback_prefill_max_context_ok( +# model: ModelInfo, +# backend: str, +# monkeypatch: pytest.MonkeyPatch, +# set_random_seed, +# max_num_seqs: int, +# max_model_len: int, +# max_num_batched_tokens: int, +# available_blocks: int, +# ): +# """Test that requests are scheduled when prefill max-context constraint is satisfied. - With holdback feature, we only check if current tkv <= max_context_len, - not if future max_tokens would fit. This request would have been blocked - in the old scheduler but can now be scheduled. +# With holdback feature, we only check if current tkv <= max_context_len, +# not if future max_tokens would fit. This request would have been blocked +# in the old scheduler but can now be scheduled. - Configuration: - * max_num_seqs: 2 - * number of prompts: 2 - * 0: len = 49, max tokens = 20, step joining = 0 - * 1: len = 70, max tokens = 10, step joining = 0 - """ +# Configuration: +# * max_num_seqs: 2 +# * number of prompts: 2 +# * 0: len = 49, max tokens = 20, step joining = 0 +# * 1: len = 70, max tokens = 10, step joining = 0 +# """ - requests = [ - create_request_for_scheduler_test( - model=model, - request_id=0, - add_step=0, - max_tokens=20, - prompt=random_prompt(model, seed=0, length=49), - use_golden_token_injection=False, - generate_hf_results=True, - ), - create_request_for_scheduler_test( - model=model, - request_id=1, - add_step=0, - max_tokens=10, - prompt=random_prompt(model, seed=1, length=70), - use_golden_token_injection=False, - generate_hf_results=True, - ), - ] +# requests = [ +# create_request_for_scheduler_test( +# model=model, +# request_id=0, +# add_step=0, +# max_tokens=20, +# prompt=random_prompt(model, seed=0, length=49), +# use_golden_token_injection=False, +# generate_hf_results=True, +# ), +# create_request_for_scheduler_test( +# model=model, +# request_id=1, +# add_step=0, +# max_tokens=10, +# prompt=random_prompt(model, seed=1, length=70), +# use_golden_token_injection=False, +# generate_hf_results=True, +# ), +# ] - checked_steps = [ - { - "step": 0, - "tkv": 0, - "waiting": ["0", "1"], - "running": [], - "request_outputs": [], - "n_used_blocks": 0, - }, - { - # Prefill sequence 0 - "step": 1, - "tkv": 49, - "waiting": ["1"], - "running": ["0"], - "request_outputs": ["0"], - "n_used_blocks": 1, - }, - { - # Decode sequence 0 - "step": 2, - "tkv": 50, - "waiting": ["1"], - "running": ["0"], - "request_outputs": ["0"], - "n_used_blocks": 1, - }, - { - # With holdback: sequence 1 CAN be scheduled now because - # prefill tkv constraint is satisfied (70 <= 128) - # Old scheduler would block this because future max_tokens - # would exceed max_context_len (70 + 98 > 168) - "step": 3, - "tkv": 70, - "waiting": [], - "running": ["1", "0"], - "request_outputs": ["1"], - "n_used_blocks": 3, - }, - { - # Both sequences decode - "step": 4, - "tkv": 115, - "waiting": [], - "running": ["1", "0"], - "request_outputs": ["1", "0"], - "n_used_blocks": 3, - }, - { - # Sequence 1 finishes - "step": 12, - "tkv": 123, - "waiting": [], - "running": ["0"], - "request_outputs": ["1", "0"], - "finished_requests": ["1"], - "n_used_blocks": 1, - }, - { - # Decode sequence 0 - # We removed the padding block induced by sequence 1 - "step": 13, - "tkv": 60, - "waiting": [], - "running": ["0"], - "request_outputs": ["0"], - "n_used_blocks": 1, - }, - { - # Decode sequence 0 - # tkv is expanding to new block - "step": 18, - "tkv": 65, - "waiting": [], - "running": ["0"], - "request_outputs": ["0"], - "n_used_blocks": 2, - }, - { - # Sequence 0 finishes - "step": 21, - "tkv": 68, - "waiting": [], - "running": [], - "request_outputs": ["0"], - "finished_requests": ["0"], - "n_used_blocks": 0, - }, - { - # tkv should be cleared one step later - "step": 22, - "tkv": 0, - "waiting": [], - "running": [], - "request_outputs": [], - "finished_requests": [], - "n_used_blocks": 0, - }, - ] +# checked_steps = [ +# { +# "step": 0, +# "tkv": 0, +# "waiting": ["0", "1"], +# "running": [], +# "request_outputs": [], +# "n_used_blocks": 0, +# }, +# { +# # Prefill sequence 0 +# "step": 1, +# "tkv": 49, +# "waiting": ["1"], +# "running": ["0"], +# "request_outputs": ["0"], +# "n_used_blocks": 1, +# }, +# { +# # Decode sequence 0 +# "step": 2, +# "tkv": 50, +# "waiting": ["1"], +# "running": ["0"], +# "request_outputs": ["0"], +# "n_used_blocks": 1, +# }, +# { +# # With holdback: sequence 1 CAN be scheduled now because +# # prefill tkv constraint is satisfied (70 <= 128) +# # Old scheduler would block this because future max_tokens +# # would exceed max_context_len (70 + 98 > 168) +# "step": 3, +# "tkv": 70, +# "waiting": [], +# "running": ["1", "0"], +# "request_outputs": ["1"], +# "n_used_blocks": 3, +# }, +# { +# # Both sequences decode +# "step": 4, +# "tkv": 115, +# "waiting": [], +# "running": ["1", "0"], +# "request_outputs": ["1", "0"], +# "n_used_blocks": 3, +# }, +# { +# # Sequence 1 finishes +# "step": 12, +# "tkv": 123, +# "waiting": [], +# "running": ["0"], +# "request_outputs": ["1", "0"], +# "finished_requests": ["1"], +# "n_used_blocks": 1, +# }, +# { +# # Decode sequence 0 +# # We removed the padding block induced by sequence 1 +# "step": 13, +# "tkv": 60, +# "waiting": [], +# "running": ["0"], +# "request_outputs": ["0"], +# "n_used_blocks": 1, +# }, +# { +# # Decode sequence 0 +# # tkv is expanding to new block +# "step": 18, +# "tkv": 65, +# "waiting": [], +# "running": ["0"], +# "request_outputs": ["0"], +# "n_used_blocks": 2, +# }, +# { +# # Sequence 0 finishes +# "step": 21, +# "tkv": 68, +# "waiting": [], +# "running": [], +# "request_outputs": ["0"], +# "finished_requests": ["0"], +# "n_used_blocks": 0, +# }, +# { +# # tkv should be cleared one step later +# "step": 22, +# "tkv": 0, +# "waiting": [], +# "running": [], +# "request_outputs": [], +# "finished_requests": [], +# "n_used_blocks": 0, +# }, +# ] - validate_scheduler_steps( - model=model, - backend=backend, - monkeypatch=monkeypatch, - requests=requests, - checked_steps=checked_steps, - max_num_seqs=max_num_seqs, - max_model_len=max_model_len, - available_blocks=available_blocks, - max_num_batched_tokens=max_num_batched_tokens, - ) +# validate_scheduler_steps( +# model=model, +# backend=backend, +# monkeypatch=monkeypatch, +# requests=requests, +# checked_steps=checked_steps, +# max_num_seqs=max_num_seqs, +# max_model_len=max_model_len, +# available_blocks=available_blocks, +# max_num_batched_tokens=max_num_batched_tokens, +# ) # NOTE Keeping the test draft just in case, but probably we won't need it at all @@ -343,7 +343,7 @@ def test_holdback_prefill_max_context_ok( @pytest.mark.full_model @pytest.mark.parametrize("max_num_seqs", [4]) @pytest.mark.parametrize("max_model_len", [128]) -@pytest.mark.parametrize("max_num_batched_tokens", [128]) +@pytest.mark.parametrize("max_num_batched_tokens", [256]) @pytest.mark.parametrize("available_blocks", [None]) def test_holdback_prefill_volumetric_ok( model: ModelInfo, @@ -363,12 +363,13 @@ def test_holdback_prefill_volumetric_ok( Configuration: * max_num_seqs: 2 * number of prompts: 2 - * 0: len = 64, max tokens = 100, step joining = 0 - * 1: len = 65, max tokens = 100, step joining = 0 + * 0: len = 15, max tokens = 60, step joining = 0 + * 1: len = 15, max tokens = 60, step joining = 0 + * 1: len = 66, max tokens = 60, step joining = 0 """ - # Prefill volume: 2 * 65 = 130 (should pass) - # Old scheduler would check future: 2 * (65 + 100) = 330 (would fail with limit 200) + # Volume right after prefill: 3 * 82 = 246 (should pass) + # Old scheduler would check future: TODO (would fail) max_batch_tkv_limit = 256 requests = [ @@ -376,19 +377,28 @@ def test_holdback_prefill_volumetric_ok( model=model, request_id=0, add_step=0, - max_tokens=100, - prompt=random_prompt(model, seed=0, length=64), + max_tokens=60, + prompt=random_prompt(model, seed=0, length=15), use_golden_token_injection=False, - generate_hf_results=False, + generate_hf_results=True, ), create_request_for_scheduler_test( model=model, request_id=1, add_step=0, - max_tokens=100, - prompt=random_prompt(model, seed=1, length=65), + max_tokens=60, + prompt=random_prompt(model, seed=1, length=15), + use_golden_token_injection=False, + generate_hf_results=True, + ), + create_request_for_scheduler_test( + model=model, + request_id=2, + add_step=0, + max_tokens=60, + prompt=random_prompt(model, seed=2, length=66), use_golden_token_injection=False, - generate_hf_results=False, + generate_hf_results=True, ), ] @@ -396,7 +406,7 @@ def test_holdback_prefill_volumetric_ok( { "step": 0, "tkv": 0, - "waiting": ["0", "1"], + "waiting": ["0", "1", "2"], "running": [], "request_outputs": [], "n_used_blocks": 0, @@ -404,135 +414,79 @@ def test_holdback_prefill_volumetric_ok( { # Prefill sequence 0 "step": 1, - "tkv": 64, - "waiting": ["1"], + "tkv": 15, + "waiting": ["1", "2"], "running": ["0"], "request_outputs": ["0"], "n_used_blocks": 1, }, { - # With holdback: sequence 1 CAN be scheduled - # Prefill volume: 2 * 65 = 130 <= 200 (passes) - # Old scheduler would block: future volume 2 * 165 = 330 > 200 + # Decode sequence 0 "step": 2, - "tkv": 65, - "waiting": [], + "tkv": 16, + "waiting": ["1", "2"], + "running": ["0"], + "request_outputs": ["0"], + "n_used_blocks": 1, + }, + { + # Prefill sequence 1 + "step": 3, + "tkv": 15, + "waiting": ["2"], "running": ["1", "0"], "request_outputs": ["1"], - "n_used_blocks": 3, + "n_used_blocks": 2, }, { - # Both sequences decode - "step": 3, - "tkv": 66, - "waiting": [], + # Decode sequences 0 and 1 + "step": 4, + "tkv": 17, + "waiting": ["2"], "running": ["1", "0"], "request_outputs": ["1", "0"], - "n_used_blocks": 3, + "n_used_blocks": 2, }, - ] - - validate_scheduler_steps( - model=model, - backend=backend, - monkeypatch=monkeypatch, - requests=requests, - checked_steps=checked_steps, - max_num_seqs=max_num_seqs, - max_model_len=max_model_len, - available_blocks=available_blocks, - max_batch_tkv_limit=max_batch_tkv_limit, - max_num_batched_tokens=max_num_batched_tokens, - ) - - -@pytest.mark.chunked_prefill -@pytest.mark.full_model -@pytest.mark.parametrize("max_num_seqs", [2]) -@pytest.mark.parametrize("max_model_len", [2048]) -@pytest.mark.parametrize("max_num_batched_tokens", [128]) -@pytest.mark.parametrize("available_blocks", [None]) -def test_holdback_prefill_volumetric_violated( - model: ModelInfo, - backend: str, - monkeypatch: pytest.MonkeyPatch, - set_random_seed, - max_num_seqs: int, - max_model_len: int, - max_num_batched_tokens: int, - available_blocks: int, -): - """Test that requests are blocked when prefill volumetric constraint is violated. - - Even with holdback, if prefill volume exceeds limit, request cannot be scheduled. - - Configuration: - * max_num_seqs: 2 - * number of prompts: 2 - * 0: len = 64, max tokens = 10, step joining = 0 - * 1: len = 65, max tokens = 10, step joining = 0 - """ - - requests = [ - create_request_for_scheduler_test( - model=model, - request_id=0, - add_step=0, - max_tokens=10, - prompt=random_prompt(model, seed=0, length=64), - use_golden_token_injection=False, - generate_hf_results=False, - ), - create_request_for_scheduler_test( - model=model, - request_id=1, - add_step=0, - max_tokens=10, - prompt=random_prompt(model, seed=1, length=65), - use_golden_token_injection=False, - generate_hf_results=False, - ), - ] - - # Prefill volume would be: 2 * 65 = 130 - max_batch_tkv_limit = 129 # Just below the prefill requirement - - checked_steps = [ { - "step": 0, - "tkv": 0, - "waiting": ["0", "1"], - "running": [], - "request_outputs": [], - "n_used_blocks": 0, + # Prefill sequence 2 + # With holdback: sequence 2 CAN be scheduled + # Decode volume: 3 * 82 = 246 <= 256 (passes) + # Old scheduler would block: future volume 3 * TODO = 330 > 256 + "step": 5, + "tkv": 66, + "waiting": [], + "running": ["2", "1", "0"], + "request_outputs": ["2"], + "n_used_blocks": 4, }, { - # Prefill sequence 0 - "step": 1, - "tkv": 64, - "waiting": ["1"], - "running": ["0"], - "request_outputs": ["0"], - "n_used_blocks": 1, + # Decode sequences 0, 1, and 2 + "step": 6, + "tkv": 82, + "waiting": [], + "running": ["2", "1", "0"], + "request_outputs": ["2", "1", "0"], + "n_used_blocks": 4, }, { - # Sequence 1 CANNOT be scheduled - # Prefill volume: 2 * 65 = 130 > 129 (limit) - "step": 2, - "tkv": 65, - "waiting": ["1"], - "running": ["0"], - "request_outputs": ["0"], - "n_used_blocks": 1, + # Decode sequences 0, 1, and 2 + "step": 9, + "tkv": 85, + "waiting": [], + "running": ["2", "1", "0"], + "request_outputs": ["2", "1", "0"], + "n_used_blocks": 4, }, { - # Sequence 0 continues decoding - "step": 3, - "tkv": 66, - "waiting": ["1"], - "running": ["0"], - "request_outputs": ["0"], - "n_used_blocks": 1, + # Decode sequences 0, 1, and 2 + # About to violate the volumetric constraint: 3 * 86 = 258 > 256 + # Holdback activates TODO + "step": 10, + "tkv": 86, + "waiting": [], + "running": ["2", "1", "0"], + "request_outputs": ["2", "1", "0"], + "n_used_blocks": 4, }, ] @@ -552,11 +506,11 @@ def test_holdback_prefill_volumetric_violated( @pytest.mark.chunked_prefill @pytest.mark.full_model -@pytest.mark.parametrize("max_num_seqs", [3]) +@pytest.mark.parametrize("max_num_seqs", [2]) @pytest.mark.parametrize("max_model_len", [2048]) @pytest.mark.parametrize("max_num_batched_tokens", [128]) @pytest.mark.parametrize("available_blocks", [None]) -def test_holdback_prefill_multiple_requests_ok( +def test_holdback_prefill_volumetric_violated( model: ModelInfo, backend: str, monkeypatch: pytest.MonkeyPatch, @@ -566,48 +520,48 @@ def test_holdback_prefill_multiple_requests_ok( max_num_batched_tokens: int, available_blocks: int, ): - """Test scheduling multiple requests with prefill constraint checking. + """Test that requests are blocked when prefill volumetric constraint is violated. - All requests can be scheduled as long as prefill constraints are satisfied, - regardless of future max_tokens requirements. + Even with holdback, if prefill volume exceeds limit, request cannot be scheduled. Configuration: - * max_num_seqs: 3 - * number of prompts: 3 - * 0: len = 50, max tokens = 200, step joining = 0 - * 1: len = 60, max tokens = 200, step joining = 0 - * 2: len = 70, max tokens = 200, step joining = 0 + * max_num_seqs: 2 + * number of prompts: 2 + * 0: len = 25, max tokens = 7, step joining = 0 + * 1: len = 25, max tokens = 6, step joining = 0 + * 1: len = 66, max tokens = 3, step joining = 0 """ - monkeypatch.setenv("SENDNN_INFERENCE_CP_INTERLEAVE_STEPS", "0") + # Volume right after prefill: 3 * 89 = 267 (fails) + max_batch_tkv_limit = 256 requests = [ create_request_for_scheduler_test( model=model, request_id=0, add_step=0, - max_tokens=200, - prompt=random_prompt(model, seed=0, length=50), + max_tokens=7, + prompt=random_prompt(model, seed=0, length=25), use_golden_token_injection=False, - generate_hf_results=False, + generate_hf_results=True, ), create_request_for_scheduler_test( model=model, request_id=1, add_step=0, - max_tokens=200, - prompt=random_prompt(model, seed=1, length=60), + max_tokens=6, + prompt=random_prompt(model, seed=1, length=25), use_golden_token_injection=False, - generate_hf_results=False, + generate_hf_results=True, ), create_request_for_scheduler_test( model=model, request_id=2, add_step=0, - max_tokens=200, - prompt=random_prompt(model, seed=2, length=70), + max_tokens=3, + prompt=random_prompt(model, seed=2, length=66), use_golden_token_injection=False, - generate_hf_results=False, + generate_hf_results=True, ), ] @@ -623,40 +577,97 @@ def test_holdback_prefill_multiple_requests_ok( { # Prefill sequence 0 "step": 1, - "tkv": 50, + "tkv": 25, "waiting": ["1", "2"], "running": ["0"], "request_outputs": ["0"], "n_used_blocks": 1, }, { - # Prefill sequence 1 - # With holdback: can schedule even though future tokens would be large + # Decode sequence 0 "step": 2, - "tkv": 60, + "tkv": 26, + "waiting": ["1", "2"], + "running": ["0"], + "request_outputs": ["0"], + "n_used_blocks": 1, + }, + { + # Prefill sequence 1 + "step": 3, + "tkv": 25, "waiting": ["2"], "running": ["1", "0"], "request_outputs": ["1"], "n_used_blocks": 2, }, + { + # Decode sequences 0 and 1 + "step": 4, + "tkv": 27, + "waiting": ["2"], + "running": ["1", "0"], + "request_outputs": ["1", "0"], + "n_used_blocks": 2, + }, + { + # Decode sequences 0 and 1 + # Cannot prefill sequence 2 + # tkv would be 3 * (28 + 64) = 276 > 256 + "step": 5, + "tkv": 28, + "waiting": ["2"], + "running": ["1", "0"], + "request_outputs": ["1", "0"], + "n_used_blocks": 2, + }, + { + # Sequences 0 and 1 both finish + "step": 8, + "tkv": 31, + "waiting": ["2"], + "running": [], + "request_outputs": ["1", "0"], + "finished_requests": ["1", "0"], + "n_used_blocks": 0, + }, { # Prefill sequence 2 - # With holdback: can schedule even though future tokens would be large - "step": 3, - "tkv": 70, + "step": 9, + "tkv": 66, "waiting": [], - "running": ["2", "1", "0"], + "running": ["2"], "request_outputs": ["2"], - "n_used_blocks": 3, + "n_used_blocks": 2, }, { - # All three sequences decode together - "step": 4, - "tkv": 71, + # Decode sequence 2 + "step": 10, + "tkv": 67, "waiting": [], - "running": ["2", "1", "0"], - "request_outputs": ["2", "1", "0"], - "n_used_blocks": 3, + "running": ["2"], + "request_outputs": ["2"], + "n_used_blocks": 2, + }, + { + # Sequence 2 finishes + "step": 11, + "tkv": 68, + "waiting": [], + "running": [], + "request_outputs": ["2"], + "finished_requests": ["2"], + "n_used_blocks": 0, + }, + { + # tkv should be cleared one step later + "step": 12, + "tkv": 0, + "waiting": [], + "running": [], + "request_outputs": [], + "finished_requests": [], + "n_used_blocks": 0, }, ] @@ -669,5 +680,6 @@ def test_holdback_prefill_multiple_requests_ok( max_num_seqs=max_num_seqs, max_model_len=max_model_len, available_blocks=available_blocks, + max_batch_tkv_limit=max_batch_tkv_limit, max_num_batched_tokens=max_num_batched_tokens, ) diff --git a/tests/hf_cache.json b/tests/hf_cache.json index 9bd016aeb..fd69acf20 100644 --- a/tests/hf_cache.json +++ b/tests/hf_cache.json @@ -603,6 +603,86 @@ "tokens": [ "cs", "ymbol" ], "logprobs": [ -2.9405062198638916, -6.677685737609863 ] } + }, + "__tokens__41505_37255_20672_12727_25130_19903_38525_14909_23426_28674": { + "100": { + "text": "\n\n### 2.3.2.2.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.1.", + "token_ids": [ 203, 203, 1482, 225, 36, 32, 37, 32, 36, 32, 36, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32 ], + "tokens": [ "\n", "\n", "###", " ", "2", ".", "3", ".", "2", ".", "2", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", "." ], + "logprobs": [ -3.0465893745422363, -1.2976518869400024, -3.163188934326172, -0.49196842312812805, -3.1609890460968018, -0.2816971242427826, -1.7173802852630615, -0.3782590329647064, -1.3585926294326782, -0.7624582052230835, -1.5881094932556152, -0.8354007005691528, -1.4533395767211914, -0.7736920118331909, -0.9893547892570496, -0.614208996295929, -1.0189998149871826, -0.46721890568733215, -0.668842613697052, -0.22856344282627106, -0.7188294529914856, -0.08179545402526855, -0.703545093536377, -0.0471746064722538, -0.42062729597091675, -0.0369371622800827, -0.23578061163425446, -0.03197914734482765, -0.1281307190656662, -0.028535017743706703, -0.07559774816036224, -0.02715468779206276, -0.05364929884672165, -0.02528834156692028, -0.03956437483429909, -0.022716999053955078, -0.028716547414660454, -0.018719438463449478, -0.022590087726712227, -0.01657642237842083, -0.018984990194439888, -0.015271991491317749, -0.015639042481780052, -0.01199073065072298, -0.013088975101709366, -0.009404288604855537, -0.011163041926920414, -0.008704103529453278, -0.00997180212289095, -0.008108317852020264, -0.009046755731105804, -0.006760344374924898, -0.008206576108932495, -0.0057660676538944244, -0.0075477901846170425, -0.00497437035664916, -0.00690976157784462, -0.004145481623709202, -0.006284238304942846, -0.0038340408354997635, -0.005858393386006355, -0.0037401027511805296, -0.005646114237606525, -0.0034246151335537434, -0.005494966637343168, -0.003046873025596142, -0.005232567898929119, -0.002723914571106434, -0.0048736585304141045, -0.0024694681633263826, -0.00458131218329072, -0.0023883646354079247, -0.0045217410661280155, -0.00239205127581954, -0.004494090098887682, -0.00235958443954587, -0.004242824390530586, -0.0023642226587980986, -0.003988764248788357, -0.0024396199733018875, -0.0038497161585837603, -0.0023695745039731264, -0.0036468682810664177, -0.0022651508916169405, -0.0034722534473985434, -0.0022508781403303146, -0.0032436635810881853, -0.0022613448090851307, -0.003047229489311576, -0.0022889384999871254, -0.0029463237151503563, -0.0023240242153406143, -0.00292671169154346, -0.002363271312788129, -0.0029105464927852154, -0.002446160651743412, -0.00289152842015028, -0.002533799270167947, -0.002845883136615157, -0.0025805288460105658 ] + } + }, + "__tokens__6605_41653_37541_12537_24352_22093_32027_38767_4614_1394": { + "100": { + "text": "\n\n## 1. \u6211\u4eec\u5c06\u4f7f\u7528\u4e00\u4e2a\u51fd\u6570\u6765\u5b9e\u73b0\u4e00\u4e2a\u51fd\u6570\n\n```python\nimport numpy as np\n\ndef get_surface_quickstart(self, x):\n \"\"\"\n \u6211\u4eec\u5c06\u4f7f\u7528\u4e00\u4e2a\u51fd\u6570\u6765\u5b9e\u73b0\u4e00\u4e2a\u51fd\u6570\n \"\"\"\n # \u6211\u4eec\u5c06\u4f7f\u7528\u4e00\u4e2a\u51fd\u6570\u6765\u5b9e\u73b0\u4e00\u4e2a\u51fd\u6570\n # \u6211\u4eec\u5c06\u4f7f\u7528\u4e00\u4e2a\u51fd\u6570\n # \u6211\u4eec\u5c06\u4f7f\u7528\u4e00\u4e2a\u51fd\u6570\n # \u6211\u4eec\u5c06\u4f7f\u7528\u4e00\u4e2a\u51fd\u6570\n # \u6211\u4eec\u5c06\u4f7f\u7528\u4e00\u4e2a\u51fd\u6570\n # ", + "token_ids": [ 203, 203, 433, 225, 35, 32, 225, 8517, 6817, 4690, 4577, 10205, 5642, 12656, 4577, 10205, 203, 203, 914, 2958, 203, 465, 6436, 619, 2065, 203, 203, 589, 622, 81, 12537, 81, 24352, 26, 784, 30, 816, 711, 284, 1524, 284, 225, 8517, 6817, 4690, 4577, 10205, 5642, 12656, 4577, 10205, 284, 1524, 284, 588, 225, 8517, 6817, 4690, 4577, 10205, 5642, 12656, 4577, 10205, 284, 588, 225, 8517, 6817, 4690, 4577, 10205, 284, 588, 225, 8517, 6817, 4690, 4577, 10205, 284, 588, 225, 8517, 6817, 4690, 4577, 10205, 284, 588, 225, 8517, 6817, 4690, 4577, 10205, 284, 588, 225 ], + "tokens": [ "\n", "\n", "##", " ", "1", ".", " ", "\u6211\u4eec", "\u5c06", "\u4f7f\u7528", "\u4e00\u4e2a", "\u51fd\u6570", "\u6765", "\u5b9e\u73b0", "\u4e00\u4e2a", "\u51fd\u6570", "\n", "\n", "```", "python", "\n", "import", " numpy", " as", " np", "\n", "\n", "def", " get", "_", "surface", "_", "quickstart", "(", "self", ",", " x", "):", "\n ", " \"\"\"", "\n ", " ", "\u6211\u4eec", "\u5c06", "\u4f7f\u7528", "\u4e00\u4e2a", "\u51fd\u6570", "\u6765", "\u5b9e\u73b0", "\u4e00\u4e2a", "\u51fd\u6570", "\n ", " \"\"\"", "\n ", " #", " ", "\u6211\u4eec", "\u5c06", "\u4f7f\u7528", "\u4e00\u4e2a", "\u51fd\u6570", "\u6765", "\u5b9e\u73b0", "\u4e00\u4e2a", "\u51fd\u6570", "\n ", " #", " ", "\u6211\u4eec", "\u5c06", "\u4f7f\u7528", "\u4e00\u4e2a", "\u51fd\u6570", "\n ", " #", " ", "\u6211\u4eec", "\u5c06", "\u4f7f\u7528", "\u4e00\u4e2a", "\u51fd\u6570", "\n ", " #", " ", "\u6211\u4eec", "\u5c06", "\u4f7f\u7528", "\u4e00\u4e2a", "\u51fd\u6570", "\n ", " #", " ", "\u6211\u4eec", "\u5c06", "\u4f7f\u7528", "\u4e00\u4e2a", "\u51fd\u6570", "\n ", " #", " " ], + "logprobs": [ -2.199794054031372, -1.4375017881393433, -2.640138864517212, -0.6885342001914978, -2.824195146560669, -0.46407780051231384, -1.4570093154907227, -3.6046652793884277, -2.542109966278076, -2.5875916481018066, -3.2663722038269043, -2.930560350418091, -1.3067225217819214, -3.0429837703704834, -2.713665008544922, -3.103975296020508, -1.153260350227356, -0.3038812577724457, -1.068111538887024, -1.6085222959518433, -0.02793041616678238, -1.364532232284546, -2.649442434310913, -0.018922410905361176, -0.005025137215852737, -0.03027990460395813, -0.7889770269393921, -0.8153836727142334, -2.613900899887085, -0.5962694883346558, -2.65573787689209, -0.1977139413356781, -3.1533970832824707, -0.4371451735496521, -2.7665152549743652, -0.5860720872879028, -2.881225347518921, -0.9788865447044373, -0.17679451406002045, -0.8672720193862915, -0.3661006689071655, -1.9630924463272095, -2.7465708255767822, -0.5309916138648987, -0.6546719670295715, -0.18320831656455994, -0.017766030505299568, -0.052741639316082, -0.11376317590475082, -0.029810355976223946, -0.03958098962903023, -0.41218724846839905, -0.7074546217918396, -0.1504538208246231, -1.7077534198760986, -1.035139560699463, -1.4580460786819458, -0.5916604399681091, -0.873860239982605, -0.4423682987689972, -0.11110946536064148, -0.6441695690155029, -0.18500685691833496, -0.178217351436615, -0.07237415760755539, -0.2097354233264923, -1.635819673538208, -0.6866054534912109, -1.791818618774414, -1.0729011297225952, -0.7346925139427185, -0.7114763259887695, -0.2690492868423462, -0.7088493704795837, -1.512905478477478, -0.4910634160041809, -0.6958803534507751, -0.326402485370636, -0.18015918135643005, -0.4319242835044861, -0.14119885861873627, -0.039465367794036865, -1.3307157754898071, -0.30063366889953613, -0.1942070722579956, -0.1208154633641243, -0.06644336134195328, -0.18892322480678558, -0.05879750847816467, -0.03347078338265419, -1.0705735683441162, -0.19743774831295013, -0.12167882919311523, -0.07732052356004715, -0.04663863405585289, -0.13200309872627258, -0.03398321196436882, -0.0282746572047472, -0.8110079765319824, -0.17539414763450623 ] + } + }, + "__tokens__46991_46587_2780_4172_41066_36174_32918_15146_29783_29825_28567_7785_21168_19343_35537_48897_46664_26747_21866_13185_1766_1349_22850_15653_18679_43833_25842_27550_11606_1173_15982_6719_25079_49087_33152_8938_43920_39162_36097_44560_37497_38817_17389_48216_47279_7923_37061_35151_22679_26068_24085_45457_24617_40871_17396_43394_44222_22660_27904_45236_35575_23918_10903_15958_34385_8163": { + "100": { + "text": "\" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \"", + "token_ids": [ 20, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313 ], + "tokens": [ "\"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"" ], + "logprobs": [ -2.742318630218506, -3.0037925243377686, -3.0905137062072754, -0.9414255023002625, -0.42565011978149414, -0.24709248542785645, -0.2613930404186249, -0.46429160237312317, -0.6351388096809387, -0.771374523639679, -0.8507365584373474, -0.8271970152854919, -0.7441465854644775, -0.6466454267501831, -0.5404894948005676, -0.4509280323982239, -0.3863682746887207, -0.34253576397895813, -0.31299206614494324, -0.29213395714759827, -0.28407159447669983, -0.28230226039886475, -0.28447502851486206, -0.28669920563697815, -0.2918243110179901, -0.29709166288375854, -0.2959142327308655, -0.29070067405700684, -0.2894592881202698, -0.2962619662284851, -0.3065527677536011, -0.3125804364681244, -0.3088591396808624, -0.3101733922958374, -0.3163264989852905, -0.3249780833721161, -0.33249661326408386, -0.3332236111164093, -0.3324698507785797, -0.3347475528717041, -0.3412342667579651, -0.35016486048698425, -0.35427409410476685, -0.3498973846435547, -0.3447745740413666, -0.3392011821269989, -0.33481502532958984, -0.33510249853134155, -0.32848823070526123, -0.31656643748283386, -0.3112860321998596, -0.30591878294944763, -0.2999021112918854, -0.2919565439224243, -0.2862960398197174, -0.2773458957672119, -0.2669422924518585, -0.260267436504364, -0.25428062677383423, -0.2512851059436798, -0.25447481870651245, -0.2535083293914795, -0.2456202507019043, -0.23807093501091003, -0.22996939718723297, -0.22121354937553406, -0.21426820755004883, -0.20618228614330292, -0.1977836787700653, -0.18722476065158844, -0.17892034351825714, -0.17277148365974426, -0.1669711172580719, -0.1632366180419922, -0.16045381128787994, -0.15518233180046082, -0.14549972116947174, -0.13666827976703644, -0.13417178392410278, -0.13527832925319672, -0.1371733844280243, -0.13487361371517181, -0.12779873609542847, -0.11836154013872147, -0.10898654907941818, -0.10413984209299088, -0.09933914989233017, -0.0895027220249176, -0.07327315211296082, -0.052362505346536636, -0.0443415641784668, -0.05291973054409027, -0.05461844429373741, -0.04578790068626404, -0.06402725726366043, -0.15075816214084625, -0.24608442187309265, -0.11664518713951111, -0.059976473450660706, -0.055951036512851715 ] + }, + "60": { + "text": "\" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \"", + "token_ids": [ 20, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313, 313 ], + "tokens": [ "\"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"" ], + "logprobs": [ -2.742318630218506, -3.0037925243377686, -3.0905137062072754, -0.9414255023002625, -0.42565011978149414, -0.24709248542785645, -0.2613930404186249, -0.46429160237312317, -0.6351388096809387, -0.771374523639679, -0.8507365584373474, -0.8271970152854919, -0.7441465854644775, -0.6466454267501831, -0.5404894948005676, -0.4509280323982239, -0.3863682746887207, -0.34253576397895813, -0.31299206614494324, -0.29213395714759827, -0.28407159447669983, -0.28230226039886475, -0.28447502851486206, -0.28669920563697815, -0.2918243110179901, -0.29709166288375854, -0.2959142327308655, -0.29070067405700684, -0.2894592881202698, -0.2962619662284851, -0.3065527677536011, -0.3125804364681244, -0.3088591396808624, -0.3101733922958374, -0.3163264989852905, -0.3249780833721161, -0.33249661326408386, -0.3332236111164093, -0.3324698507785797, -0.3347475528717041, -0.3412342667579651, -0.35016486048698425, -0.35427409410476685, -0.3498973846435547, -0.3447745740413666, -0.3392011821269989, -0.33481502532958984, -0.33510249853134155, -0.32848823070526123, -0.31656643748283386, -0.3112860321998596, -0.30591878294944763, -0.2999021112918854, -0.2919565439224243, -0.2862960398197174, -0.2773458957672119, -0.2669422924518585, -0.260267436504364, -0.25428062677383423, -0.2512851059436798 ] + }, + "2": { + "text": "\" \"", + "token_ids": [ 20, 313 ], + "tokens": [ "\"", " \"" ], + "logprobs": [ -2.742318630218506, -3.0037922859191895 ] + } + }, + "__tokens__41505_37255_20672_12727_25130_19903_38525_14909_23426_28674_44635_24806_13853_37149_30394": { + "100": { + "text": "\n#define CHD_PRE_COMP_TYPE_V1 0x00000001\n#define CHD_PRE_COMP_TYPE_V2 0x00000002\n#define CHD_PRE_COMP_TYPE_V3 0x00000004\n#define CHD_PRE_COMP_TYPE_V4 0x00000008", + "token_ids": [ 203, 21, 1234, 5557, 54, 81, 2883, 81, 5762, 81, 2319, 81, 72, 35, 225, 34, 106, 34, 34, 34, 34, 34, 34, 34, 35, 203, 21, 1234, 5557, 54, 81, 2883, 81, 5762, 81, 2319, 81, 72, 36, 225, 34, 106, 34, 34, 34, 34, 34, 34, 34, 36, 203, 21, 1234, 5557, 54, 81, 2883, 81, 5762, 81, 2319, 81, 72, 37, 225, 34, 106, 34, 34, 34, 34, 34, 34, 34, 38, 203, 21, 1234, 5557, 54, 81, 2883, 81, 5762, 81, 2319, 81, 72, 38, 225, 34, 106, 34, 34, 34, 34, 34, 34, 34, 42 ], + "tokens": [ "\n", "#", "define", " CH", "D", "_", "PRE", "_", "COMP", "_", "TYPE", "_", "V", "1", " ", "0", "x", "0", "0", "0", "0", "0", "0", "0", "1", "\n", "#", "define", " CH", "D", "_", "PRE", "_", "COMP", "_", "TYPE", "_", "V", "2", " ", "0", "x", "0", "0", "0", "0", "0", "0", "0", "2", "\n", "#", "define", " CH", "D", "_", "PRE", "_", "COMP", "_", "TYPE", "_", "V", "3", " ", "0", "x", "0", "0", "0", "0", "0", "0", "0", "4", "\n", "#", "define", " CH", "D", "_", "PRE", "_", "COMP", "_", "TYPE", "_", "V", "4", " ", "0", "x", "0", "0", "0", "0", "0", "0", "0", "8" ], + "logprobs": [ -4.148627281188965, -1.9023098945617676, -0.5024883151054382, -3.5166265964508057, -1.1848623752593994, -0.44509002566337585, -4.0265913009643555, -2.1063308715820312, -4.351094722747803, -1.4263076782226562, -3.4072775840759277, -0.46827957034111023, -4.332708835601807, -1.7475272417068481, -1.3255337476730347, -0.8408962488174438, -0.37473323941230774, -0.7385716438293457, -0.47072669863700867, -0.18981407582759857, -0.36916476488113403, -0.1256936490535736, -0.17075695097446442, -0.17706118524074554, -0.4836171567440033, -0.5464026927947998, -0.2326568365097046, -0.04812277853488922, -0.03288242593407631, -0.004547967109829187, -0.0012984187342226505, -0.10411255806684494, -0.003740933956578374, -0.04705769941210747, -0.0035643160808831453, -0.18792352080345154, -0.002217930741608143, -0.0816296860575676, -0.16547933220863342, -0.033810585737228394, -0.03572966530919075, -0.018084051087498665, -0.017124062404036522, -0.010007210075855255, -0.016362886875867844, -0.028098611161112785, -0.033767715096473694, -0.023985574021935463, -0.06287175416946411, -0.06378459185361862, -0.014418580569326878, -0.08894632756710052, -0.02926105447113514, -0.028644727542996407, -0.0009793015196919441, -0.0002786724944598973, -0.16831068694591522, -0.0022170981392264366, -0.03330787271261215, -0.002313439268618822, -0.041456446051597595, -0.0002917817619163543, -0.15613055229187012, -0.03973784297704697, -0.008251740597188473, -0.005037354305386543, -0.0033077073749154806, -0.006288502831012011, -0.0025725625455379486, -0.003108076984062791, -0.005129747558385134, -0.00485254218801856, -0.005447543226182461, -0.026349563151597977, -0.12644031643867493, -0.008795449510216713, -0.10167873650789261, -0.01642856001853943, -0.015500782988965511, -0.00047994061606004834, -8.880697714630514e-05, -0.07469187676906586, -0.00070296844933182, -0.01724405214190483, -0.0008412636234425008, -0.014866752550005913, -9.274052717955783e-05, -0.07821717858314514, -0.03616850823163986, -0.004844712559133768, -0.001525192055851221, -0.002193070948123932, -0.0034923297353088856, -0.0012396040838211775, -0.0017492959741503, -0.0022299441043287516, -0.0032256022095680237, -0.003230711678043008, -0.029909281060099602, -0.011151960119605064 ] + } + }, + "__tokens__6605_41653_37541_12537_24352_22093_32027_38767_4614_1394_41079_21271_37467_104_21892": { + "100": { + "text": " 3.0.1 2019-01-01 14:30:00 2019-01-01 14:30:00 2019-01-01 14:30:00 2019-01-01 14:30:00 2019-01-01 14", + "token_ids": [ 225, 37, 32, 34, 32, 35, 225, 36, 34, 35, 43, 31, 34, 35, 31, 34, 35, 225, 35, 38, 44, 37, 34, 44, 34, 34, 225, 36, 34, 35, 43, 31, 34, 35, 31, 34, 35, 225, 35, 38, 44, 37, 34, 44, 34, 34, 225, 36, 34, 35, 43, 31, 34, 35, 31, 34, 35, 225, 35, 38, 44, 37, 34, 44, 34, 34, 225, 36, 34, 35, 43, 31, 34, 35, 31, 34, 35, 225, 35, 38, 44, 37, 34, 44, 34, 34, 225, 36, 34, 35, 43, 31, 34, 35, 31, 34, 35, 225, 35, 38 ], + "tokens": [ " ", "3", ".", "0", ".", "1", " ", "2", "0", "1", "9", "-", "0", "1", "-", "0", "1", " ", "1", "4", ":", "3", "0", ":", "0", "0", " ", "2", "0", "1", "9", "-", "0", "1", "-", "0", "1", " ", "1", "4", ":", "3", "0", ":", "0", "0", " ", "2", "0", "1", "9", "-", "0", "1", "-", "0", "1", " ", "1", "4", ":", "3", "0", ":", "0", "0", " ", "2", "0", "1", "9", "-", "0", "1", "-", "0", "1", " ", "1", "4", ":", "3", "0", ":", "0", "0", " ", "2", "0", "1", "9", "-", "0", "1", "-", "0", "1", " ", "1", "4" ], + "logprobs": [ -2.393317222595215, -2.2741634845733643, -0.5255809426307678, -1.454126000404358, -1.2284232378005981, -1.6320126056671143, -1.6596072912216187, -1.7695199251174927, -0.41706928610801697, -0.41384077072143555, -1.5353412628173828, -0.541284441947937, -0.3393153250217438, -1.750977873802185, -0.018852457404136658, -1.096146583557129, -1.5840284824371338, -0.4938022792339325, -0.7846470475196838, -2.0936717987060547, -0.39854463934898376, -1.633031964302063, -1.9479743242263794, -0.3339877128601074, -1.1566519737243652, -1.060813546180725, -0.9951281547546387, -2.2271621227264404, -0.1960439682006836, -0.1647973656654358, -0.09352599829435349, -0.3445914685726166, -0.014470632188022137, -0.1592160165309906, -0.001257463125512004, -0.010343162342905998, -0.10037858039140701, -0.05559161305427551, -0.07893017679452896, -0.14610762894153595, -0.016683464869856834, -0.0467216856777668, -0.025850284844636917, -0.030198249965906143, -0.043905384838581085, -0.08884358406066895, -1.2519291639328003, -0.5712020397186279, -0.013030261732637882, -0.04303565248847008, -0.04684977978467941, -0.029735142365098, -0.0020624573808163404, -0.016849223524332047, -0.0018703126115724444, -0.005184776149690151, -0.020718814805150032, -0.040916938334703445, -0.009638085961341858, -0.057150136679410934, -0.0024661386851221323, -0.010109765455126762, -0.004823832772672176, -0.005025255959481001, -0.0070279063656926155, -0.009060695767402649, -0.7706539630889893, -0.1780472695827484, -0.001004786929115653, -0.01082677487283945, -0.0162928719073534, -0.010154962539672852, -0.0010992205934599042, -0.010214435867965221, -0.0004924515378661454, -0.0026298719458281994, -0.013098269701004028, -0.016590023413300514, -0.0034246151335537434, -0.020235290750861168, -0.0010221739066764712, -0.002377542434260249, -0.0017613149248063564, -0.0027014450170099735, -0.00199333718046546, -0.003969291225075722, -0.3933572769165039, -0.04973369836807251, -0.00044967554276809096, -0.005075783468782902, -0.008561218157410622, -0.005584354046732187, -0.0008662762120366096, -0.0056595089845359325, -0.0003175231395289302, -0.0014316319720819592, -0.009665240533649921, -0.0076560406014323235, -0.0022322041913866997, -0.013572084717452526 ] + } + }, + "__tokens__41505_37255_20672_12727_25130_19903_38525_14909_23426_28674_44635_24806_13853_37149_30394_12313_44715_48305_39823_44343_15245_35872_44179_33619_23207": { + "60": { + "text": "\n\n### 2.1.2.3.1.1.1.1.1.2.1.2.2.3.1.3.1.4.1.5.1.6.1.7.1.8.1.9.", + "token_ids": [ 203, 203, 1482, 225, 36, 32, 35, 32, 36, 32, 37, 32, 35, 32, 35, 32, 35, 32, 35, 32, 35, 32, 36, 32, 35, 32, 36, 32, 36, 32, 37, 32, 35, 32, 37, 32, 35, 32, 38, 32, 35, 32, 39, 32, 35, 32, 40, 32, 35, 32, 41, 32, 35, 32, 42, 32, 35, 32, 43, 32 ], + "tokens": [ "\n", "\n", "###", " ", "2", ".", "1", ".", "2", ".", "3", ".", "1", ".", "1", ".", "1", ".", "1", ".", "1", ".", "2", ".", "1", ".", "2", ".", "2", ".", "3", ".", "1", ".", "3", ".", "1", ".", "4", ".", "1", ".", "5", ".", "1", ".", "6", ".", "1", ".", "7", ".", "1", ".", "8", ".", "1", ".", "9", "." ], + "logprobs": [ -4.0449628829956055, -1.4918135404586792, -2.295555591583252, -1.1818832159042358, -2.2455224990844727, -0.1181272566318512, -1.7651878595352173, -0.559072732925415, -1.477505087852478, -0.441697895526886, -1.8287240266799927, -0.6087567806243896, -1.709043264389038, -0.26909974217414856, -1.10608971118927, -0.28878718614578247, -1.0484988689422607, -0.2591942846775055, -0.5988566875457764, -0.18785445392131805, -0.5857826471328735, -0.10326796770095825, -0.720543384552002, -0.05641297996044159, -0.48668402433395386, -0.14942237734794617, -0.839197039604187, -0.07304967194795609, -0.551040530204773, -0.07871787250041962, -1.0794485807418823, -0.08451627939939499, -1.239895224571228, -0.052667904645204544, -0.6149333715438843, -0.02352249063551426, -0.5463603734970093, -0.019887220114469528, -0.2718304395675659, -0.03291968256235123, -0.3950521945953369, -0.01926017552614212, -0.7182247638702393, -0.023578496649861336, -0.3283310532569885, -0.019009435549378395, -0.7635195851325989, -0.016471009701490402, -0.16172637045383453, -0.011872699484229088, -0.8374265432357788, -0.016047487035393715, -0.16847145557403564, -0.012802669778466225, -0.30135419964790344, -0.014644630253314972, -0.1540326178073883, -0.02250838465988636, -0.18895263969898224, -0.015077201649546623 ] + }, + "7": { + "text": "\n\n### 2.1", + "token_ids": [ 203, 203, 1482, 225, 36, 32, 35 ], + "tokens": [ "\n", "\n", "###", " ", "2", ".", "1" ], + "logprobs": [ -4.0449628829956055, -1.4918135404586792, -2.295555591583252, -1.1818832159042358, -2.2455224990844727, -0.1181272566318512, -1.7651878595352173 ] + } + }, + "__tokens__6605_41653_37541_12537_24352_22093_32027_38767_4614_1394_41079_21271_37467_104_21892_35465_11244_46461_44307_1504_1251_26611_46161_18737_10647": { + "60": { + "text": "\n STDERR SimulatorAlnode graphsWise ADVANCED_GRAPH_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME_NAME", + "token_ids": [ 11244, 46461, 44307, 1504, 1251, 26611, 46161, 6988, 72, 48408, 81, 17326, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474, 81, 2474 ], + "tokens": [ "\n ", " STDERR", " Simulator", "Al", "node", " graphs", "Wise", " AD", "V", "ANCED", "_", "GRAPH", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME", "_", "NAME" ], + "logprobs": [ -2.7261359691619873, -1.2988193035125732, -0.39854758977890015, -0.07328899204730988, -0.11031155288219452, -0.2497384399175644, -0.16161644458770752, -4.789858818054199, -1.4781008958816528, -0.7245264649391174, -2.005262613296509, -3.8723177909851074, -1.1007609367370605, -3.0019583702087402, -2.3345510959625244, -3.587618827819824, -2.288602590560913, -2.764263391494751, -0.8430484533309937, -0.8916659951210022, -0.2132130116224289, -0.40108659863471985, -0.1460379958152771, -0.2521405518054962, -0.12207847088575363, -0.18518200516700745, -0.11607035994529724, -0.13806740939617157, -0.11627934128046036, -0.10058428347110748, -0.08686070144176483, -0.0747978538274765, -0.06742114573717117, -0.05440983176231384, -0.05565362796187401, -0.04014641419053078, -0.04035218060016632, -0.032005585730075836, -0.02955831214785576, -0.027025436982512474, -0.02407553419470787, -0.02166801318526268, -0.022037271410226822, -0.017526863142848015, -0.020062964409589767, -0.015297937206923962, -0.017838051542639732, -0.013740241527557373, -0.015049018897116184, -0.012525111436843872, -0.012721222825348377, -0.011666995473206043, -0.012215563096106052, -0.01085000578314066, -0.010758611373603344, -0.010379500687122345, -0.009193351492285728, -0.009988208301365376, -0.00830033142119646, -0.009442431852221489 ] + }, + "6": { + "text": "\n STDERR SimulatorAlnode graphs", + "token_ids": [ 11244, 46461, 44307, 1504, 1251, 26611 ], + "tokens": [ "\n ", " STDERR", " Simulator", "Al", "node", " graphs" ], + "logprobs": [ -2.7261359691619873, -1.2988193035125732, -0.39854758977890015, -0.07328899204730988, -0.11031155288219452, -0.2497384399175644 ] + } } } }, From 9777fc688f4610ea4f73cd67089e52f670a88ddb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 27 May 2026 19:30:58 +0000 Subject: [PATCH 009/106] remove max-context-len constraint entirely MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 17 +- .../test_spyre_holdback_scheduler_steps.py | 319 ------------------ 2 files changed, 2 insertions(+), 334 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 53f482c48..beb8ed501 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -489,7 +489,6 @@ def _satisfies_last_chunk_constraints(self, request: Request) -> bool: in the decode batch, and if all the other spyre-related conditions are satisfied.""" decoding_requests = [r for r in self.running if r not in self.ongoing_prefills] - max_context_len = self.model_config.max_model_len # check that there is space in the current decode batch num_running = len(decoding_requests) @@ -504,28 +503,16 @@ def _satisfies_last_chunk_constraints(self, request: Request) -> bool: n_blocks = math.floor(max(self.tkv, prompt_len) / self.block_size) new_req_tkv = n_blocks * self.block_size + prompt_len % self.block_size - # check that no tkv exceeds max_context_len (immediate constraint check) - cond2 = new_req_tkv <= max_context_len - # check cond2 for all other sequences in the current decode batch - for req in decoding_requests: - # current tkv of the (left aligned) decode sequence - dec_req_tkv = n_blocks * self.block_size + req.num_computed_tokens % self.block_size - cond2_current = dec_req_tkv <= max_context_len - cond2 = cond2 and cond2_current - # early exiting loop if violated 2nd condition - if not cond2: - return False - # check that batch size x tkv is smaller than the max supported number # Note: using max_tkv is a conservative upper bound here. For the # optimal check we need model runner to return per sequence tkvs - cond3 = lambda: self.check_batch_tkv_limit( + cond2 = lambda: self.check_batch_tkv_limit( request=request, new_req_tkv=new_req_tkv, running=decoding_requests, ) - return cond1 and cond2 and cond3() + return cond1 and cond2() def _has_scheduling_priority(self, request): decoding_requests = [r for r in self.running if r not in self.ongoing_prefills] diff --git a/tests/e2e/test_spyre_holdback_scheduler_steps.py b/tests/e2e/test_spyre_holdback_scheduler_steps.py index 5fe1de541..beafcaf0f 100644 --- a/tests/e2e/test_spyre_holdback_scheduler_steps.py +++ b/tests/e2e/test_spyre_holdback_scheduler_steps.py @@ -20,325 +20,6 @@ from spyre_util import ModelInfo -# @pytest.mark.chunked_prefill -# @pytest.mark.full_model -# @pytest.mark.parametrize("max_num_seqs", [2]) -# @pytest.mark.parametrize("max_model_len", [128]) -# @pytest.mark.parametrize("max_num_batched_tokens", [128]) -# @pytest.mark.parametrize("available_blocks", [None]) -# def test_holdback_prefill_max_context_ok( -# model: ModelInfo, -# backend: str, -# monkeypatch: pytest.MonkeyPatch, -# set_random_seed, -# max_num_seqs: int, -# max_model_len: int, -# max_num_batched_tokens: int, -# available_blocks: int, -# ): -# """Test that requests are scheduled when prefill max-context constraint is satisfied. - -# With holdback feature, we only check if current tkv <= max_context_len, -# not if future max_tokens would fit. This request would have been blocked -# in the old scheduler but can now be scheduled. - -# Configuration: -# * max_num_seqs: 2 -# * number of prompts: 2 -# * 0: len = 49, max tokens = 20, step joining = 0 -# * 1: len = 70, max tokens = 10, step joining = 0 -# """ - -# requests = [ -# create_request_for_scheduler_test( -# model=model, -# request_id=0, -# add_step=0, -# max_tokens=20, -# prompt=random_prompt(model, seed=0, length=49), -# use_golden_token_injection=False, -# generate_hf_results=True, -# ), -# create_request_for_scheduler_test( -# model=model, -# request_id=1, -# add_step=0, -# max_tokens=10, -# prompt=random_prompt(model, seed=1, length=70), -# use_golden_token_injection=False, -# generate_hf_results=True, -# ), -# ] - -# checked_steps = [ -# { -# "step": 0, -# "tkv": 0, -# "waiting": ["0", "1"], -# "running": [], -# "request_outputs": [], -# "n_used_blocks": 0, -# }, -# { -# # Prefill sequence 0 -# "step": 1, -# "tkv": 49, -# "waiting": ["1"], -# "running": ["0"], -# "request_outputs": ["0"], -# "n_used_blocks": 1, -# }, -# { -# # Decode sequence 0 -# "step": 2, -# "tkv": 50, -# "waiting": ["1"], -# "running": ["0"], -# "request_outputs": ["0"], -# "n_used_blocks": 1, -# }, -# { -# # With holdback: sequence 1 CAN be scheduled now because -# # prefill tkv constraint is satisfied (70 <= 128) -# # Old scheduler would block this because future max_tokens -# # would exceed max_context_len (70 + 98 > 168) -# "step": 3, -# "tkv": 70, -# "waiting": [], -# "running": ["1", "0"], -# "request_outputs": ["1"], -# "n_used_blocks": 3, -# }, -# { -# # Both sequences decode -# "step": 4, -# "tkv": 115, -# "waiting": [], -# "running": ["1", "0"], -# "request_outputs": ["1", "0"], -# "n_used_blocks": 3, -# }, -# { -# # Sequence 1 finishes -# "step": 12, -# "tkv": 123, -# "waiting": [], -# "running": ["0"], -# "request_outputs": ["1", "0"], -# "finished_requests": ["1"], -# "n_used_blocks": 1, -# }, -# { -# # Decode sequence 0 -# # We removed the padding block induced by sequence 1 -# "step": 13, -# "tkv": 60, -# "waiting": [], -# "running": ["0"], -# "request_outputs": ["0"], -# "n_used_blocks": 1, -# }, -# { -# # Decode sequence 0 -# # tkv is expanding to new block -# "step": 18, -# "tkv": 65, -# "waiting": [], -# "running": ["0"], -# "request_outputs": ["0"], -# "n_used_blocks": 2, -# }, -# { -# # Sequence 0 finishes -# "step": 21, -# "tkv": 68, -# "waiting": [], -# "running": [], -# "request_outputs": ["0"], -# "finished_requests": ["0"], -# "n_used_blocks": 0, -# }, -# { -# # tkv should be cleared one step later -# "step": 22, -# "tkv": 0, -# "waiting": [], -# "running": [], -# "request_outputs": [], -# "finished_requests": [], -# "n_used_blocks": 0, -# }, -# ] - -# validate_scheduler_steps( -# model=model, -# backend=backend, -# monkeypatch=monkeypatch, -# requests=requests, -# checked_steps=checked_steps, -# max_num_seqs=max_num_seqs, -# max_model_len=max_model_len, -# available_blocks=available_blocks, -# max_num_batched_tokens=max_num_batched_tokens, -# ) - - -# NOTE Keeping the test draft just in case, but probably we won't need it at all -# and can remove the constraint along with the tests -# @pytest.mark.chunked_prefill -# @pytest.mark.full_model -# @pytest.mark.parametrize("max_num_seqs", [2]) -# @pytest.mark.parametrize("max_model_len", [128]) -# @pytest.mark.parametrize("max_num_batched_tokens", [128]) -# @pytest.mark.parametrize("available_blocks", [None]) -# def test_holdback_prefill_max_context_violated( -# model: ModelInfo, -# backend: str, -# monkeypatch: pytest.MonkeyPatch, -# set_random_seed, -# max_num_seqs: int, -# max_model_len: int, -# max_num_batched_tokens: int, -# available_blocks: int, -# ): -# """Test that requests are blocked when prefill max-context constraint is violated. - -# Even with holdback, if the prefill tkv would exceed max_context_len, -# the request cannot be scheduled. - -# Configuration: -# * max_num_seqs: 2 -# * number of prompts: 2 -# * 0: len = 60, max tokens = 10, step joining = 0 -# * 1: len = 75, max tokens = 5, step joining = 0 (exceeds max_model_len) -# """ - -# requests = [ -# create_request_for_scheduler_test( -# model=model, -# request_id=0, -# add_step=0, -# max_tokens=10, -# prompt=random_prompt(model, seed=0, length=60), -# use_golden_token_injection=False, -# generate_hf_results=False, -# ), -# create_request_for_scheduler_test( -# model=model, -# request_id=1, -# add_step=6, -# max_tokens=10, -# prompt=random_prompt(model, seed=1, length=75), -# use_golden_token_injection=False, -# generate_hf_results=False, -# ), -# ] - -# checked_steps = [ -# { -# "step": 0, -# "tkv": 0, -# "waiting": ["0"], -# "running": [], -# "request_outputs": [], -# "n_used_blocks": 0, -# }, -# { -# # Prefill sequence 0 -# "step": 1, -# "tkv": 60, -# "waiting": [], -# "running": ["0"], -# "request_outputs": ["0"], -# "n_used_blocks": 1, -# }, -# { -# # Decode 1 sequence 0 -# "step": 2, -# "tkv": 61, -# "waiting": [], -# "running": ["0"], -# "request_outputs": ["0"], -# "n_used_blocks": 1, -# }, -# { -# # Decode 5 sequence 0 -# # Request 1 joins the waiting queue -# "step": 6, -# "tkv": 65, -# "waiting": ["1"], -# "running": ["0"], -# "request_outputs": ["0"], -# "n_used_blocks": 2, -# }, -# { -# # Decode 6 sequence 0 -# # Sequence 1 CANNOT be scheduled because the padding-induced to -# # request 0 would shift its tkv beyond max_context_len -# # (64 + 66 = 130 > 128) -# "step": 7, -# "tkv": 75, -# "waiting": [], -# "running": ["1", "0"], -# "request_outputs": ["0"], -# "n_used_blocks": 2, -# }, -# # { -# # # Decode 6 sequence 0 -# # # Sequence 1 CANNOT be scheduled because the padding-induced to -# # # request 0 would shift its tkv beyond max_context_len -# # # (64 + 66 = 130 > 128) -# # "step": 7, -# # "tkv": 66, -# # "waiting": ["1"], -# # "running": ["0"], -# # "request_outputs": ["0"], -# # "n_used_blocks": 2, -# # }, -# { -# # Sequence 0 finishes -# "step": 10, -# "tkv": 70, -# "waiting": ["1"], -# "running": [], -# "request_outputs": ["0"], -# "finished_requests": ["0"], -# "n_used_blocks": 2, -# }, -# { -# # Prefill sequence 1 -# "step": 11, -# "tkv": 70, -# "waiting": ["1"], -# "running": [], -# "request_outputs": ["0"], -# "finished_requests": ["0"], -# "n_used_blocks": 2, -# }, -# # { -# # # Continue decoding sequence 0 -# # "step": 3, -# # "tkv": 66, -# # "waiting": ["1"], -# # "running": ["0"], -# # "request_outputs": ["0"], -# # "n_used_blocks": 1, -# # }, -# ] - -# validate_scheduler_steps( -# model=model, -# backend=backend, -# monkeypatch=monkeypatch, -# requests=requests, -# checked_steps=checked_steps, -# max_num_seqs=max_num_seqs, -# max_model_len=max_model_len, -# available_blocks=available_blocks, -# max_num_batched_tokens=max_num_batched_tokens, -# ) - - @pytest.mark.chunked_prefill @pytest.mark.full_model @pytest.mark.parametrize("max_num_seqs", [4]) From 86a4aa9264daf6926c0eb5fb9176652cc432365e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Thu, 28 May 2026 21:08:26 +0000 Subject: [PATCH 010/106] scheduler predict next tkv MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 47 +++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index beb8ed501..0c9a09ad4 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -561,6 +561,53 @@ def check_batch_tkv_limit(self, request: Request, new_req_tkv: int, running) -> current_batch_tkv = batch_size * current_max_tkv return current_batch_tkv <= self.max_batch_tkv_limit + def predict_next_decode_tkv(self, running_requests: list[Request]) -> int: + """ + Predicts the TKV after the next decode step for a given batch of running + requests. + + This method replicates the TKV calculation logic from the model runner's + _prepare_decode method, accounting for: + - Block alignment (left-padding to make batch rectangular) + - The next token that will be generated (+1) + - Maximum TKV across all requests in the batch + + Args: + running_requests: List of Request objects currently in the decode batch + + Returns: + The predicted TKV value after the next decode step + """ + if not running_requests: + return 0 + + # Step 1: Find the maximum number of blocks across all requests + max_n_blocks = 0 + for request in running_requests: + block_ids_per_kv_cache_group = self.kv_cache_manager.get_block_ids(request.request_id) + assert len(block_ids_per_kv_cache_group) == 1 + num_blocks = len(block_ids_per_kv_cache_group[0]) + max_n_blocks = max(max_n_blocks, num_blocks) + + # Step 2: Calculate TKV for each request and find the maximum + max_tkv = 0 + for request in running_requests: + # Get the number of blocks for this request + block_ids_per_kv_cache_group = self.kv_cache_manager.get_block_ids(request.request_id) + num_blocks = len(block_ids_per_kv_cache_group[0]) + + # Calculate left padding blocks needed for alignment + left_pad_blocks_count = max_n_blocks - num_blocks + left_padding = left_pad_blocks_count * self.block_size + + # Calculate TKV for this request (including the next token) + req_tkv = left_padding + request.num_computed_tokens + 1 + + # Track the maximum TKV + max_tkv = max(max_tkv, req_tkv) + + return max_tkv + def finish_requests( self, request_ids: Union[str, Iterable[str], None], From 88a7c4d1ad2de45e21e1d7e05c45457e9788471f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 29 May 2026 10:13:15 +0200 Subject: [PATCH 011/106] rename holdback to pausing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/worker/spyre_model_runner.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index 8da51623d..a779c66f1 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -1449,21 +1449,21 @@ def _update_batch(self, scheduler_output: SchedulerOutput): req_data = scheduler_output.scheduled_cached_reqs # Synchronize input_batch with scheduler output: remove requests - # that are not in scheduler output. This handles hold back cases - # where scheduler temporarily removes requests from running queue + # that are not in scheduler output. This handles pausing of decode + # requests where scheduler temporarily removes them from running queue scheduled_req_ids = set(req_data.req_ids) current_batch_req_ids = set(self.input_batch.req_id_to_index.keys()) # Find requests that are in input_batch but not in scheduler output (held back) - heldback_req_ids = current_batch_req_ids - scheduled_req_ids - for req_id in heldback_req_ids: + paused_req_ids = current_batch_req_ids - scheduled_req_ids + for req_id in paused_req_ids: # Only remove if it's not a finished request (finished requests are handled separately) if req_id not in (scheduler_output.finished_req_ids or []): logger.info("Removing held back request %s from input_batch", req_id) self.input_batch.remove_request(req_id) # Find requests that are in scheduler output but not in input_batch - # (restored from holding back) + # (restore from pausing) restored_req_ids = scheduled_req_ids - current_batch_req_ids for req_id in restored_req_ids: # Add back the request that was held back From ed9b1ab4aad3607b1f9e325fc952b8f994a10f0f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 29 May 2026 11:42:30 +0200 Subject: [PATCH 012/106] rename handle pausing and resuming MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 73 +++++++++++++++++++++++++++ 1 file changed, 73 insertions(+) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 0c9a09ad4..ce6d2dbb5 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -193,6 +193,10 @@ def __init__(self, *args, **kwargs) -> None: # keep a list to be able to batch prefills in the future. self.ongoing_prefills: list[Request] = [] + # Track requests that were temporarily paused from decoding due to + # batch TKV constraint and moved back to waiting queue + self.paused_decoding_requests: list[Request] = [] + # Prefills interleaving: if the feature flag is set, prefill operations # are interleaved with a decode step. This allows to minimize currently # decoding requests @@ -393,6 +397,9 @@ def schedule(self) -> "SchedulerOutput": self.previous_step_was_prefill = False running_holdback = [] + if not self.previous_step_was_prefill: + self._handle_decode_requests_pausing() + # delegate to super of SpyreScheduler: base V1 Scheduler outputs = super(SpyreScheduler, self).schedule() @@ -561,6 +568,72 @@ def check_batch_tkv_limit(self, request: Request, new_req_tkv: int, running) -> current_batch_tkv = batch_size * current_max_tkv return current_batch_tkv <= self.max_batch_tkv_limit + def _can_decode_all_requests(self, decoding_requests: list[Request]) -> bool: + """ + Check if all decoding requests can be decoded in the next step without + violating the max batch TKV limit. + """ + if not decoding_requests: + return True + + next_predicted_tkv = self.predict_next_decode_tkv(decoding_requests) + + # the tkv should never get beyond max_model_len + assert next_predicted_tkv <= self.max_model_len + + # check batch tkv limit: batch_size * predicted_tkv must not exceed limit + batch_size = len(decoding_requests) + predicted_batch_tkv = batch_size * next_predicted_tkv + + return predicted_batch_tkv <= self.max_batch_tkv_limit + + def _handle_decode_requests_pausing(self) -> None: + """ + Manage pausing and resuming of decode requests based on batch TKV constraints. + + This method: + 1. Pauses requests with the fewest decoded tokens when batch TKV limit is exceeded + 2. Resumes previously paused requests (oldest first) when capacity is available + """ + decoding_requests = [r for r in self.running if r not in self.ongoing_prefills] + + had_to_remove = False + initial_had_requests = len(decoding_requests) > 0 + + # If we can't decode all requests due to batch TKV limits, iteratively + # remove requests with the fewest decoded tokens and pause them until + # the remaining batch fits within constraints + while not self._can_decode_all_requests(decoding_requests): + had_to_remove = True + # Remove the request with the fewest decoded tokens + # Decoded tokens = num_computed_tokens - num_prompt_tokens + request_to_remove = min( + decoding_requests, key=lambda r: r.num_computed_tokens - r.num_prompt_tokens + ) + decoding_requests.remove(request_to_remove) + self.running.remove(request_to_remove) + self.paused_decoding_requests.append(request_to_remove) + + # It shouldn't be possible to remove all requests if we started with some + assert not initial_had_requests or len(decoding_requests) > 0 + + # If we didn't have to remove any requests, try to add back previously + # paused requests (oldest first) as long as they fit within constraints + if not had_to_remove: + while self.paused_decoding_requests: + # Try adding the oldest paused request (first in list) + request_to_add = self.paused_decoding_requests[0] + test_requests = decoding_requests + [request_to_add] + + if self._can_decode_all_requests(test_requests): + # Can add this request back + self.paused_decoding_requests.pop(0) + self.running.append(request_to_add) + decoding_requests.append(request_to_add) + else: + # Can't add any more requests + break + def predict_next_decode_tkv(self, running_requests: list[Request]) -> int: """ Predicts the TKV after the next decode step for a given batch of running From e39dc3290c271471f8ce6804ed7f69efc1aaaf99 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 29 May 2026 14:36:21 +0200 Subject: [PATCH 013/106] bugfix: wrong tkv when expanding to new block MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index ce6d2dbb5..e2f1020da 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -655,11 +655,19 @@ def predict_next_decode_tkv(self, running_requests: list[Request]) -> int: return 0 # Step 1: Find the maximum number of blocks across all requests + # Account for requests that will need a new block after the next token max_n_blocks = 0 for request in running_requests: block_ids_per_kv_cache_group = self.kv_cache_manager.get_block_ids(request.request_id) assert len(block_ids_per_kv_cache_group) == 1 num_blocks = len(block_ids_per_kv_cache_group[0]) + + # Check if the next token will require a new block + next_token_count = request.num_computed_tokens + 1 + if next_token_count % self.block_size == 1: + # The next token will fill the current block and require a new one + num_blocks += 1 + max_n_blocks = max(max_n_blocks, num_blocks) # Step 2: Calculate TKV for each request and find the maximum @@ -669,6 +677,12 @@ def predict_next_decode_tkv(self, running_requests: list[Request]) -> int: block_ids_per_kv_cache_group = self.kv_cache_manager.get_block_ids(request.request_id) num_blocks = len(block_ids_per_kv_cache_group[0]) + # Check if the next token will require a new block + next_token_count = request.num_computed_tokens + 1 + if next_token_count % self.block_size == 1: + # The next token will fill the current block and require a new one + num_blocks += 1 + # Calculate left padding blocks needed for alignment left_pad_blocks_count = max_n_blocks - num_blocks left_padding = left_pad_blocks_count * self.block_size From 117f7f269adb68490e949a1bab12d74061cd75b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 29 May 2026 15:04:23 +0200 Subject: [PATCH 014/106] continued: rename holdback to pausing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/worker/spyre_model_runner.py | 8 ++++---- ...ps.py => test_spyre_decode_pausing_scheduler_steps.py} | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) rename tests/e2e/{test_spyre_holdback_scheduler_steps.py => test_spyre_decode_pausing_scheduler_steps.py} (97%) diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index a779c66f1..9a761d883 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -1454,21 +1454,21 @@ def _update_batch(self, scheduler_output: SchedulerOutput): scheduled_req_ids = set(req_data.req_ids) current_batch_req_ids = set(self.input_batch.req_id_to_index.keys()) - # Find requests that are in input_batch but not in scheduler output (held back) + # Find requests that are in input_batch but not in scheduler output (paused) paused_req_ids = current_batch_req_ids - scheduled_req_ids for req_id in paused_req_ids: # Only remove if it's not a finished request (finished requests are handled separately) if req_id not in (scheduler_output.finished_req_ids or []): - logger.info("Removing held back request %s from input_batch", req_id) + logger.info("Removing paused request %s from input_batch", req_id) self.input_batch.remove_request(req_id) # Find requests that are in scheduler output but not in input_batch # (restore from pausing) restored_req_ids = scheduled_req_ids - current_batch_req_ids for req_id in restored_req_ids: - # Add back the request that was held back + # Add back the request that was paused if req_id in self.requests: - logger.info("Restoring held back request %s to input_batch", req_id) + logger.info("Restoring paused request %s to input_batch", req_id) req_state = self.requests[req_id] self.input_batch.add_request(req_state) diff --git a/tests/e2e/test_spyre_holdback_scheduler_steps.py b/tests/e2e/test_spyre_decode_pausing_scheduler_steps.py similarity index 97% rename from tests/e2e/test_spyre_holdback_scheduler_steps.py rename to tests/e2e/test_spyre_decode_pausing_scheduler_steps.py index beafcaf0f..c3b269518 100644 --- a/tests/e2e/test_spyre_holdback_scheduler_steps.py +++ b/tests/e2e/test_spyre_decode_pausing_scheduler_steps.py @@ -1,8 +1,8 @@ -"""Verification of the holdback feature in the chunked prefill scheduler. +"""Verification of the decoding requests pausing feature in the chunked prefill scheduler. This tests the relaxed constraint checking where requests are scheduled if prefill constraints are satisfied (not future constraints). Requests that -would violate constraints during decode will be held back at that time. +would violate constraints during decode will be paused at that time. The two main constraints checked at prefill time are: 1. Max-context constraint: current tkv <= max_context_len @@ -26,7 +26,7 @@ @pytest.mark.parametrize("max_model_len", [128]) @pytest.mark.parametrize("max_num_batched_tokens", [256]) @pytest.mark.parametrize("available_blocks", [None]) -def test_holdback_prefill_volumetric_ok( +def test_pausing_prefill_volumetric_ok( model: ModelInfo, backend: str, monkeypatch: pytest.MonkeyPatch, @@ -191,7 +191,7 @@ def test_holdback_prefill_volumetric_ok( @pytest.mark.parametrize("max_model_len", [2048]) @pytest.mark.parametrize("max_num_batched_tokens", [128]) @pytest.mark.parametrize("available_blocks", [None]) -def test_holdback_prefill_volumetric_violated( +def test_pausing_prefill_volumetric_violated( model: ModelInfo, backend: str, monkeypatch: pytest.MonkeyPatch, From cd281785169976a2e2654d663b24a6e9a7c76bc9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 29 May 2026 14:34:11 +0000 Subject: [PATCH 015/106] complete decode pausing tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- ...est_spyre_decode_pause_scheduler_steps.py} | 90 ++++++++++++++++--- tests/hf_cache.json | 48 ++++++++++ 2 files changed, 124 insertions(+), 14 deletions(-) rename tests/e2e/{test_spyre_decode_pausing_scheduler_steps.py => test_spyre_decode_pause_scheduler_steps.py} (80%) diff --git a/tests/e2e/test_spyre_decode_pausing_scheduler_steps.py b/tests/e2e/test_spyre_decode_pause_scheduler_steps.py similarity index 80% rename from tests/e2e/test_spyre_decode_pausing_scheduler_steps.py rename to tests/e2e/test_spyre_decode_pause_scheduler_steps.py index c3b269518..865395422 100644 --- a/tests/e2e/test_spyre_decode_pausing_scheduler_steps.py +++ b/tests/e2e/test_spyre_decode_pause_scheduler_steps.py @@ -26,7 +26,7 @@ @pytest.mark.parametrize("max_model_len", [128]) @pytest.mark.parametrize("max_num_batched_tokens", [256]) @pytest.mark.parametrize("available_blocks", [None]) -def test_pausing_prefill_volumetric_ok( +def test_volumetric_decode_pausing( model: ModelInfo, backend: str, monkeypatch: pytest.MonkeyPatch, @@ -44,9 +44,9 @@ def test_pausing_prefill_volumetric_ok( Configuration: * max_num_seqs: 2 * number of prompts: 2 - * 0: len = 15, max tokens = 60, step joining = 0 - * 1: len = 15, max tokens = 60, step joining = 0 - * 1: len = 66, max tokens = 60, step joining = 0 + * 0: len = 15, max tokens = 11, step joining = 0 + * 1: len = 15, max tokens = 13, step joining = 0 + * 2: len = 66, max tokens = 10, step joining = 0 """ # Volume right after prefill: 3 * 82 = 246 (should pass) @@ -58,7 +58,7 @@ def test_pausing_prefill_volumetric_ok( model=model, request_id=0, add_step=0, - max_tokens=60, + max_tokens=11, prompt=random_prompt(model, seed=0, length=15), use_golden_token_injection=False, generate_hf_results=True, @@ -67,7 +67,7 @@ def test_pausing_prefill_volumetric_ok( model=model, request_id=1, add_step=0, - max_tokens=60, + max_tokens=13, prompt=random_prompt(model, seed=1, length=15), use_golden_token_injection=False, generate_hf_results=True, @@ -76,7 +76,7 @@ def test_pausing_prefill_volumetric_ok( model=model, request_id=2, add_step=0, - max_tokens=60, + max_tokens=10, prompt=random_prompt(model, seed=2, length=66), use_golden_token_injection=False, generate_hf_results=True, @@ -159,16 +159,77 @@ def test_pausing_prefill_volumetric_ok( "n_used_blocks": 4, }, { - # Decode sequences 0, 1, and 2 + # Decode sequences 0 and 1 # About to violate the volumetric constraint: 3 * 86 = 258 > 256 - # Holdback activates TODO + # Holdback activates: request 2 gets paused "step": 10, - "tkv": 86, + "tkv": 22, # tkv of request 0 without left-padding "waiting": [], - "running": ["2", "1", "0"], - "request_outputs": ["2", "1", "0"], + "running": ["1", "0"], + "request_outputs": ["1", "0"], "n_used_blocks": 4, }, + { + # Decode sequences 0 and 1 + # Sequences 0 finishes + "step": 13, + "tkv": 25, + "waiting": [], + "running": ["1"], + "request_outputs": ["1", "0"], + "finished_requests": ["0"], + "n_used_blocks": 3, + }, + { + # Decode sequences 1 and 2 + # Sequence 2 can now resume + "step": 14, + "tkv": 89, # 25 + 64 = tkv of request 1 with left-padding + "waiting": [], + "running": ["1", "2"], + "request_outputs": ["1", "2"], + "n_used_blocks": 3, + }, + { + # Decode sequences 1 and 2 + # Sequence 1 finishes + "step": 16, + "tkv": 91, + "waiting": [], + "running": ["2"], + "request_outputs": ["1", "2"], + "finished_requests": ["1"], + "n_used_blocks": 2, + }, + { + # Decode sequence 2 + "step": 17, + "tkv": 74, # tkv of request 2 + "waiting": [], + "running": ["2"], + "request_outputs": ["2"], + "n_used_blocks": 2, + }, + { + # Sequence 2 finishes + "step": 18, + "tkv": 75, + "waiting": [], + "running": [], + "request_outputs": ["2"], + "finished_requests": ["2"], + "n_used_blocks": 0, + }, + { + # tkv should be cleared one step later + "step": 19, + "tkv": 0, + "waiting": [], + "running": [], + "request_outputs": [], + "finished_requests": [], + "n_used_blocks": 0, + }, ] validate_scheduler_steps( @@ -191,7 +252,7 @@ def test_pausing_prefill_volumetric_ok( @pytest.mark.parametrize("max_model_len", [2048]) @pytest.mark.parametrize("max_num_batched_tokens", [128]) @pytest.mark.parametrize("available_blocks", [None]) -def test_pausing_prefill_volumetric_violated( +def test_prefill_volumetric_violated( model: ModelInfo, backend: str, monkeypatch: pytest.MonkeyPatch, @@ -201,7 +262,8 @@ def test_pausing_prefill_volumetric_violated( max_num_batched_tokens: int, available_blocks: int, ): - """Test that requests are blocked when prefill volumetric constraint is violated. + """Test that requests are blocked when the volumetric constraint is immediately + violated after prefill. Even with holdback, if prefill volume exceeds limit, request cannot be scheduled. diff --git a/tests/hf_cache.json b/tests/hf_cache.json index fd69acf20..9b401967b 100644 --- a/tests/hf_cache.json +++ b/tests/hf_cache.json @@ -638,6 +638,18 @@ "token_ids": [ 20, 313 ], "tokens": [ "\"", " \"" ], "logprobs": [ -2.742318630218506, -3.0037922859191895 ] + }, + "3": { + "text": "\" \" \"", + "token_ids": [ 20, 313, 313 ], + "tokens": [ "\"", " \"", " \"" ], + "logprobs": [ -2.742318630218506, -3.0037922859191895, -3.0905137062072754 ] + }, + "10": { + "text": "\" \" \" \" \" \" \" \" \" \"", + "token_ids": [ 20, 313, 313, 313, 313, 313, 313, 313, 313, 313 ], + "tokens": [ "\"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"", " \"" ], + "logprobs": [ -2.742318630218506, -3.0037922859191895, -3.0905137062072754, -0.9414255023002625, -0.42565011978149414, -0.24709248542785645, -0.2613930404186249, -0.46429160237312317, -0.6351388096809387, -0.771374523639679 ] } }, "__tokens__41505_37255_20672_12727_25130_19903_38525_14909_23426_28674_44635_24806_13853_37149_30394": { @@ -646,6 +658,24 @@ "token_ids": [ 203, 21, 1234, 5557, 54, 81, 2883, 81, 5762, 81, 2319, 81, 72, 35, 225, 34, 106, 34, 34, 34, 34, 34, 34, 34, 35, 203, 21, 1234, 5557, 54, 81, 2883, 81, 5762, 81, 2319, 81, 72, 36, 225, 34, 106, 34, 34, 34, 34, 34, 34, 34, 36, 203, 21, 1234, 5557, 54, 81, 2883, 81, 5762, 81, 2319, 81, 72, 37, 225, 34, 106, 34, 34, 34, 34, 34, 34, 34, 38, 203, 21, 1234, 5557, 54, 81, 2883, 81, 5762, 81, 2319, 81, 72, 38, 225, 34, 106, 34, 34, 34, 34, 34, 34, 34, 42 ], "tokens": [ "\n", "#", "define", " CH", "D", "_", "PRE", "_", "COMP", "_", "TYPE", "_", "V", "1", " ", "0", "x", "0", "0", "0", "0", "0", "0", "0", "1", "\n", "#", "define", " CH", "D", "_", "PRE", "_", "COMP", "_", "TYPE", "_", "V", "2", " ", "0", "x", "0", "0", "0", "0", "0", "0", "0", "2", "\n", "#", "define", " CH", "D", "_", "PRE", "_", "COMP", "_", "TYPE", "_", "V", "3", " ", "0", "x", "0", "0", "0", "0", "0", "0", "0", "4", "\n", "#", "define", " CH", "D", "_", "PRE", "_", "COMP", "_", "TYPE", "_", "V", "4", " ", "0", "x", "0", "0", "0", "0", "0", "0", "0", "8" ], "logprobs": [ -4.148627281188965, -1.9023098945617676, -0.5024883151054382, -3.5166265964508057, -1.1848623752593994, -0.44509002566337585, -4.0265913009643555, -2.1063308715820312, -4.351094722747803, -1.4263076782226562, -3.4072775840759277, -0.46827957034111023, -4.332708835601807, -1.7475272417068481, -1.3255337476730347, -0.8408962488174438, -0.37473323941230774, -0.7385716438293457, -0.47072669863700867, -0.18981407582759857, -0.36916476488113403, -0.1256936490535736, -0.17075695097446442, -0.17706118524074554, -0.4836171567440033, -0.5464026927947998, -0.2326568365097046, -0.04812277853488922, -0.03288242593407631, -0.004547967109829187, -0.0012984187342226505, -0.10411255806684494, -0.003740933956578374, -0.04705769941210747, -0.0035643160808831453, -0.18792352080345154, -0.002217930741608143, -0.0816296860575676, -0.16547933220863342, -0.033810585737228394, -0.03572966530919075, -0.018084051087498665, -0.017124062404036522, -0.010007210075855255, -0.016362886875867844, -0.028098611161112785, -0.033767715096473694, -0.023985574021935463, -0.06287175416946411, -0.06378459185361862, -0.014418580569326878, -0.08894632756710052, -0.02926105447113514, -0.028644727542996407, -0.0009793015196919441, -0.0002786724944598973, -0.16831068694591522, -0.0022170981392264366, -0.03330787271261215, -0.002313439268618822, -0.041456446051597595, -0.0002917817619163543, -0.15613055229187012, -0.03973784297704697, -0.008251740597188473, -0.005037354305386543, -0.0033077073749154806, -0.006288502831012011, -0.0025725625455379486, -0.003108076984062791, -0.005129747558385134, -0.00485254218801856, -0.005447543226182461, -0.026349563151597977, -0.12644031643867493, -0.008795449510216713, -0.10167873650789261, -0.01642856001853943, -0.015500782988965511, -0.00047994061606004834, -8.880697714630514e-05, -0.07469187676906586, -0.00070296844933182, -0.01724405214190483, -0.0008412636234425008, -0.014866752550005913, -9.274052717955783e-05, -0.07821717858314514, -0.03616850823163986, -0.004844712559133768, -0.001525192055851221, -0.002193070948123932, -0.0034923297353088856, -0.0012396040838211775, -0.0017492959741503, -0.0022299441043287516, -0.0032256022095680237, -0.003230711678043008, -0.029909281060099602, -0.011151960119605064 ] + }, + "60": { + "text": "\n#define CHD_PRE_COMP_TYPE_V1 0x00000001\n#define CHD_PRE_COMP_TYPE_V2 0x00000002\n#define CHD_PRE_COMP_", + "token_ids": [ 203, 21, 1234, 5557, 54, 81, 2883, 81, 5762, 81, 2319, 81, 72, 35, 225, 34, 106, 34, 34, 34, 34, 34, 34, 34, 35, 203, 21, 1234, 5557, 54, 81, 2883, 81, 5762, 81, 2319, 81, 72, 36, 225, 34, 106, 34, 34, 34, 34, 34, 34, 34, 36, 203, 21, 1234, 5557, 54, 81, 2883, 81, 5762, 81 ], + "tokens": [ "\n", "#", "define", " CH", "D", "_", "PRE", "_", "COMP", "_", "TYPE", "_", "V", "1", " ", "0", "x", "0", "0", "0", "0", "0", "0", "0", "1", "\n", "#", "define", " CH", "D", "_", "PRE", "_", "COMP", "_", "TYPE", "_", "V", "2", " ", "0", "x", "0", "0", "0", "0", "0", "0", "0", "2", "\n", "#", "define", " CH", "D", "_", "PRE", "_", "COMP", "_" ], + "logprobs": [ -4.148627281188965, -1.9023098945617676, -0.5024883151054382, -3.5166265964508057, -1.1848623752593994, -0.44509002566337585, -4.0265913009643555, -2.1063308715820312, -4.351094722747803, -1.4263076782226562, -3.4072775840759277, -0.46827957034111023, -4.332708835601807, -1.7475272417068481, -1.3255337476730347, -0.8408962488174438, -0.37473323941230774, -0.7385716438293457, -0.47072669863700867, -0.18981407582759857, -0.36916476488113403, -0.1256936490535736, -0.17075695097446442, -0.17706118524074554, -0.4836171567440033, -0.5464026927947998, -0.2326568365097046, -0.04812277853488922, -0.03288242593407631, -0.004547967109829187, -0.0012984187342226505, -0.10411255806684494, -0.003740933956578374, -0.04705769941210747, -0.0035643160808831453, -0.18792352080345154, -0.002217930741608143, -0.0816296860575676, -0.16547933220863342, -0.033810585737228394, -0.03572966530919075, -0.018084051087498665, -0.017124062404036522, -0.010007210075855255, -0.016362886875867844, -0.028098611161112785, -0.033767715096473694, -0.023985574021935463, -0.06287175416946411, -0.06378459185361862, -0.014418580569326878, -0.08894632756710052, -0.02926105447113514, -0.028644727542996407, -0.0009793015196919441, -0.0002786724944598973, -0.16831068694591522, -0.0022170981392264366, -0.03330787271261215, -0.002313439268618822 ] + }, + "10": { + "text": "\n#define CHD_PRE_COMP_", + "token_ids": [ 203, 21, 1234, 5557, 54, 81, 2883, 81, 5762, 81 ], + "tokens": [ "\n", "#", "define", " CH", "D", "_", "PRE", "_", "COMP", "_" ], + "logprobs": [ -4.148627281188965, -1.9023098945617676, -0.5024883151054382, -3.5166265964508057, -1.1848623752593994, -0.44509002566337585, -4.0265913009643555, -2.1063308715820312, -4.351094722747803, -1.4263076782226562 ] + }, + "11": { + "text": "\n#define CHD_PRE_COMP_TYPE", + "token_ids": [ 203, 21, 1234, 5557, 54, 81, 2883, 81, 5762, 81, 2319 ], + "tokens": [ "\n", "#", "define", " CH", "D", "_", "PRE", "_", "COMP", "_", "TYPE" ], + "logprobs": [ -4.148627281188965, -1.9023098945617676, -0.5024883151054382, -3.5166265964508057, -1.1848623752593994, -0.44509002566337585, -4.0265913009643555, -2.1063308715820312, -4.351094722747803, -1.4263076782226562, -3.4072775840759277 ] } }, "__tokens__6605_41653_37541_12537_24352_22093_32027_38767_4614_1394_41079_21271_37467_104_21892": { @@ -654,6 +684,24 @@ "token_ids": [ 225, 37, 32, 34, 32, 35, 225, 36, 34, 35, 43, 31, 34, 35, 31, 34, 35, 225, 35, 38, 44, 37, 34, 44, 34, 34, 225, 36, 34, 35, 43, 31, 34, 35, 31, 34, 35, 225, 35, 38, 44, 37, 34, 44, 34, 34, 225, 36, 34, 35, 43, 31, 34, 35, 31, 34, 35, 225, 35, 38, 44, 37, 34, 44, 34, 34, 225, 36, 34, 35, 43, 31, 34, 35, 31, 34, 35, 225, 35, 38, 44, 37, 34, 44, 34, 34, 225, 36, 34, 35, 43, 31, 34, 35, 31, 34, 35, 225, 35, 38 ], "tokens": [ " ", "3", ".", "0", ".", "1", " ", "2", "0", "1", "9", "-", "0", "1", "-", "0", "1", " ", "1", "4", ":", "3", "0", ":", "0", "0", " ", "2", "0", "1", "9", "-", "0", "1", "-", "0", "1", " ", "1", "4", ":", "3", "0", ":", "0", "0", " ", "2", "0", "1", "9", "-", "0", "1", "-", "0", "1", " ", "1", "4", ":", "3", "0", ":", "0", "0", " ", "2", "0", "1", "9", "-", "0", "1", "-", "0", "1", " ", "1", "4", ":", "3", "0", ":", "0", "0", " ", "2", "0", "1", "9", "-", "0", "1", "-", "0", "1", " ", "1", "4" ], "logprobs": [ -2.393317222595215, -2.2741634845733643, -0.5255809426307678, -1.454126000404358, -1.2284232378005981, -1.6320126056671143, -1.6596072912216187, -1.7695199251174927, -0.41706928610801697, -0.41384077072143555, -1.5353412628173828, -0.541284441947937, -0.3393153250217438, -1.750977873802185, -0.018852457404136658, -1.096146583557129, -1.5840284824371338, -0.4938022792339325, -0.7846470475196838, -2.0936717987060547, -0.39854463934898376, -1.633031964302063, -1.9479743242263794, -0.3339877128601074, -1.1566519737243652, -1.060813546180725, -0.9951281547546387, -2.2271621227264404, -0.1960439682006836, -0.1647973656654358, -0.09352599829435349, -0.3445914685726166, -0.014470632188022137, -0.1592160165309906, -0.001257463125512004, -0.010343162342905998, -0.10037858039140701, -0.05559161305427551, -0.07893017679452896, -0.14610762894153595, -0.016683464869856834, -0.0467216856777668, -0.025850284844636917, -0.030198249965906143, -0.043905384838581085, -0.08884358406066895, -1.2519291639328003, -0.5712020397186279, -0.013030261732637882, -0.04303565248847008, -0.04684977978467941, -0.029735142365098, -0.0020624573808163404, -0.016849223524332047, -0.0018703126115724444, -0.005184776149690151, -0.020718814805150032, -0.040916938334703445, -0.009638085961341858, -0.057150136679410934, -0.0024661386851221323, -0.010109765455126762, -0.004823832772672176, -0.005025255959481001, -0.0070279063656926155, -0.009060695767402649, -0.7706539630889893, -0.1780472695827484, -0.001004786929115653, -0.01082677487283945, -0.0162928719073534, -0.010154962539672852, -0.0010992205934599042, -0.010214435867965221, -0.0004924515378661454, -0.0026298719458281994, -0.013098269701004028, -0.016590023413300514, -0.0034246151335537434, -0.020235290750861168, -0.0010221739066764712, -0.002377542434260249, -0.0017613149248063564, -0.0027014450170099735, -0.00199333718046546, -0.003969291225075722, -0.3933572769165039, -0.04973369836807251, -0.00044967554276809096, -0.005075783468782902, -0.008561218157410622, -0.005584354046732187, -0.0008662762120366096, -0.0056595089845359325, -0.0003175231395289302, -0.0014316319720819592, -0.009665240533649921, -0.0076560406014323235, -0.0022322041913866997, -0.013572084717452526 ] + }, + "60": { + "text": " 3.0.1 2019-01-01 14:30:00 2019-01-01 14:30:00 2019-01-01 14", + "token_ids": [ 225, 37, 32, 34, 32, 35, 225, 36, 34, 35, 43, 31, 34, 35, 31, 34, 35, 225, 35, 38, 44, 37, 34, 44, 34, 34, 225, 36, 34, 35, 43, 31, 34, 35, 31, 34, 35, 225, 35, 38, 44, 37, 34, 44, 34, 34, 225, 36, 34, 35, 43, 31, 34, 35, 31, 34, 35, 225, 35, 38 ], + "tokens": [ " ", "3", ".", "0", ".", "1", " ", "2", "0", "1", "9", "-", "0", "1", "-", "0", "1", " ", "1", "4", ":", "3", "0", ":", "0", "0", " ", "2", "0", "1", "9", "-", "0", "1", "-", "0", "1", " ", "1", "4", ":", "3", "0", ":", "0", "0", " ", "2", "0", "1", "9", "-", "0", "1", "-", "0", "1", " ", "1", "4" ], + "logprobs": [ -2.393317222595215, -2.2741634845733643, -0.5255809426307678, -1.454126000404358, -1.2284232378005981, -1.6320126056671143, -1.6596072912216187, -1.7695199251174927, -0.41706928610801697, -0.41384077072143555, -1.5353412628173828, -0.541284441947937, -0.3393153250217438, -1.750977873802185, -0.018852457404136658, -1.096146583557129, -1.5840284824371338, -0.4938022792339325, -0.7846470475196838, -2.0936717987060547, -0.39854463934898376, -1.633031964302063, -1.9479743242263794, -0.3339877128601074, -1.1566519737243652, -1.060813546180725, -0.9951281547546387, -2.2271621227264404, -0.1960439682006836, -0.1647973656654358, -0.09352599829435349, -0.3445914685726166, -0.014470632188022137, -0.1592160165309906, -0.001257463125512004, -0.010343162342905998, -0.10037858039140701, -0.05559161305427551, -0.07893017679452896, -0.14610762894153595, -0.016683464869856834, -0.0467216856777668, -0.025850284844636917, -0.030198249965906143, -0.043905384838581085, -0.08884358406066895, -1.2519291639328003, -0.5712020397186279, -0.013030261732637882, -0.04303565248847008, -0.04684977978467941, -0.029735142365098, -0.0020624573808163404, -0.016849223524332047, -0.0018703126115724444, -0.005184776149690151, -0.020718814805150032, -0.040916938334703445, -0.009638085961341858, -0.057150136679410934 ] + }, + "12": { + "text": " 3.0.1 2019-", + "token_ids": [ 225, 37, 32, 34, 32, 35, 225, 36, 34, 35, 43, 31 ], + "tokens": [ " ", "3", ".", "0", ".", "1", " ", "2", "0", "1", "9", "-" ], + "logprobs": [ -2.393317222595215, -2.2741634845733643, -0.5255809426307678, -1.454126000404358, -1.2284232378005981, -1.6320124864578247, -1.6596072912216187, -1.7695199251174927, -0.41706928610801697, -0.41384077072143555, -1.5353412628173828, -0.541284441947937 ] + }, + "13": { + "text": " 3.0.1 2019-0", + "token_ids": [ 225, 37, 32, 34, 32, 35, 225, 36, 34, 35, 43, 31, 34 ], + "tokens": [ " ", "3", ".", "0", ".", "1", " ", "2", "0", "1", "9", "-", "0" ], + "logprobs": [ -2.393317222595215, -2.2741634845733643, -0.5255809426307678, -1.454126000404358, -1.2284232378005981, -1.6320124864578247, -1.6596072912216187, -1.7695199251174927, -0.41706928610801697, -0.41384077072143555, -1.5353412628173828, -0.541284441947937, -0.3393153250217438 ] } }, "__tokens__41505_37255_20672_12727_25130_19903_38525_14909_23426_28674_44635_24806_13853_37149_30394_12313_44715_48305_39823_44343_15245_35872_44179_33619_23207": { From 2053a3b8be425150f1d4aa4a7d5d8c3a211aff53 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 29 May 2026 16:40:58 +0000 Subject: [PATCH 016/106] handle paused decoding requests abort MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index e2f1020da..5d460edc8 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -700,7 +700,10 @@ def finish_requests( request_ids: Union[str, Iterable[str], None], finished_status: RequestStatus, ) -> list[tuple[str, int]]: - """Handles removing finished requests from ongoing_prefills""" + """ + Handles removing finished requests from ongoing_prefills and + paused_decoding_requests + """ if isinstance(request_ids, str): request_ids = (request_ids,) @@ -717,6 +720,13 @@ def finish_requests( else [r for r in self.ongoing_prefills if r.request_id not in request_ids] ) + # Also remove from paused_decoding_requests + self.paused_decoding_requests = ( + [] + if request_ids is None + else [r for r in self.paused_decoding_requests if r.request_id not in request_ids] + ) + return aborted_requests def calc_cached_tokens(self, prompt_len: int) -> tuple[int, int]: From 6a4c60593f2d7da64535750e74a3671d3ec55c52 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 29 May 2026 19:50:50 +0000 Subject: [PATCH 017/106] fix tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- tests/v1/core/test_scheduler_structured_outputs.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/v1/core/test_scheduler_structured_outputs.py b/tests/v1/core/test_scheduler_structured_outputs.py index 0c0a65694..710d3582f 100644 --- a/tests/v1/core/test_scheduler_structured_outputs.py +++ b/tests/v1/core/test_scheduler_structured_outputs.py @@ -39,6 +39,7 @@ def mocked_scheduler(): scheduler.skipped_waiting = FCFSRequestQueue() scheduler.running = [] scheduler.ongoing_prefills = [] + scheduler.paused_decoding_requests = [] scheduler.chunk_size = 128 scheduler.do_interleaving = False scheduler.previous_step_was_prefill = False From 85556569ece42c9dc188582e4727a67633208173 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 29 May 2026 19:52:08 +0000 Subject: [PATCH 018/106] only add requests that were paused MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/worker/spyre_model_runner.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index 9a761d883..374867566 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -154,6 +154,9 @@ def __init__( # Requests self.requests: dict[str, RequestStateT] = {} + # Track paused requests to ensure we only restore previously paused ones + self.paused_req_ids: set[str] = set() + @abstractmethod def build_input_batch(self) -> InputBatchT: raise NotImplementedError @@ -1461,16 +1464,20 @@ def _update_batch(self, scheduler_output: SchedulerOutput): if req_id not in (scheduler_output.finished_req_ids or []): logger.info("Removing paused request %s from input_batch", req_id) self.input_batch.remove_request(req_id) + # Track that this request was paused + self.paused_req_ids.add(req_id) # Find requests that are in scheduler output but not in input_batch # (restore from pausing) restored_req_ids = scheduled_req_ids - current_batch_req_ids for req_id in restored_req_ids: - # Add back the request that was paused - if req_id in self.requests: + # Only restore requests that were previously paused + if req_id in self.paused_req_ids and req_id in self.requests: logger.info("Restoring paused request %s to input_batch", req_id) req_state = self.requests[req_id] self.input_batch.add_request(req_state) + # Remove from paused tracking since it's now restored + self.paused_req_ids.discard(req_id) for i, req_id in enumerate(req_data.req_ids): req_state: SamplingRequestState = self.requests[req_id] @@ -1487,6 +1494,8 @@ def _update_batch(self, scheduler_output: SchedulerOutput): if scheduler_output.finished_req_ids: for req_id in scheduler_output.finished_req_ids: self.input_batch.remove_request(req_id) + # Clean up paused tracking for finished requests + self.paused_req_ids.discard(req_id) # TODO: Processing multiple removals at once can break alignment # of logitprocs. Refactor so that we can batch removals to the # `input_batch` From 01e43b27a9133ddae02d62336e5fcd518b654c01 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 29 May 2026 20:04:37 +0000 Subject: [PATCH 019/106] fix tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- tests/v1/worker/test_prefix_caching_worker.py | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/tests/v1/worker/test_prefix_caching_worker.py b/tests/v1/worker/test_prefix_caching_worker.py index c9426a7d5..b5549cfb7 100644 --- a/tests/v1/worker/test_prefix_caching_worker.py +++ b/tests/v1/worker/test_prefix_caching_worker.py @@ -202,17 +202,17 @@ def test_multi_chunk_partial_match_misaligned( # Schedule decodes of requests 0 and 1 model_runner_output_6 = pc_model_runner.execute_running_requests() pc_model_runner.assert_block_tables_and_slot_mappings( - block_tables=[[1, 2, 3, 4, 5, 6, 11], [1, 2, 3, 7, 8, 9, 10]], - slot_mappings=[[11], [10]], + block_tables=[[1, 2, 3, 7, 8, 9, 10], [1, 2, 3, 4, 5, 6, 11]], + slot_mappings=[[10], [11]], slot_slice=slice(0, 1), ) pc_model_runner.verify_model_runner_output( model_runner_output_6, - req_ids=["0", "1"], + req_ids=["1", "0"], num_sampled_token_ids=2, tkv=385, n_free_blocks=17, - left_padding={"0": 0, "1": 0}, + left_padding={"1": 0, "0": 0}, ) @@ -313,17 +313,17 @@ def test_first_chunk_recomputation( # Schedule decodes of requests 0 and 1 model_runner_output_3 = pc_model_runner.execute_running_requests() pc_model_runner.assert_block_tables_and_slot_mappings( - block_tables=[[1, 2, 5], [1, 3, 4]], - slot_mappings=[[5], [4]], + block_tables=[[1, 3, 4], [1, 2, 5]], + slot_mappings=[[4], [5]], slot_slice=slice(0, 1), ) pc_model_runner.verify_model_runner_output( model_runner_output_3, - req_ids=["0", "1"], + req_ids=["1", "0"], num_sampled_token_ids=2, tkv=129, n_free_blocks=17, - left_padding={"0": 0, "1": 0}, + left_padding={"1": 0, "0": 0}, ) @@ -442,16 +442,16 @@ def test_middle_chunk_recomputation_with_padding( # Schedule decodes of requests 0 and 1 model_runner_output_4 = pc_model_runner.execute_running_requests() pc_model_runner.assert_block_tables_and_slot_mappings( - block_tables=[[0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 12], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]], - slot_mappings=[[12], [11]], + block_tables=[[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], [0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 12]], + slot_mappings=[[11], [12]], slot_slice=slice(0, 1), ) pc_model_runner.verify_model_runner_output( model_runner_output_4, - req_ids=["0", "1"], + req_ids=["1", "0"], num_sampled_token_ids=2, tkv=641, n_free_blocks=33, - left_padding={"0": 128, "1": 0}, + left_padding={"1": 0, "0": 128}, ) From 867e38c82d0ce7b3610b5c893592f0c10a72faa9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Sat, 30 May 2026 22:13:28 +0000 Subject: [PATCH 020/106] fix tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../v1/worker/spyre_model_runner.py | 19 +- tests/e2e/test_spyre_cp_scheduler_steps.py | 317 +++++------------- ...test_spyre_decode_pause_scheduler_steps.py | 22 +- tests/e2e/test_spyre_pc_scheduler_steps.py | 6 +- tests/hf_cache.json | 62 ++++ 5 files changed, 167 insertions(+), 259 deletions(-) diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index 374867566..bca4151cc 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -1007,6 +1007,9 @@ def _prepare_chunked_prefill(self, req_id: str) -> SamplingForwardInputs: input_tokens = input_tokens.unsqueeze(0).clone() input_positions = input_positions.unsqueeze(0).clone() + # Set tkv to prompt length at each prefill step + self.tkv = prompt_len + # NOTE(wallas): Looks like we need to use multiple of blocks for prefill # so, later we use model.n_pads_right to get right logits. # In my naive mind this would be the `request_tkv` below, @@ -1287,16 +1290,10 @@ def add_new_request(self, request: NewRequestData): assert sampling_params is not None, "sampling_params are required for this model runner" assert prompt_token_ids is not None, "prompt token ids are required for this model runner" - is_new_batch = self.input_batch.num_reqs == 0 - prompt_len = len(prompt_token_ids) mm_features = getattr(request, "mm_features", None) self.prefill_batch.clear_requests() - # set the new tkv to the prompt length if starting a new decode batch - if is_new_batch: - self.tkv = prompt_len - block_ids_per_kv_cache_group = request.block_ids assert len(block_ids_per_kv_cache_group) == 1 @@ -1351,16 +1348,6 @@ def _maybe_prepare_last_prefill(self, req_id: str, scheduler_output: SchedulerOu ) request.cached_mm_embeddings = None - # Last prefill: we might need to update the tkv - req_n_blocks = math.ceil(prompt_len / self.block_size) - cur_n_blocks = math.ceil(self.tkv / self.block_size) - new_n_blocks = max(req_n_blocks, cur_n_blocks) - assert new_n_blocks > 0 - base_n_tokens = (new_n_blocks - 1) * self.block_size - req_tkv_new_block = base_n_tokens + (prompt_len - 1) % self.block_size + 1 - cur_tkv_new_block = base_n_tokens + (self.tkv - 1) % self.block_size + 1 - self.tkv = max(req_tkv_new_block, cur_tkv_new_block) - # Last prefill we need to setup the logitsprocessors to sampling prefill_index = self.input_batch.add_request(request) for logitsproc in self.input_batch.logitsprocs_wrappers: diff --git a/tests/e2e/test_spyre_cp_scheduler_steps.py b/tests/e2e/test_spyre_cp_scheduler_steps.py index ac712991d..0c5279955 100644 --- a/tests/e2e/test_spyre_cp_scheduler_steps.py +++ b/tests/e2e/test_spyre_cp_scheduler_steps.py @@ -7,7 +7,12 @@ """ import pytest -from scheduling_utils import check_scheduler_inference_steps +from scheduling_utils import ( + check_scheduler_inference_steps, + create_request_for_scheduler_test, + random_prompt, + validate_scheduler_steps, +) from spyre_util import ModelInfo @@ -28,23 +33,38 @@ def test_prefill_tkv_too_big( max_num_batched_tokens: int, available_blocks: int, ): - """Scenario where the requested prompt is too long for current tkv value + """Here we ensure that the tkv never goes beyond max_model_len, even in an + edge case. - Note that as we could prefill the prompt straight away, however, - in this test the max model length is decreased to a value where - the tkv of the decode batch would be shifted beyond the max model length, - we therefore have to wait with scheduling. + Edge case: due to a long-prompt joining the decode batch, the currently + decoding request needs to be left-padded, bringing the max-tokens beyond + max-model-len. We make sure the left-padding is removed on time when + expanding to a new block, keeping the tkv in acceptable range always. Configuration: * max_num_seqs: 2 * number of prompts: 2 - * 0: len = 49, max tokens = 17, step joining = 0 - * 1: len = 70, max tokens = 17, step joining = 0 + * 0: len = 60, max tokens = 10, step joining = 0 + * 1: len = 111, max tokens = 17, step joining = 0 """ - seqs_max_tokens = [17, 17] - prompts_lengths = [49, 70] - steps_add_reqs = [0, 0] + request1 = create_request_for_scheduler_test( + model=model, + request_id=0, + add_step=0, + max_tokens=10, + prompt=random_prompt(model=model, seed=0, length=60), + use_golden_token_injection=True, + ) + + request2 = create_request_for_scheduler_test( + model=model, + request_id=1, + add_step=0, + max_tokens=17, + prompt=random_prompt(model=model, seed=0, length=111), + use_golden_token_injection=True, + ) checked_steps = [ { @@ -57,66 +77,86 @@ def test_prefill_tkv_too_big( }, { # Prefill sequence 0 - # total blocks in use: 1 "step": 1, - "tkv": 49, # prompt len + "tkv": 60, "waiting": ["1"], "running": ["0"], "request_outputs": ["0"], "n_used_blocks": 1, }, - # Here we cannot schedule sequence 1. By shifting sequence 0 by - # 1 block its max tkv would exceed the max model length { # Decode sequence 0 - # total blocks in use: 1 (writing into right pads) "step": 2, - "tkv": 50, + "tkv": 61, "waiting": ["1"], "running": ["0"], "request_outputs": ["0"], "n_used_blocks": 1, }, { - # Prefill sequence 1, tkv large enough to prefill w/o tkv shift - # total blocks in use: 1 + 2 - "step": 17, - # add 64 to tkv of seq 0 (64) to have it in the same block as seq 1 - "tkv": 128, + # Prefill sequence 1 + "step": 3, + "tkv": 111, # tkv of the prefilling sequence 1 "waiting": [], "running": ["1", "0"], "request_outputs": ["1"], - # 2 + 2 (prefill (2 block) + 17 decodes in the last block) "n_used_blocks": 3, }, + { + # Decode sequences 0 and 1 + "step": 4, + "tkv": 126, # left-padding of sequence 0: 64 + 62 = 126 + "waiting": [], + "running": ["1", "0"], + "request_outputs": ["1", "0"], + "n_used_blocks": 3, + }, + { + # Decode sequences 0 and 1 + # Last step before tkv would overflow max_context_length + "step": 6, + "tkv": 128, + "waiting": [], + "running": ["1", "0"], + "request_outputs": ["1", "0"], + "n_used_blocks": 3, + }, + { + # Decode sequences 0 and 1 + # Sequence 0 now needs two blocks. Instead of adding one on the + # right (which would overflow the tkv), we remove it's left-padding + # block, bringing back the tkv to a satisfactory value + "step": 7, + "tkv": 115, # corresponds now to tkv of request 1 + "waiting": [], + "running": ["1", "0"], + "request_outputs": ["1", "0"], + "n_used_blocks": 4, + }, { # Decode sequences 0 and 1 # Sequence 0 finishes - "step": 18, - # remove left padding of seq 0, and keep its tkv in the same block - # as seq 1: 129 - 64 = 65 - # tkv of seq 1 is now max - "tkv": 71, + "step": 11, + "tkv": 119, "waiting": [], "running": ["1"], "request_outputs": ["1", "0"], "finished_requests": ["0"], - "n_used_blocks": 2, # seq 0 needs another block for the last token + "n_used_blocks": 2, }, { - # Decode sequence 1 - # total blocks in use: 4 - 2 = 2 - "step": 19, - "tkv": 72, + # Decode sequences 1 + "step": 12, + "tkv": 120, "waiting": [], "running": ["1"], "request_outputs": ["1"], "n_used_blocks": 2, }, { - # Sequence 1 finishes at step 33 - "step": 33, - "tkv": 86, + # Sequence 1 finishes + "step": 19, + "tkv": 127, "waiting": [], "running": [], "request_outputs": ["1"], @@ -125,7 +165,7 @@ def test_prefill_tkv_too_big( }, { # Tkv should be cleared one step later - "step": 34, + "step": 20, "tkv": 0, "waiting": [], "running": [], @@ -134,13 +174,11 @@ def test_prefill_tkv_too_big( }, ] - check_scheduler_inference_steps( + validate_scheduler_steps( model=model, backend=backend, monkeypatch=monkeypatch, - seqs_max_tokens=seqs_max_tokens, - prompts_lengths=prompts_lengths, - steps_add_reqs=steps_add_reqs, + requests=[request1, request2], checked_steps=checked_steps, max_num_seqs=max_num_seqs, max_model_len=max_model_len, @@ -453,7 +491,7 @@ def test_cp_prefill_interleave1( { # Chunk 0 of request 1 prefill "step": 3, - "tkv": 11, + "tkv": 512, "waiting": [], "running": ["1", "0"], "request_outputs": [], @@ -472,7 +510,7 @@ def test_cp_prefill_interleave1( # Chunk 1 of request 1 prefill # tkv of decode batch (tkv not updated until last chunk) "step": 5, - "tkv": 12, + "tkv": 512, "waiting": [], "running": ["1", "0"], "request_outputs": [], @@ -491,7 +529,7 @@ def test_cp_prefill_interleave1( # Chunk 2 of request 1 prefill # tkv of decode batch (tkv not updated until last chunk) "step": 7, - "tkv": 13, + "tkv": 512, "waiting": [], "running": ["1", "0"], "request_outputs": [], @@ -631,7 +669,7 @@ def test_cp_prefill_no_interleave( # Chunk 0 of request 1 prefill # tkv of decode batch (tkv not updated until last chunk) "step": 2, - "tkv": 10, + "tkv": 512, "waiting": [], "running": ["1", "0"], "request_outputs": [], @@ -641,7 +679,7 @@ def test_cp_prefill_no_interleave( # Chunk 1 of request 1 prefill # tkv of decode batch (tkv not updated until last chunk) "step": 3, - "tkv": 10, + "tkv": 512, "waiting": [], "running": ["1", "0"], "request_outputs": [], @@ -651,7 +689,7 @@ def test_cp_prefill_no_interleave( # Chunk 2 of request 1 prefill # tkv of decode batch (tkv not updated until last chunk) "step": 4, - "tkv": 10, + "tkv": 512, "waiting": [], "running": ["1", "0"], "request_outputs": [], @@ -832,7 +870,7 @@ def test_cp_prefill_interleave2( # Chunk 0 of request 1 prefill # tkv of decode batch (tkv not updated until last chunk) "step": 4, - "tkv": 12, + "tkv": 512, "waiting": [], "running": ["1", "0"], "request_outputs": [], @@ -851,7 +889,7 @@ def test_cp_prefill_interleave2( # Chunk 1 of request 1 prefill # tkv of decode batch (tkv not updated until last chunk) "step": 6, - "tkv": 13, + "tkv": 512, "waiting": [], "running": ["1", "0"], "request_outputs": [], @@ -870,7 +908,7 @@ def test_cp_prefill_interleave2( # Chunk 2 of request 1 prefill # tkv of decode batch (tkv not updated until last chunk) "step": 8, - "tkv": 14, + "tkv": 512, "waiting": [], "running": ["1", "0"], "request_outputs": [], @@ -951,182 +989,3 @@ def test_cp_prefill_interleave2( random_prompts=True, max_num_batched_tokens=max_num_batched_tokens, ) - - -# TODO had to move test at the end, having it after test_prefill_tkv_too_big -# was breaking the ordering ("error in test ordering!") -# looks like an issue with sorting the runtime configurations -@pytest.mark.chunked_prefill -@pytest.mark.full_model -# These values are all parameterized for test sorting -@pytest.mark.parametrize("max_num_seqs", [4]) -@pytest.mark.parametrize("max_model_len", [128]) # restricted to violate scheduler condition -@pytest.mark.parametrize("max_num_batched_tokens", [128]) -@pytest.mark.parametrize("available_blocks", [None]) -def test_prefill_tkv_too_big2( - model: ModelInfo, - backend: str, - monkeypatch: pytest.MonkeyPatch, - set_random_seed, - max_num_seqs: int, - max_model_len: int, - max_num_batched_tokens: int, - available_blocks: int, -): - """Scenario where the requested number of output is too big for current - tkv value. We need to wait for a previous long prompt request to finish and - have tkv reduced to a previous block before being able to schedule the - new request. - - Configuration: - * max_num_seqs: 4 - * number of prompts: 3 - * 0: len = 20, max tokens = 5, step joining = 0 - * 1: len = 80, max tokens = 3, step joining = 0 - * 2: len = 16, max tokens = 50, step joining = 0 - """ - - monkeypatch.setenv("SENDNN_INFERENCE_CP_INTERLEAVE_STEPS", "0") - - seqs_max_tokens = [5, 3, 50] - prompts_lengths = [20, 80, 16] - steps_add_reqs = [0, 0, 0] - - checked_steps = [ - { - "step": 0, - "tkv": 0, - "waiting": ["0", "1", "2"], - "running": [], - "request_outputs": [], - "n_used_blocks": 0, - }, - { - # Prefill sequence 0 - "step": 1, - "tkv": 20, - "waiting": ["1", "2"], - "running": ["0"], - "request_outputs": ["0"], - "n_used_blocks": 1, - }, - # tkv should be updated at the end of the last chunked prefill - # here we have only one chunk, so it will be updated directly - { - # Prefill sequence 1 - "step": 2, - "tkv": 84, # 64 (1 block padding) + 20 (prompt of seq 0) = 84 - "waiting": ["2"], - "running": ["1", "0"], - "request_outputs": ["1"], - # 1 + 2 (prefill (2 block) + 3 decodes in the last block) - "n_used_blocks": 3, - }, - # Here we cannot schedule sequence 2. Current tkv being in the second - # block, the number of requested tokens can't fit in the remaining space - # (64 (full block left padding) + 16 (prompt) + 50 (decode) = 130 > 128) - { - # Decode 1 of sequence 0 - # Decode 1 of sequence 1 - "step": 3, - "tkv": 85, - "waiting": ["2"], - "running": ["1", "0"], - "request_outputs": ["1", "0"], - "n_used_blocks": 3, - }, - { - # Decode 2 of sequence 0 - # Decode 2 of sequence 1 - # Sequence 1 finishes - "step": 4, - "tkv": 86, - "waiting": ["2"], - "running": ["0"], - "request_outputs": ["1", "0"], - "finished_requests": ["1"], - "n_used_blocks": 1, - }, - # The tkv value used here is computed before the model forward pass and - # token sampling of this step. As a result, it does not yet reflect - # sequences that finish in the current step. In this case, tkv=86 still - # includes sequence 1, which completes in this step, and this will only - # be accounted for in the next step. Therefore, sequence 2 cannot be - # prefilled yet at this point. - { - # Decode 3 of sequence 0 - "step": 5, - "tkv": 23, # 20 (prompt len) + 3 (decodes) = 23 - "waiting": ["2"], - "running": ["0"], - "request_outputs": ["0"], - "n_used_blocks": 1, - }, - # Sequence 2 can be scheduled for prefill, now that tkv is moved back to - # the first block. - { - # Prefill sequence 2 - "step": 6, - "tkv": 23, - "waiting": [], - "running": ["2", "0"], - "request_outputs": ["2"], - # 3 - 2 (finished seq 1) + 2 (prefill + 50 decodes in new block) - "n_used_blocks": 2, - }, - { - # Decode 4 of sequence 0 - # Decode 1 of sequence 2 - # Sequence 0 finishes - "step": 7, - "tkv": 24, - "waiting": [], - "running": ["2"], - "request_outputs": ["2", "0"], - "finished_requests": ["0"], - "n_used_blocks": 1, - }, - { - # Decode 2 of sequence 2 - "step": 8, - "tkv": 18, # 16 (prompt len) + 2 (decodes) = 18 - "waiting": [], - "running": ["2"], - "request_outputs": ["2"], - "n_used_blocks": 1, - }, - { - # Decode 49 of sequence 2 - # Sequence 2 finishes - "step": 55, - "tkv": 65, # 16 (prompt len) + 49 (decodes) = 65 - "waiting": [], - "running": [], - "request_outputs": ["2"], - "finished_requests": ["2"], - "n_used_blocks": 0, - }, - { - # tkv should be cleared one step later - "step": 56, - "tkv": 0, - "waiting": [], - "running": [], - "request_outputs": [], - "n_used_blocks": 0, - }, - ] - - check_scheduler_inference_steps( - model=model, - backend=backend, - monkeypatch=monkeypatch, - seqs_max_tokens=seqs_max_tokens, - prompts_lengths=prompts_lengths, - steps_add_reqs=steps_add_reqs, - checked_steps=checked_steps, - max_num_seqs=max_num_seqs, - max_model_len=max_model_len, - available_blocks=available_blocks, - max_num_batched_tokens=max_num_batched_tokens, - ) diff --git a/tests/e2e/test_spyre_decode_pause_scheduler_steps.py b/tests/e2e/test_spyre_decode_pause_scheduler_steps.py index 865395422..267ee72e5 100644 --- a/tests/e2e/test_spyre_decode_pause_scheduler_steps.py +++ b/tests/e2e/test_spyre_decode_pause_scheduler_steps.py @@ -24,9 +24,9 @@ @pytest.mark.full_model @pytest.mark.parametrize("max_num_seqs", [4]) @pytest.mark.parametrize("max_model_len", [128]) -@pytest.mark.parametrize("max_num_batched_tokens", [256]) +@pytest.mark.parametrize("max_num_batched_tokens", [128]) @pytest.mark.parametrize("available_blocks", [None]) -def test_volumetric_decode_pausing( +def test_max_batch_tkv_decode_pausing( model: ModelInfo, backend: str, monkeypatch: pytest.MonkeyPatch, @@ -36,10 +36,11 @@ def test_volumetric_decode_pausing( max_num_batched_tokens: int, available_blocks: int, ): - """Test that requests are scheduled when prefill volumetric constraint is satisfied. + """Test that requests are scheduled and removed on time before max batch tkv + exceeds the limit. - With holdback, we only check current_max_tkv * batch_size <= limit, - not future volumetric constraints. + With pausing, we only check current_max_tkv * batch_size <= limit, + not future max_batch_tkv values. Configuration: * max_num_seqs: 2 @@ -50,7 +51,6 @@ def test_volumetric_decode_pausing( """ # Volume right after prefill: 3 * 82 = 246 (should pass) - # Old scheduler would check future: TODO (would fail) max_batch_tkv_limit = 256 requests = [ @@ -132,7 +132,7 @@ def test_volumetric_decode_pausing( # Prefill sequence 2 # With holdback: sequence 2 CAN be scheduled # Decode volume: 3 * 82 = 246 <= 256 (passes) - # Old scheduler would block: future volume 3 * TODO = 330 > 256 + # Old scheduler would block: future volume 3 * 91 = 273 > 256 "step": 5, "tkv": 66, "waiting": [], @@ -249,10 +249,10 @@ def test_volumetric_decode_pausing( @pytest.mark.chunked_prefill @pytest.mark.full_model @pytest.mark.parametrize("max_num_seqs", [2]) -@pytest.mark.parametrize("max_model_len", [2048]) +@pytest.mark.parametrize("max_model_len", [128]) @pytest.mark.parametrize("max_num_batched_tokens", [128]) @pytest.mark.parametrize("available_blocks", [None]) -def test_prefill_volumetric_violated( +def test_prefill_exceeds_max_batch_tkv( model: ModelInfo, backend: str, monkeypatch: pytest.MonkeyPatch, @@ -265,14 +265,14 @@ def test_prefill_volumetric_violated( """Test that requests are blocked when the volumetric constraint is immediately violated after prefill. - Even with holdback, if prefill volume exceeds limit, request cannot be scheduled. + Even with pausing, if prefill volume exceeds limit, request cannot be scheduled. Configuration: * max_num_seqs: 2 * number of prompts: 2 * 0: len = 25, max tokens = 7, step joining = 0 * 1: len = 25, max tokens = 6, step joining = 0 - * 1: len = 66, max tokens = 3, step joining = 0 + * 2: len = 66, max tokens = 3, step joining = 0 """ # Volume right after prefill: 3 * 89 = 267 (fails) diff --git a/tests/e2e/test_spyre_pc_scheduler_steps.py b/tests/e2e/test_spyre_pc_scheduler_steps.py index d7496e8ce..624522884 100644 --- a/tests/e2e/test_spyre_pc_scheduler_steps.py +++ b/tests/e2e/test_spyre_pc_scheduler_steps.py @@ -397,7 +397,7 @@ def test_prefix_hit_decoded_block_within_batch( "step": 68, # seq 1 tkv (193) is in 4th block. Need to pad seq 0 tkv to 4th # block as well: 192 + 64 = 256 - "tkv": 256, + "tkv": 193, "waiting": [], "running": ["1", "0"], "request_outputs": ["1"], @@ -1650,7 +1650,7 @@ def test_first_chunk_partial_match( request_id=1, add_step=0, max_tokens=2, - prompt=prompt2, + prompt=prompt2, # prompt_len: 64 + (192 * 2) - 64 - 64 = 320 use_golden_token_injection=True, ) @@ -1675,7 +1675,7 @@ def test_first_chunk_partial_match( }, { # prefill seq 1. This step was crashing before "step": 2, - "tkv": 64, + "tkv": 320, "waiting": [], "running": ["1", "0"], "request_outputs": [], diff --git a/tests/hf_cache.json b/tests/hf_cache.json index 9b401967b..6a29a01e6 100644 --- a/tests/hf_cache.json +++ b/tests/hf_cache.json @@ -731,6 +731,22 @@ "tokens": [ "\n ", " STDERR", " Simulator", "Al", "node", " graphs" ], "logprobs": [ -2.7261359691619873, -1.2988193035125732, -0.39854758977890015, -0.07328899204730988, -0.11031155288219452, -0.2497384399175644 ] } + }, + "__tokens__41505_37255_20672_12727_25130_19903_38525_14909_23426_28674_44635_24806_13853_37149_30394_12313_44715_48305_39823_44343_15245_35872_44179_33619_23207_4950_21340_30026_44876_47510_23446_42531_12804_39568_26970_691_35375_19603_40542_32841_57_24260_42644_11989_15985_42785_9392_27894_11729_47556_39478_22019_3954_15732_24966_45850_5361_27096_34729_26908": { + "10": { + "text": "scape; \u201curalSOFTWAREffset Landscape; \u201c", + "token_ids": [ 10946, 45, 7850, 5361, 27096, 34729, 26908, 10946, 45, 7850 ], + "tokens": [ "scape", ";", " \u201c", "ural", "SOFTWARE", "ffset", " Land", "scape", ";", " \u201c" ], + "logprobs": [ -0.9898954033851624, -4.104252815246582, -2.3384828567504883, -2.246047258377075, -1.1454188823699951, -2.8821964263916016, -0.6978211402893066, -0.01959829591214657, -0.49376457929611206, -0.3851028084754944 ] + } + }, + "__tokens__41505_37255_20672_12727_25130_19903_38525_14909_23426_28674_44635_24806_13853_37149_30394_12313_44715_48305_39823_44343_15245_35872_44179_33619_23207_4950_21340_30026_44876_47510_23446_42531_12804_39568_26970_691_35375_19603_40542_32841_57_24260_42644_11989_15985_42785_9392_27894_11729_47556_39478_22019_3954_15732_24966_45850_5361_27096_34729_26908_40032_26556_47374_29648_28882_21872_29309_18919_28294_14270_9309_9178_30119_32276_23422_4415_37237_43095_45386_41408_44147_45371_26572_19233_34666_13548_39893_41754_43993_28990_46682_28493_22146_32452_48968_45069_38993_4049_30119_23910_30973_41537_11946_35954_5758_10836_39055_16345_40103_4945_7194": { + "17": { + "text": "lestutheseksenoticedoubles$\u201d\u201d\u201d\u201d", + "token_ids": [ 274, 270, 303, 292, 277, 5505, 272, 366, 13350, 97, 1187, 101, 22, 4830, 4830, 4830, 4830 ], + "tokens": [ "le", "st", "ut", "he", "se", "ks", "en", "ot", "iced", "o", "uble", "s", "$", "\u201d", "\u201d", "\u201d", "\u201d" ], + "logprobs": [ -2.5207459926605225, -3.87141489982605, -3.726951837539673, -3.8319180011749268, -2.156740427017212, -2.9964098930358887, -3.5194807052612305, -3.5848548412323, -2.74958872795105, -2.9147884845733643, -1.9768964052200317, -2.396059513092041, -3.6019279956817627, -2.8196353912353516, -2.8711845874786377, -1.8511260747909546, -0.4670819640159607 ] + } } } }, @@ -1329,6 +1345,52 @@ "tokens": [ "\n", "how", " do", " I", " add", " multiple", " new", " columns" ], "logprobs": [ -0.5748504996299744, -1.667298436164856, -0.15292026102542877, -0.07519898563623428, -0.25922346115112305, -0.09096132218837738, -0.573093831539154, -0.009411255829036236 ] } + }, + "__tokens__41505_37255_20672_12727_25130_19903_38525_14909_23426_28674_44635_24806_13853_37149_30394_12313_44715_48305_39823_44343_15245_35872_44179_33619_23207": { + "7": { + "text": "\n\n#include Date: Mon, 1 Jun 2026 22:49:08 +0000 Subject: [PATCH 021/106] pause/resume concept in intput batch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../v1/sample/spyre_logits_processor.py | 99 +++- .../v1/worker/spyre_input_batch.py | 151 ++++- .../v1/worker/spyre_model_runner.py | 12 +- tests/e2e/test_logits_processors.py | 537 +++++++++++++++++- tests/v1/worker/test_spyre_input_batch.py | 135 ++++- 5 files changed, 901 insertions(+), 33 deletions(-) diff --git a/sendnn_inference/v1/sample/spyre_logits_processor.py b/sendnn_inference/v1/sample/spyre_logits_processor.py index d44317838..a0865b2b2 100644 --- a/sendnn_inference/v1/sample/spyre_logits_processor.py +++ b/sendnn_inference/v1/sample/spyre_logits_processor.py @@ -1,4 +1,5 @@ import itertools +from dataclasses import dataclass, field from typing import Sequence, Union import torch @@ -11,11 +12,62 @@ LogitsProcessor, _load_custom_logitsprocs, ) -from vllm.v1.sample.logits_processor.state import LogitsProcessors +from vllm.v1.sample.logits_processor.state import BatchUpdateBuilder, LogitsProcessors logger = init_logger(__name__) +@dataclass(frozen=True) +class SpyreBatchUpdate(BatchUpdate): + """Extends BatchUpdate with pause/resume events for chunked-prefill holdback.""" + + # (dense_index, req_id) pairs — request was temporarily removed from the + # active batch; its logitproc state should be saved and not destroyed. + paused: list[tuple[int, str]] = field(default_factory=list) + # (dense_index, req_id) pairs — request is returning to the active batch; + # its previously saved logitproc state should be restored at dense_index. + resumed: list[tuple[int, str]] = field(default_factory=list) + + +class SpyreBatchUpdateBuilder(BatchUpdateBuilder): + """Extends BatchUpdateBuilder with pause/resume tracking.""" + + def __init__(self) -> None: + super().__init__() + self._paused: list[tuple[int, str]] = [] + self._resumed: list[tuple[int, str]] = [] + + def pause_append(self, dense_index: int, req_id: str) -> None: + self._paused.append((dense_index, req_id)) + + def resume_append(self, dense_index: int, req_id: str) -> None: + self._resumed.append((dense_index, req_id)) + + def get_and_reset(self, batch_size: int) -> SpyreBatchUpdate | None: + paused, self._paused = self._paused, [] + resumed, self._resumed = self._resumed, [] + base = super().get_and_reset(batch_size) + if base is None and not paused and not resumed: + return None + if base is None: + return SpyreBatchUpdate( + batch_size=batch_size, + removed=[], + added=[], + moved=[], + paused=paused, + resumed=resumed, + ) + return SpyreBatchUpdate( + batch_size=base.batch_size, + removed=base.removed, + added=base.added, + moved=base.moved, + paused=paused, + resumed=resumed, + ) + + def build_logitsprocs_for_cb( vllm_config: "VllmConfig", device: torch.device, @@ -43,7 +95,12 @@ def build_logitsprocs_for_cb( class LogitProcessorWrapper(LogitsProcessor): - """Logit processor to inject expected token during generation for tests""" + """Per-request logits processor manager for the persistent CB batch. + + Maintains one inner LogitsProcessor instance per active dense slot. + Pause/resume events save and restore exact per-request state so that + temporarily held-back requests do not lose their generation history. + """ def __init__( self, @@ -53,29 +110,26 @@ def __init__( is_pin_memory: bool, batch_size: int, ): - self.logitprocs: list[LogitsProcessor] = [ - logit_processor(vllm_config, device, is_pin_memory) for _ in range(batch_size) - ] - + self._factory = lambda: logit_processor(vllm_config, device, is_pin_memory) + self.logitprocs: list[LogitsProcessor] = [self._factory() for _ in range(batch_size)] + # Saved logitproc objects for paused requests, keyed by req_id. + self._saved: dict[str, LogitsProcessor] = {} self._is_argmax_invariant: bool = self.logitprocs[0].is_argmax_invariant() - self._prefill_index: int | None = None def is_argmax_invariant(self) -> bool: - """Never impacts greedy sampling""" return self._is_argmax_invariant - def update_state(self, batch_update: BatchUpdate | None): - # This method keeps the indices consistent of request while the - # persistent batch is changing. - - # Some LogitsProcessors, eg. MinTokensLogitsProcessor, require - # update_state to be called even if batch_update is None + def update_state(self, batch_update: BatchUpdate | None) -> None: + # Some LogitsProcessors (e.g. MinTokensLogitsProcessor) require + # update_state to be called even when batch_update is None. update_called = {i: False for i in range(len(self.logitprocs))} if batch_update is not None: for index, params, prompt_tok_ids, out_tok_ids in batch_update.added: update_called[index] = True + if self.logitprocs[index] is None: + self.logitprocs[index] = self._factory() self.logitprocs[index].update_state( BatchUpdate( batch_size=1, @@ -91,14 +145,23 @@ def update_state(self, batch_update: BatchUpdate | None): BatchUpdate(batch_size=1, removed=[0], moved=[], added=[]) ) + for index, req_id in getattr(batch_update, "resumed", []): + if req_id in self._saved: + self.logitprocs[index] = self._saved.pop(req_id) + + for index, req_id in getattr(batch_update, "paused", []): + self._saved[req_id] = self.logitprocs[index] + self.logitprocs[index] = None + for adx, bdx, _ in batch_update.moved: update_called[adx], update_called[bdx] = update_called[bdx], update_called[adx] = ( self.logitprocs[adx], self.logitprocs[bdx], ) = self.logitprocs[bdx], self.logitprocs[adx] - for index in [i for i, called in update_called.items() if not called]: - self.logitprocs[index].update_state(None) + for index, called in update_called.items(): + if not called and self.logitprocs[index] is not None: + self.logitprocs[index].update_state(None) def apply(self, logits: torch.Tensor) -> torch.Tensor: if self._prefill_index is not None: @@ -106,10 +169,8 @@ def apply(self, logits: torch.Tensor) -> torch.Tensor: self._prefill_index = None return logits - batch_size = logits.shape[0] - for i in range(batch_size): + for i in range(logits.shape[0]): logits[i] = self.logitprocs[i].apply(logits[i].unsqueeze(0)) - return logits def set_prefill_index(self, idx: int) -> None: diff --git a/sendnn_inference/v1/worker/spyre_input_batch.py b/sendnn_inference/v1/worker/spyre_input_batch.py index b22b434bb..831045ceb 100644 --- a/sendnn_inference/v1/worker/spyre_input_batch.py +++ b/sendnn_inference/v1/worker/spyre_input_batch.py @@ -12,10 +12,13 @@ from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams, SamplingType from vllm.v1.pool.metadata import PoolingMetadata -from vllm.v1.sample.logits_processor import BatchUpdateBuilder, LogitsProcessors, MoveDirectionality +from vllm.v1.sample.logits_processor import LogitsProcessors, MoveDirectionality from vllm.v1.sample.metadata import SamplingMetadata -from sendnn_inference.v1.sample.spyre_logits_processor import LogitProcessorWrapper +from sendnn_inference.v1.sample.spyre_logits_processor import ( + LogitProcessorWrapper, + SpyreBatchUpdateBuilder, +) class RequestState(Protocol): @@ -291,7 +294,7 @@ def __init__( # Internal representation of per-step batch state changes, used for # reordering persistent batch and generating logitsprocs batch state # updates. Should reset each step. - self.batch_update_builder = BatchUpdateBuilder() + self.batch_update_builder = SpyreBatchUpdateBuilder() self.logitsprocs = logitsprocs or LogitsProcessors() self.logitsprocs_wrappers = [ @@ -505,7 +508,7 @@ def remove_request(self, req_id: str): # Remove and move up self.batch_update_builder.removed_append(dense_index) - end_dense_idx = min(self._num_requests + 1, self.max_num_reqs - 1) + end_dense_idx = min(self._num_requests, self.max_num_reqs - 1) for tmp_dense in range(dense_index, end_dense_idx): self.batch_update_builder.moved.append( (tmp_dense, tmp_dense + 1, MoveDirectionality.UNIDIRECTIONAL) @@ -532,6 +535,146 @@ def remove_request(self, req_id: str): self.bad_words_token_ids.pop(req_index, None) + def pause_request(self, req_id: str) -> None: + """Temporarily remove a request from the active batch. + + Unlike remove_request: + - No 'removed' batch-update event is emitted, so LogitProcessorWrapper + saves the exact logitproc state rather than destroying it. + - The slot is freed (cleared from _req_ids) so new requests can use it. + """ + # Pop from id map and clear the slot + req_index = self.req_id_to_index.pop(req_id, None) + if req_index is None: + return + + # Must compute dense_index before masking + dense_index = self.req_idx_to_dense_index(req_index) + + # Free the slot for new requests + self._req_ids[req_index] = None + self.req_indices_mask[req_index] = False + self._num_requests -= 1 + + # Tell LogitProcessorWrapper to save state at this dense position. + self.batch_update_builder.pause_append(dense_index, req_id) + + # Emit shift moves so every other processor (MinP, LogitBias, …) sees + # the correct dense-index after the gap is closed. + end_dense_idx = min(self._num_requests, self.max_num_reqs - 1) + for tmp_dense in range(dense_index, end_dense_idx): + self.batch_update_builder.moved.append( + (tmp_dense, tmp_dense + 1, MoveDirectionality.UNIDIRECTIONAL) + ) + + self.req_output_token_ids.pop(dense_index) + + self.greedy_reqs.discard(req_id) + self.random_reqs.discard(req_id) + self.top_p_reqs.discard(req_id) + self.top_k_reqs.discard(req_id) + + self.frequency_penalties_reqs.discard(req_id) + self.presence_penalties_reqs.discard(req_id) + self.repetition_penalties_reqs.discard(req_id) + self.generators.pop(req_index, None) + self.num_logprobs.pop(req_id, None) + + self.has_allowed_token_ids.discard(req_id) + + if self.allowed_token_ids_mask is not None: + self.allowed_token_ids_mask[req_index].fill_(False) + + self.bad_words_token_ids.pop(req_index, None) + + def resume_request(self, req_id: str, request: "SamplingRequestState") -> None: + """Restore a previously paused request to the active batch. + + Emits an 'added' event so builtin processors (MinP, LogitBias, …) + re-register the request's sampling params. The accompanying 'resumed' + event then tells LogitProcessorWrapper to overwrite that freshly + initialised slot with the exact saved state, preserving history. + """ + # Get an available slot (same as add_request) + req_index = self.get_available_index() + assert req_index is not None + assert req_index < self.max_num_reqs + + # Set up the slot + self._req_ids[req_index] = req_id + self.req_indices_mask[req_index] = True + self.req_id_to_index[req_id] = req_index + self._num_requests += 1 + + # Copy prompt and output token ids + num_prompt_tokens = len(request.prompt_token_ids) + self.num_prompt_tokens[req_index] = num_prompt_tokens + self.token_ids_cpu[req_index, :num_prompt_tokens] = request.prompt_token_ids + start_idx = num_prompt_tokens + end_idx = start_idx + len(request.output_token_ids) + self.token_ids_cpu[req_index, start_idx:end_idx] = request.output_token_ids + + dense_index = self.req_idx_to_dense_index(req_index) + self.req_output_token_ids.insert(dense_index, request.output_token_ids) + + # Tell LogitProcessorWrapper to restore the saved state at dense_index. + # No 'added' event needed - the saved LogitsProcessor instance already + # contains all inner processor states with correct history. + tmp_dense = self._num_requests - 1 + self.batch_update_builder.resume_append(tmp_dense, req_id) + + # Bubble whatever is at the tail position to the correct dense position. + # This maintains correct indexing for all logits processors. + while tmp_dense > dense_index: + self.batch_update_builder.moved.append( + (tmp_dense, tmp_dense - 1, MoveDirectionality.SWAP) + ) + tmp_dense -= 1 + + sampling_params = request.sampling_params + if sampling_params.sampling_type == SamplingType.GREEDY: + self.temperature_cpu[req_index] = -1.0 + self.greedy_reqs.add(req_id) + else: + self.temperature_cpu[req_index] = sampling_params.temperature + self.random_reqs.add(req_id) + + self.top_p_cpu[req_index] = sampling_params.top_p + if sampling_params.top_p < 1: + self.top_p_reqs.add(req_id) + top_k = sampling_params.top_k + if 0 < top_k < self.vocab_size: + self.top_k_reqs.add(req_id) + else: + top_k = self.vocab_size + self.top_k_cpu[req_index] = top_k + self.frequency_penalties_cpu[req_index] = sampling_params.frequency_penalty + if sampling_params.frequency_penalty != 0.0: + self.frequency_penalties_reqs.add(req_id) + self.presence_penalties_cpu[req_index] = sampling_params.presence_penalty + if sampling_params.presence_penalty != 0.0: + self.presence_penalties_reqs.add(req_id) + self.repetition_penalties_cpu[req_index] = sampling_params.repetition_penalty + if sampling_params.repetition_penalty != 1.0: + self.repetition_penalties_reqs.add(req_id) + + if request.generator is not None: + self.generators[req_index] = request.generator + + if sampling_params.logprobs is not None: + self.num_logprobs[req_id] = sampling_params.logprobs + + if sampling_params.allowed_token_ids: + self.has_allowed_token_ids.add(req_id) + if self.allowed_token_ids_mask is None: + self.allowed_token_ids_mask = torch.zeros( + self.max_num_reqs, self.vocab_size, dtype=torch.bool, device=self.device + ) + self.allowed_token_ids_mask[req_index][sampling_params.allowed_token_ids] = True + + if sampling_params.bad_words_token_ids: + self.bad_words_token_ids[req_index] = sampling_params.bad_words_token_ids + def refresh_metadata(self): """Apply batch updates, reset input batch at end of step diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index bca4151cc..1b4f94b8d 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -1447,12 +1447,12 @@ def _update_batch(self, scheduler_output: SchedulerOutput): # Find requests that are in input_batch but not in scheduler output (paused) paused_req_ids = current_batch_req_ids - scheduled_req_ids for req_id in paused_req_ids: - # Only remove if it's not a finished request (finished requests are handled separately) + # Only pause if it's not a finished request (finished requests are handled separately) if req_id not in (scheduler_output.finished_req_ids or []): - logger.info("Removing paused request %s from input_batch", req_id) - self.input_batch.remove_request(req_id) - # Track that this request was paused + logger.info("Pausing request %s from input_batch", req_id) + self.input_batch.pause_request(req_id) self.paused_req_ids.add(req_id) + self.input_batch.refresh_metadata() # Find requests that are in scheduler output but not in input_batch # (restore from pausing) @@ -1462,9 +1462,9 @@ def _update_batch(self, scheduler_output: SchedulerOutput): if req_id in self.paused_req_ids and req_id in self.requests: logger.info("Restoring paused request %s to input_batch", req_id) req_state = self.requests[req_id] - self.input_batch.add_request(req_state) - # Remove from paused tracking since it's now restored + self.input_batch.resume_request(req_id, req_state) self.paused_req_ids.discard(req_id) + self.input_batch.refresh_metadata() for i, req_id in enumerate(req_data.req_ids): req_state: SamplingRequestState = self.requests[req_id] diff --git a/tests/e2e/test_logits_processors.py b/tests/e2e/test_logits_processors.py index b973487a9..ca608f2cd 100644 --- a/tests/e2e/test_logits_processors.py +++ b/tests/e2e/test_logits_processors.py @@ -5,6 +5,7 @@ from vllm import LLM, SamplingParams from vllm.config import VllmConfig from vllm.v1.sample.logits_processor import BatchUpdate, LogitsProcessor, MoveDirectionality +from sendnn_inference.v1.sample.spyre_logits_processor import SpyreBatchUpdate def test_custom_logits_processor( @@ -58,7 +59,7 @@ def apply(self, logits: torch.Tensor) -> torch.Tensor: # TODO: validate that this test case is valid for chunked prefill -def test_logits_processor(model: ModelInfo, backend, monkeypatch, max_model_len, mode: str): +def test_logits_processor_cp(model: ModelInfo, backend, monkeypatch, max_model_len, mode: str): """ Test if the state of logits processors are correct due to the switch of prefill/decode in a step engine. The LLM is initialized with bs=2, @@ -84,6 +85,8 @@ class SpyLogitsProcessor(LogitsProcessor): def __init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool): self.req_info: dict[int, SamplingParams] = {} + # Saved state for paused requests, keyed by req_id. + self._paused_info: dict[str, SamplingParams] = {} def is_argmax_invariant(self) -> bool: return False @@ -95,7 +98,19 @@ def update_state(self, batch_update: BatchUpdate | None): for index, params, _, _ in batch_update.added: self.req_info[index] = params nonlocal spy_outputs - spy_outputs[params.max_tokens] = [] + # Use setdefault so a resume-triggered added event does not + # reset tokens already collected before the pause. + spy_outputs.setdefault(params.max_tokens, []) + + if SpyreBatchUpdate is not None: + for dense_index, req_id in getattr(batch_update, "resumed", []): + if req_id in self._paused_info: + self.req_info[dense_index] = self._paused_info.pop(req_id) + + if SpyreBatchUpdate is not None: + for dense_index, req_id in getattr(batch_update, "paused", []): + if dense_index in self.req_info: + self._paused_info[req_id] = self.req_info.pop(dense_index) if self.req_info: # Process removed requests. @@ -137,7 +152,7 @@ def apply(self, logits: torch.Tensor) -> torch.Tensor: max_num_batched_tokens=128, enable_prefix_caching=mode == "pc", ) - prompt = ["Hello Logits Processors"] * 3 + prompt = ["1 2 3 4 5 6 7 8 9 " * 10] * 3 params0 = SamplingParams(max_tokens=5, temperature=0, logprobs=0, ignore_eos=True) params1 = SamplingParams(max_tokens=10, temperature=0, logprobs=0, ignore_eos=True) params2 = SamplingParams(max_tokens=7, temperature=0, logprobs=0, ignore_eos=True) @@ -151,3 +166,519 @@ def apply(self, logits: torch.Tensor) -> torch.Tensor: assert spy_outputs[5] == outputs[0].outputs[0].token_ids assert spy_outputs[10] == outputs[1].outputs[0].token_ids assert spy_outputs[7] == outputs[2].outputs[0].token_ids + + +def test_logits_processor_advanced( + model: ModelInfo, backend, monkeypatch, max_model_len, mode: str +): + """ + Complex test for logits processor state management with controlled SchedulerOutput. + + Tests multiple simultaneous operations: + - Adding new requests while finishing others + - Pausing and resuming requests + - Verifying correct index management + - Ensuring no state overwrites occur + + This test simulates various scheduler scenarios where requests can be: + 1. Added and finished in the same step + 2. Paused and resumed + 3. Resumed while others finish + 4. Multiple operations happening simultaneously + """ + from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput + from vllm.v1.request import Request + from tests.v1.worker.mock_model import InstrumentedModelRunner + + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + + # Track all state transitions and verify correctness + state_log: list[dict] = [] + + class StateTrackingLogitsProcessor(LogitsProcessor): + """ + Tracks the complete state of the batch at each update. + Verifies that indices are correct and no overwrites occur. + """ + + def __init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool): + # Map dense_index -> (req_id, token_count) + self.active_requests: dict[int, tuple[str, int]] = {} + # Map req_id -> (last_dense_index, token_count) for paused requests + self.paused_requests: dict[str, tuple[int, int]] = {} + self.step_count = 0 + + def is_argmax_invariant(self) -> bool: + return False + + def update_state(self, batch_update: BatchUpdate | None): + if not batch_update: + return + + self.step_count += 1 + step_info = { + "step": self.step_count, + "batch_size": batch_update.batch_size, + "operations": [], + "active_before": dict(self.active_requests), + "paused_before": dict(self.paused_requests), + } + + # Process added requests + for dense_index, params, prompt_toks, output_toks in batch_update.added: + req_id = f"req_{params.max_tokens}" # Use max_tokens as identifier + token_count = len(output_toks) + + # Verify no overwrite + if dense_index in self.active_requests: + old_req = self.active_requests[dense_index] + step_info["operations"].append( + { + "type": "ERROR_OVERWRITE", + "index": dense_index, + "old_req": old_req, + "new_req": (req_id, token_count), + } + ) + + self.active_requests[dense_index] = (req_id, token_count) + step_info["operations"].append( + {"type": "added", "index": dense_index, "req_id": req_id, "tokens": token_count} + ) + + # Process resumed requests (restore from pause) + if hasattr(batch_update, "resumed"): + for dense_index, req_id in batch_update.resumed: + if req_id in self.paused_requests: + old_index, token_count = self.paused_requests.pop(req_id) + self.active_requests[dense_index] = (req_id, token_count) + step_info["operations"].append( + { + "type": "resumed", + "index": dense_index, + "req_id": req_id, + "old_index": old_index, + "tokens": token_count, + } + ) + + # Process paused requests (save state) + if hasattr(batch_update, "paused"): + for dense_index, req_id in batch_update.paused: + if dense_index in self.active_requests: + req_info = self.active_requests.pop(dense_index) + self.paused_requests[req_id] = (dense_index, req_info[1]) + step_info["operations"].append( + { + "type": "paused", + "index": dense_index, + "req_id": req_id, + "tokens": req_info[1], + } + ) + + # Process removed requests + for dense_index in batch_update.removed: + if dense_index in self.active_requests: + req_info = self.active_requests.pop(dense_index) + step_info["operations"].append( + { + "type": "removed", + "index": dense_index, + "req_id": req_info[0], + "tokens": req_info[1], + } + ) + + # Process moved requests (always swaps like LogitProcessorWrapper) + for src_idx, dst_idx, _ in batch_update.moved: + src_req = self.active_requests.get(src_idx) + dst_req = self.active_requests.get(dst_idx) + + # Always swap both positions (matching LogitProcessorWrapper behavior) + if src_req: + self.active_requests[dst_idx] = src_req + else: + self.active_requests.pop(dst_idx, None) + if dst_req: + self.active_requests[src_idx] = dst_req + else: + self.active_requests.pop(src_idx, None) + + step_info["operations"].append( + { + "type": "moved", + "src": src_idx, + "dst": dst_idx, + "src_req": src_req, + "dst_req": dst_req, + } + ) + + step_info["active_after"] = dict(self.active_requests) + step_info["paused_after"] = dict(self.paused_requests) + + # Verify indices are consecutive immediately after each update + if self.active_requests: + indices = sorted(self.active_requests.keys()) + expected = list(range(len(indices))) + if indices != expected: + step_info["operations"].append( + { + "type": "ERROR_NON_CONSECUTIVE", + "actual_indices": indices, + "expected_indices": expected, + "message": ( + f"Indices are not consecutive: {indices}, expected {expected}" + ), + } + ) + + nonlocal state_log + state_log.append(step_info) + + def apply(self, logits: torch.Tensor) -> torch.Tensor: + # Increment token count for all active requests + for idx in self.active_requests: + req_id, token_count = self.active_requests[idx] + self.active_requests[idx] = (req_id, token_count + 1) + return logits + + patch_environment( + backend=backend, + monkeypatch=monkeypatch, + ) + + # Build the model runner with our tracking processor + runner = InstrumentedModelRunner.build( + monkeypatch=monkeypatch, + enable_prefix_caching=mode == "pc", + model_name=model.name, + max_num_seqs=4, # Allow up to 4 concurrent requests + max_model_len=max_model_len, + max_num_batched_tokens=128, + ) + + # Replace logits processors with our tracking one + from sendnn_inference.v1.sample.spyre_logits_processor import build_logitsprocs_for_cb + + runner.input_batch.logitsprocs = build_logitsprocs_for_cb( + vllm_config=runner.vllm_config, + device=runner.device, + is_pin_memory=runner.pin_memory, + is_pooling_model=False, + batch_size=4, + custom_logitsprocs=[StateTrackingLogitsProcessor], + ) + + # Helper to create requests + def make_request(req_id: str, prompt_len: int, max_tokens: int) -> Request: + return Request( + request_id=req_id, + prompt_token_ids=[42] * prompt_len, + sampling_params=SamplingParams(max_tokens=max_tokens, temperature=0), + pooling_params=None, + ) + + # Helper to create new request data + def make_new_req_data(req_id: str, prompt_len: int, max_tokens: int) -> NewRequestData: + req = make_request(req_id, prompt_len, max_tokens) + block_ids = list(range(1, (prompt_len + 63) // 64 + 1)) + return NewRequestData.from_request(req, block_ids=(block_ids,)) + + # Helper to create cached request data + def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> CachedRequestData: + cached = CachedRequestData.make_empty() + cached.req_ids = list(req_states.keys()) + cached.num_computed_tokens = [state[0] for state in req_states.values()] + cached.new_block_ids = [None for _ in req_states] # Simplified for test + return cached + + # Add first request: req 0 + req1 = make_new_req_data("req1", 50, 5) + sched_out = SchedulerOutput( + scheduled_new_reqs=[req1], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={"req1": 50}, + total_num_scheduled_tokens=50, + finished_req_ids=set(), + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + runner.execute_model(sched_out) + + # Decode request 0 + cached = make_cached_req_data({"req1": (50, [])}) + sched_out = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=cached, + num_scheduled_tokens={"req1": 1}, + total_num_scheduled_tokens=1, + finished_req_ids=set(), + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + runner.execute_model(sched_out) + + # Add second request (long one, needs three chunks) + # Chunked-prefill 1/3 of request 2 + req2 = make_new_req_data("req2", 266, 10) + sched_out = SchedulerOutput( + scheduled_new_reqs=[req2], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={"req2": 128}, + total_num_scheduled_tokens=128, + finished_req_ids=set(), + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + runner.execute_model(sched_out) + + # Decode request 0 + cached = make_cached_req_data({"req1": (51, [])}) + sched_out = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=cached, + num_scheduled_tokens={"req1": 1}, + total_num_scheduled_tokens=1, + finished_req_ids=set(), + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + runner.execute_model(sched_out) + + # Chunked-prefill 2/3 of request 2 + cached = make_cached_req_data({"req2": (128, [])}) + sched_out = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=cached, + num_scheduled_tokens={"req2": 128}, + total_num_scheduled_tokens=128, + finished_req_ids=set(), + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + runner.execute_model(sched_out) + + # Decode request 0 + cached = make_cached_req_data({"req1": (52, [])}) + sched_out = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=cached, + num_scheduled_tokens={"req1": 1}, + total_num_scheduled_tokens=1, + finished_req_ids=set(), + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + runner.execute_model(sched_out) + + # Chunked-prefill 3/3 of request 2 + cached = make_cached_req_data({"req2": (256, [])}) + sched_out = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=cached, + num_scheduled_tokens={"req2": 138}, + total_num_scheduled_tokens=138, + finished_req_ids=set(), + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + runner.execute_model(sched_out) + + # Decode requests 1 and 2 + cached = make_cached_req_data({"req1": (53, []), "req2": (266, [])}) + sched_out = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=cached, + num_scheduled_tokens={"req1": 1, "req2": 1}, + total_num_scheduled_tokens=2, + finished_req_ids=set(), + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + runner.execute_model(sched_out) + + # Now finish req1 and add req3 in the same step + req3 = make_new_req_data("req3", 50, 7) + sched_out = SchedulerOutput( + scheduled_new_reqs=[req3], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={"req3": 50}, + total_num_scheduled_tokens=50, + finished_req_ids={"req1"}, + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + runner.execute_model(sched_out) + + # Verify no overwrites occurred + for step in state_log: + for op in step["operations"]: + assert op["type"] != "ERROR_OVERWRITE", ( + f"Overwrite detected at step {step['step']}: {op}" + ) + + # Clean up + cached = make_cached_req_data({"req2": (267, []), "req3": (51, [])}) + sched_out = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=cached, + num_scheduled_tokens={}, + total_num_scheduled_tokens=0, + finished_req_ids={"req2", "req3"}, + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + runner.execute_model(sched_out) + + # Start fresh with req4 + req4 = make_new_req_data("req4", 50, 8) + sched_out = SchedulerOutput( + scheduled_new_reqs=[req4], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={"req4": 50}, + total_num_scheduled_tokens=50, + finished_req_ids=set(), + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + runner.execute_model(sched_out) + + # Decode request 4 + cached = make_cached_req_data({"req4": (50, [])}) + sched_out = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=cached, + num_scheduled_tokens={"req4": 1}, + total_num_scheduled_tokens=1, + finished_req_ids=set(), + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + runner.execute_model(sched_out) + + # Prefill request 5 + req5 = make_new_req_data("req5", 50, 12) + sched_out = SchedulerOutput( + scheduled_new_reqs=[req5], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={"req5": 50}, + total_num_scheduled_tokens=50, + finished_req_ids=set(), + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + runner.execute_model(sched_out) + + # Simulate pause of req4 (scheduler removes it temporarily) + # In real scenario, scheduler would not include req4 in cached_reqs + cached = make_cached_req_data({"req5": (50, [])}) + sched_out = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=cached, + num_scheduled_tokens={"req5": 1}, + total_num_scheduled_tokens=1, + finished_req_ids=set(), + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + runner.execute_model(sched_out) + + # Resume req4 and finish req5 simultaneously + cached = make_cached_req_data({"req4": (51, []), "req5": (51, [])}) + sched_out = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=cached, + num_scheduled_tokens={"req4": 1, "req5": 1}, + total_num_scheduled_tokens=2, + finished_req_ids={"req5"}, + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + runner.execute_model(sched_out) + + # Final verification + print("\n=== State Log Summary ===") + for step in state_log: + print(f"\nStep {step['step']}:") + print(f" Batch size: {step['batch_size']}") + print(f" Operations: {len(step['operations'])}") + for op in step["operations"]: + print(f" - {op['type']}: {op}") + print(f" Active after: {step['active_after']}") + print(f" Paused after: {step['paused_after']}") + + # Verify no errors occurred during state transitions + errors = [ + (step["step"], op) + for step in state_log + for op in step["operations"] + if op["type"].startswith("ERROR") + ] + + if errors: + print("\n=== ERRORS DETECTED ===") + for step_num, error_op in errors: + print(f"Step {step_num}: {error_op['type']}") + print(f" Details: {error_op}") + + assert len(errors) == 0, f"Found {len(errors)} errors in state transitions: {errors}" + + # Additional verification: all steps should have consecutive indices + for step in state_log: + if step["active_after"]: + indices = sorted(step["active_after"].keys()) + expected = list(range(len(indices))) + assert indices == expected, ( + f"Step {step['step']}: Non-contiguous indices {indices}, expected {expected}" + ) + + print("\n=== Test Passed ===") + print(f"Total steps: {len(state_log)}") + print(f"Total operations: {sum(len(s['operations']) for s in state_log)}") diff --git a/tests/v1/worker/test_spyre_input_batch.py b/tests/v1/worker/test_spyre_input_batch.py index 9e5f8afb7..c6031a1be 100644 --- a/tests/v1/worker/test_spyre_input_batch.py +++ b/tests/v1/worker/test_spyre_input_batch.py @@ -7,9 +7,18 @@ from vllm.sampling_params import SamplingParams from vllm.utils.platform_utils import is_pin_memory_available from vllm.utils.torch_utils import make_tensor_with_pad -from vllm.v1.sample.logits_processor import LogitsProcessors +from vllm.v1.sample.logits_processor import ( + BatchUpdate, + LogitsProcessor, + LogitsProcessors, + MoveDirectionality, +) from vllm.v1.sample.metadata import SamplingMetadata +from sendnn_inference.v1.sample.spyre_logits_processor import ( + LogitProcessorWrapper, + SpyreBatchUpdate, +) from sendnn_inference.v1.worker.spyre_input_batch import SamplingInputBatch, SamplingRequestState VOCAB_SIZE = 1024 @@ -280,3 +289,127 @@ def test_sampling_metadata_topk_edges(): assert input_batch.top_k[0] == VOCAB_SIZE assert input_batch.top_k[1] == VOCAB_SIZE + + +@pytest.mark.cpu +@pytest.mark.worker +def test_logitproc_wrapper_pause_resume(): + """ + Verifies that LogitProcessorWrapper saves and restores the exact logitproc + object across pause → (add new request) → resume, guaranteeing that: + - B's per-request state is not overwritten by the new request D. + - After resume, wrapper.logitprocs[dense_index] is the *same Python object* + that existed before the pause (identity check), so all accumulated state + (grammar progress, token counts, etc.) is intact. + + The batch events used here are exactly what pause_request / resume_request + emit for a 4-slot batch with A(slot=0), B(slot=1), C(slot=2), D(slot=3): + + pause B → paused(1,"B"), UNIDIRECTIONAL(1,2), UNIDIRECTIONAL(2,3) + add D → added(2, D_params) + resume B → added(3, B_params), SWAP(3,2), SWAP(2,1), resumed(1,"B") + """ + + class TrackingLogitsProcessor(LogitsProcessor): + """Minimal processor that records which request it belongs to.""" + + def __init__(self, vllm_config, device, is_pin_memory): + self.label: str | None = None + + def is_argmax_invariant(self) -> bool: + return True + + def update_state(self, batch_update: BatchUpdate | None) -> None: + if batch_update is None: + return + for _, params, _, _ in batch_update.added: + self.label = f"t{params.max_tokens}" + for _ in batch_update.removed: + self.label = None + + def apply(self, logits: torch.Tensor) -> torch.Tensor: + return logits + + wrapper = LogitProcessorWrapper( + logit_processor=TrackingLogitsProcessor, + vllm_config=None, + device=torch.device("cpu"), + is_pin_memory=False, + batch_size=4, + ) + + def p(max_tokens: int) -> SamplingParams: + return SamplingParams(max_tokens=max_tokens, temperature=0) + + # ── Step 1: add A(dense=0), B(dense=1), C(dense=2) ────────────────────── + wrapper.update_state( + SpyreBatchUpdate( + batch_size=3, + added=[(0, p(5), [], []), (1, p(10), [], []), (2, p(7), [], [])], + removed=[], + moved=[], + paused=[], + resumed=[], + ) + ) + assert wrapper.logitprocs[0].label == "t5" + assert wrapper.logitprocs[1].label == "t10" + assert wrapper.logitprocs[2].label == "t7" + + b_obj = wrapper.logitprocs[1] # hold ref to B's original logitproc object + + # ── Step 2: pause B (dense=1) ──────────────────────────────────────────── + # pause_request emits: paused(1,"B"), UNIDIRECTIONAL(1,2), UNIDIRECTIONAL(2,3) + wrapper.update_state( + SpyreBatchUpdate( + batch_size=2, + added=[], + removed=[], + moved=[ + (1, 2, MoveDirectionality.UNIDIRECTIONAL), + (2, 3, MoveDirectionality.UNIDIRECTIONAL), + ], + paused=[(1, "B")], + resumed=[], + ) + ) + assert wrapper._saved.get("B") is b_obj, "B's logitproc must be in _saved" + assert wrapper.logitprocs[0].label == "t5" # A unchanged at dense=0 + assert wrapper.logitprocs[1].label == "t7" # C shifted to dense=1 + + # ── Step 3: add D (dense=2) while B is paused ──────────────────────────── + # add_request(D) emits: added(2, D_params) — no moves since 2 == tmp_dense + wrapper.update_state( + SpyreBatchUpdate( + batch_size=3, + added=[(2, p(3), [], [])], + removed=[], + moved=[], + paused=[], + resumed=[], + ) + ) + assert wrapper.logitprocs[2].label == "t3" # D at dense=2 + assert wrapper._saved.get("B") is b_obj, "B's state must survive D being added" + + # ── Step 4: resume B (dense=1) ─────────────────────────────────────────── + # resume_request emits: added(3,B_params), SWAP(3,2), SWAP(2,1), resumed(1,"B") + wrapper.update_state( + SpyreBatchUpdate( + batch_size=4, + added=[], + removed=[], + moved=[ + (3, 2, MoveDirectionality.SWAP), + (2, 1, MoveDirectionality.SWAP), + ], + paused=[], + resumed=[(3, "B")], + ) + ) + assert wrapper.logitprocs[1] is b_obj, "B's exact logitproc object must be restored" + assert wrapper.logitprocs[0].label == "t5" # A at dense=0 + assert wrapper.logitprocs[1].label == "t10" # B at dense=1 (saved state) + assert wrapper.logitprocs[2].label == "t7" # C at dense=2 + assert wrapper.logitprocs[3].label == "t3" # D at dense=3 + assert "B" not in wrapper._saved # saved state consumed on resume From 2195185bac700d7950cac8fe04c2483d44306bd1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 2 Jun 2026 11:41:24 +0000 Subject: [PATCH 022/106] small bugs catched MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../v1/sample/spyre_logits_processor.py | 11 +- .../v1/worker/spyre_model_runner.py | 6 +- tests/e2e/test_logits_processors.py | 115 ++++++++++-------- 3 files changed, 76 insertions(+), 56 deletions(-) diff --git a/sendnn_inference/v1/sample/spyre_logits_processor.py b/sendnn_inference/v1/sample/spyre_logits_processor.py index a0865b2b2..1188d722a 100644 --- a/sendnn_inference/v1/sample/spyre_logits_processor.py +++ b/sendnn_inference/v1/sample/spyre_logits_processor.py @@ -146,18 +146,19 @@ def update_state(self, batch_update: BatchUpdate | None) -> None: ) for index, req_id in getattr(batch_update, "resumed", []): - if req_id in self._saved: - self.logitprocs[index] = self._saved.pop(req_id) + assert req_id in self._saved + self.logitprocs[index] = self._saved.pop(req_id) for index, req_id in getattr(batch_update, "paused", []): self._saved[req_id] = self.logitprocs[index] self.logitprocs[index] = None for adx, bdx, _ in batch_update.moved: - update_called[adx], update_called[bdx] = update_called[bdx], update_called[adx] = ( - self.logitprocs[adx], + update_called[adx], update_called[bdx] = update_called[bdx], update_called[adx] + self.logitprocs[adx], self.logitprocs[bdx] = ( self.logitprocs[bdx], - ) = self.logitprocs[bdx], self.logitprocs[adx] + self.logitprocs[adx], + ) for index, called in update_called.items(): if not called and self.logitprocs[index] is not None: diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index 1b4f94b8d..22f131c5b 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -1443,6 +1443,7 @@ def _update_batch(self, scheduler_output: SchedulerOutput): # requests where scheduler temporarily removes them from running queue scheduled_req_ids = set(req_data.req_ids) current_batch_req_ids = set(self.input_batch.req_id_to_index.keys()) + need_metadata_refresh = True # Find requests that are in input_batch but not in scheduler output (paused) paused_req_ids = current_batch_req_ids - scheduled_req_ids @@ -1453,6 +1454,7 @@ def _update_batch(self, scheduler_output: SchedulerOutput): self.input_batch.pause_request(req_id) self.paused_req_ids.add(req_id) self.input_batch.refresh_metadata() + need_metadata_refresh = False # Find requests that are in scheduler output but not in input_batch # (restore from pausing) @@ -1465,6 +1467,7 @@ def _update_batch(self, scheduler_output: SchedulerOutput): self.input_batch.resume_request(req_id, req_state) self.paused_req_ids.discard(req_id) self.input_batch.refresh_metadata() + need_metadata_refresh = False for i, req_id in enumerate(req_data.req_ids): req_state: SamplingRequestState = self.requests[req_id] @@ -1487,7 +1490,8 @@ def _update_batch(self, scheduler_output: SchedulerOutput): # of logitprocs. Refactor so that we can batch removals to the # `input_batch` self.input_batch.refresh_metadata() - else: + need_metadata_refresh = False + if need_metadata_refresh: # Due to logits processor we need to refresh metadata at each step self.input_batch.refresh_metadata() diff --git a/tests/e2e/test_logits_processors.py b/tests/e2e/test_logits_processors.py index ca608f2cd..676b2bb7e 100644 --- a/tests/e2e/test_logits_processors.py +++ b/tests/e2e/test_logits_processors.py @@ -395,11 +395,11 @@ def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> Cached return cached # Add first request: req 0 - req1 = make_new_req_data("req1", 50, 5) + req0 = make_new_req_data("req0", 50, 5) sched_out = SchedulerOutput( - scheduled_new_reqs=[req1], + scheduled_new_reqs=[req0], scheduled_cached_reqs=CachedRequestData.make_empty(), - num_scheduled_tokens={"req1": 50}, + num_scheduled_tokens={"req0": 50}, total_num_scheduled_tokens=50, finished_req_ids=set(), kv_connector_metadata=None, @@ -411,11 +411,11 @@ def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> Cached runner.execute_model(sched_out) # Decode request 0 - cached = make_cached_req_data({"req1": (50, [])}) + cached = make_cached_req_data({"req0": (50, [])}) sched_out = SchedulerOutput( scheduled_new_reqs=[], scheduled_cached_reqs=cached, - num_scheduled_tokens={"req1": 1}, + num_scheduled_tokens={"req0": 1}, total_num_scheduled_tokens=1, finished_req_ids=set(), kv_connector_metadata=None, @@ -427,12 +427,12 @@ def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> Cached runner.execute_model(sched_out) # Add second request (long one, needs three chunks) - # Chunked-prefill 1/3 of request 2 - req2 = make_new_req_data("req2", 266, 10) + # Chunked-prefill 1/3 of request 1 + req1 = make_new_req_data("req1", 266, 10) sched_out = SchedulerOutput( - scheduled_new_reqs=[req2], + scheduled_new_reqs=[req1], scheduled_cached_reqs=CachedRequestData.make_empty(), - num_scheduled_tokens={"req2": 128}, + num_scheduled_tokens={"req1": 128}, total_num_scheduled_tokens=128, finished_req_ids=set(), kv_connector_metadata=None, @@ -444,11 +444,11 @@ def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> Cached runner.execute_model(sched_out) # Decode request 0 - cached = make_cached_req_data({"req1": (51, [])}) + cached = make_cached_req_data({"req0": (51, [])}) sched_out = SchedulerOutput( scheduled_new_reqs=[], scheduled_cached_reqs=cached, - num_scheduled_tokens={"req1": 1}, + num_scheduled_tokens={"req0": 1}, total_num_scheduled_tokens=1, finished_req_ids=set(), kv_connector_metadata=None, @@ -459,12 +459,12 @@ def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> Cached ) runner.execute_model(sched_out) - # Chunked-prefill 2/3 of request 2 - cached = make_cached_req_data({"req2": (128, [])}) + # Chunked-prefill 2/3 of request 1 + cached = make_cached_req_data({"req1": (128, [])}) sched_out = SchedulerOutput( scheduled_new_reqs=[], scheduled_cached_reqs=cached, - num_scheduled_tokens={"req2": 128}, + num_scheduled_tokens={"req1": 128}, total_num_scheduled_tokens=128, finished_req_ids=set(), kv_connector_metadata=None, @@ -476,11 +476,11 @@ def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> Cached runner.execute_model(sched_out) # Decode request 0 - cached = make_cached_req_data({"req1": (52, [])}) + cached = make_cached_req_data({"req0": (52, [])}) sched_out = SchedulerOutput( scheduled_new_reqs=[], scheduled_cached_reqs=cached, - num_scheduled_tokens={"req1": 1}, + num_scheduled_tokens={"req0": 1}, total_num_scheduled_tokens=1, finished_req_ids=set(), kv_connector_metadata=None, @@ -491,12 +491,12 @@ def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> Cached ) runner.execute_model(sched_out) - # Chunked-prefill 3/3 of request 2 - cached = make_cached_req_data({"req2": (256, [])}) + # Chunked-prefill 3/3 of request 1 + cached = make_cached_req_data({"req1": (256, [])}) sched_out = SchedulerOutput( scheduled_new_reqs=[], scheduled_cached_reqs=cached, - num_scheduled_tokens={"req2": 138}, + num_scheduled_tokens={"req1": 138}, total_num_scheduled_tokens=138, finished_req_ids=set(), kv_connector_metadata=None, @@ -507,12 +507,12 @@ def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> Cached ) runner.execute_model(sched_out) - # Decode requests 1 and 2 - cached = make_cached_req_data({"req1": (53, []), "req2": (266, [])}) + # Decode requests 0 and 1 + cached = make_cached_req_data({"req0": (53, []), "req1": (266, [])}) sched_out = SchedulerOutput( scheduled_new_reqs=[], scheduled_cached_reqs=cached, - num_scheduled_tokens={"req1": 1, "req2": 1}, + num_scheduled_tokens={"req0": 1, "req1": 1}, total_num_scheduled_tokens=2, finished_req_ids=set(), kv_connector_metadata=None, @@ -523,12 +523,28 @@ def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> Cached ) runner.execute_model(sched_out) - # Now finish req1 and add req3 in the same step - req3 = make_new_req_data("req3", 50, 7) + # Decode request 0, pause request 1 + cached = make_cached_req_data({"req0": (54, [])}) sched_out = SchedulerOutput( - scheduled_new_reqs=[req3], + scheduled_new_reqs=[], + scheduled_cached_reqs=cached, + num_scheduled_tokens={"req0": 1}, + total_num_scheduled_tokens=1, + finished_req_ids=set(), + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + runner.execute_model(sched_out) + + # Now finish req1, pause req0, and add req2 + req2 = make_new_req_data("req2", 50, 7) + sched_out = SchedulerOutput( + scheduled_new_reqs=[req2], scheduled_cached_reqs=CachedRequestData.make_empty(), - num_scheduled_tokens={"req3": 50}, + num_scheduled_tokens={"req2": 50}, total_num_scheduled_tokens=50, finished_req_ids={"req1"}, kv_connector_metadata=None, @@ -546,14 +562,13 @@ def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> Cached f"Overwrite detected at step {step['step']}: {op}" ) - # Clean up - cached = make_cached_req_data({"req2": (267, []), "req3": (51, [])}) + # Clean up: finish remaining requests sched_out = SchedulerOutput( scheduled_new_reqs=[], - scheduled_cached_reqs=cached, + scheduled_cached_reqs=CachedRequestData.make_empty(), num_scheduled_tokens={}, total_num_scheduled_tokens=0, - finished_req_ids={"req2", "req3"}, + finished_req_ids={"req1", "req2"}, kv_connector_metadata=None, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, @@ -562,12 +577,12 @@ def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> Cached ) runner.execute_model(sched_out) - # Start fresh with req4 - req4 = make_new_req_data("req4", 50, 8) + # Start fresh with req3 + req3 = make_new_req_data("req3", 50, 8) sched_out = SchedulerOutput( - scheduled_new_reqs=[req4], + scheduled_new_reqs=[req3], scheduled_cached_reqs=CachedRequestData.make_empty(), - num_scheduled_tokens={"req4": 50}, + num_scheduled_tokens={"req3": 50}, total_num_scheduled_tokens=50, finished_req_ids=set(), kv_connector_metadata=None, @@ -578,12 +593,12 @@ def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> Cached ) runner.execute_model(sched_out) - # Decode request 4 - cached = make_cached_req_data({"req4": (50, [])}) + # Decode request 3 + cached = make_cached_req_data({"req3": (50, [])}) sched_out = SchedulerOutput( scheduled_new_reqs=[], scheduled_cached_reqs=cached, - num_scheduled_tokens={"req4": 1}, + num_scheduled_tokens={"req3": 1}, total_num_scheduled_tokens=1, finished_req_ids=set(), kv_connector_metadata=None, @@ -594,12 +609,12 @@ def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> Cached ) runner.execute_model(sched_out) - # Prefill request 5 - req5 = make_new_req_data("req5", 50, 12) + # Prefill request 4 + req4 = make_new_req_data("req4", 50, 12) sched_out = SchedulerOutput( - scheduled_new_reqs=[req5], + scheduled_new_reqs=[req4], scheduled_cached_reqs=CachedRequestData.make_empty(), - num_scheduled_tokens={"req5": 50}, + num_scheduled_tokens={"req4": 50}, total_num_scheduled_tokens=50, finished_req_ids=set(), kv_connector_metadata=None, @@ -610,13 +625,13 @@ def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> Cached ) runner.execute_model(sched_out) - # Simulate pause of req4 (scheduler removes it temporarily) - # In real scenario, scheduler would not include req4 in cached_reqs - cached = make_cached_req_data({"req5": (50, [])}) + # Simulate pause of req3 (scheduler removes it temporarily) + # In real scenario, scheduler would not include req3 in cached_reqs + cached = make_cached_req_data({"req4": (50, [])}) sched_out = SchedulerOutput( scheduled_new_reqs=[], scheduled_cached_reqs=cached, - num_scheduled_tokens={"req5": 1}, + num_scheduled_tokens={"req4": 1}, total_num_scheduled_tokens=1, finished_req_ids=set(), kv_connector_metadata=None, @@ -627,14 +642,14 @@ def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> Cached ) runner.execute_model(sched_out) - # Resume req4 and finish req5 simultaneously - cached = make_cached_req_data({"req4": (51, []), "req5": (51, [])}) + # Resume req3 and finish req4 simultaneously + cached = make_cached_req_data({"req3": (51, [])}) sched_out = SchedulerOutput( scheduled_new_reqs=[], scheduled_cached_reqs=cached, - num_scheduled_tokens={"req4": 1, "req5": 1}, - total_num_scheduled_tokens=2, - finished_req_ids={"req5"}, + num_scheduled_tokens={"req3": 1}, + total_num_scheduled_tokens=1, + finished_req_ids={"req4"}, kv_connector_metadata=None, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, From 41ee95e4eb79f6a60144702d3a8d3639848d0a10 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 2 Jun 2026 12:40:16 +0000 Subject: [PATCH 023/106] handle finish of paused requests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../v1/sample/spyre_logits_processor.py | 18 ++++++++++++++++-- .../v1/worker/spyre_input_batch.py | 4 ++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/sendnn_inference/v1/sample/spyre_logits_processor.py b/sendnn_inference/v1/sample/spyre_logits_processor.py index 1188d722a..5417da3b1 100644 --- a/sendnn_inference/v1/sample/spyre_logits_processor.py +++ b/sendnn_inference/v1/sample/spyre_logits_processor.py @@ -19,7 +19,7 @@ @dataclass(frozen=True) class SpyreBatchUpdate(BatchUpdate): - """Extends BatchUpdate with pause/resume events for chunked-prefill holdback.""" + """Extends BatchUpdate with pause/resume lifecycle events.""" # (dense_index, req_id) pairs — request was temporarily removed from the # active batch; its logitproc state should be saved and not destroyed. @@ -27,6 +27,9 @@ class SpyreBatchUpdate(BatchUpdate): # (dense_index, req_id) pairs — request is returning to the active batch; # its previously saved logitproc state should be restored at dense_index. resumed: list[tuple[int, str]] = field(default_factory=list) + # req_ids for requests that finished while paused and whose saved state + # should be discarded without being restored to an active slot. + finished_paused: list[str] = field(default_factory=list) class SpyreBatchUpdateBuilder(BatchUpdateBuilder): @@ -36,6 +39,7 @@ def __init__(self) -> None: super().__init__() self._paused: list[tuple[int, str]] = [] self._resumed: list[tuple[int, str]] = [] + self._finished_paused: list[str] = [] def pause_append(self, dense_index: int, req_id: str) -> None: self._paused.append((dense_index, req_id)) @@ -43,11 +47,15 @@ def pause_append(self, dense_index: int, req_id: str) -> None: def resume_append(self, dense_index: int, req_id: str) -> None: self._resumed.append((dense_index, req_id)) + def finished_paused_append(self, req_id: str) -> None: + self._finished_paused.append(req_id) + def get_and_reset(self, batch_size: int) -> SpyreBatchUpdate | None: paused, self._paused = self._paused, [] resumed, self._resumed = self._resumed, [] + finished_paused, self._finished_paused = self._finished_paused, [] base = super().get_and_reset(batch_size) - if base is None and not paused and not resumed: + if base is None and not paused and not resumed and not finished_paused: return None if base is None: return SpyreBatchUpdate( @@ -57,6 +65,7 @@ def get_and_reset(self, batch_size: int) -> SpyreBatchUpdate | None: moved=[], paused=paused, resumed=resumed, + finished_paused=finished_paused, ) return SpyreBatchUpdate( batch_size=base.batch_size, @@ -65,6 +74,7 @@ def get_and_reset(self, batch_size: int) -> SpyreBatchUpdate | None: moved=base.moved, paused=paused, resumed=resumed, + finished_paused=finished_paused, ) @@ -153,6 +163,10 @@ def update_state(self, batch_update: BatchUpdate | None) -> None: self._saved[req_id] = self.logitprocs[index] self.logitprocs[index] = None + for req_id in getattr(batch_update, "finished_paused", []): + assert req_id in self._saved + self._saved.pop(req_id) + for adx, bdx, _ in batch_update.moved: update_called[adx], update_called[bdx] = update_called[bdx], update_called[adx] self.logitprocs[adx], self.logitprocs[bdx] = ( diff --git a/sendnn_inference/v1/worker/spyre_input_batch.py b/sendnn_inference/v1/worker/spyre_input_batch.py index 831045ceb..60831cb0d 100644 --- a/sendnn_inference/v1/worker/spyre_input_batch.py +++ b/sendnn_inference/v1/worker/spyre_input_batch.py @@ -493,6 +493,10 @@ def remove_request(self, req_id: str): For the continuous batching, the removed request indices can be overwritten by new requests """ + req_index = self.req_id_to_index.get(req_id) + if req_index is None: + self.batch_update_builder.finished_paused_append(req_id) + return req_index = super().remove_request(req_id) if req_index is None: From 944ed67e1f0fac1c64747c56952436ca004d5027 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 2 Jun 2026 15:01:39 +0000 Subject: [PATCH 024/106] clean the tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- tests/e2e/test_logits_processors.py | 640 ++++++---------------------- tests/logits_processor_utils.py | 383 +++++++++++++++++ 2 files changed, 516 insertions(+), 507 deletions(-) create mode 100644 tests/logits_processor_utils.py diff --git a/tests/e2e/test_logits_processors.py b/tests/e2e/test_logits_processors.py index 676b2bb7e..a549a7495 100644 --- a/tests/e2e/test_logits_processors.py +++ b/tests/e2e/test_logits_processors.py @@ -1,11 +1,15 @@ import torch from llm_cache import patch_environment from llm_cache_util import force_engine_shutdown +from logits_processor_utils import ( + DummyLogitsProcessor, + NoOpLogitsProcessor, + SpyLogitsProcessor, + StateTrackingLogitsProcessorWrapper, + execute_step, +) from spyre_util import ModelInfo from vllm import LLM, SamplingParams -from vllm.config import VllmConfig -from vllm.v1.sample.logits_processor import BatchUpdate, LogitsProcessor, MoveDirectionality -from sendnn_inference.v1.sample.spyre_logits_processor import SpyreBatchUpdate def test_custom_logits_processor( @@ -18,18 +22,7 @@ def test_custom_logits_processor( monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") has_invoked_logits_processor = False - class DummyLogitsProcessor(LogitsProcessor): - def __init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool): - # Required to register LogitsProcessor - pass - - def is_argmax_invariant(self) -> bool: - return False - - def update_state(self, batch_update: BatchUpdate | None): - # Required to register LogitsProcessor - pass - + class TestDummyLogitsProcessor(DummyLogitsProcessor): def apply(self, logits: torch.Tensor) -> torch.Tensor: nonlocal has_invoked_logits_processor has_invoked_logits_processor = True @@ -47,7 +40,7 @@ def apply(self, logits: torch.Tensor) -> torch.Tensor: max_num_seqs=max_num_seqs, max_num_batched_tokens=128, enable_prefix_caching=mode == "pc", - logits_processors=[DummyLogitsProcessor], + logits_processors=[TestDummyLogitsProcessor], ) prompt = "Hello Logits Processors" params = SamplingParams(max_tokens=5, temperature=0, logprobs=0) @@ -78,65 +71,11 @@ def test_logits_processor_cp(model: ModelInfo, backend, monkeypatch, max_model_l # is the max_tokens to ease identify the requests spy_outputs: dict[int, list[int]] = {} - class SpyLogitsProcessor(LogitsProcessor): - """ - This logits processor collect the tokens - """ - - def __init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool): - self.req_info: dict[int, SamplingParams] = {} - # Saved state for paused requests, keyed by req_id. - self._paused_info: dict[str, SamplingParams] = {} - - def is_argmax_invariant(self) -> bool: - return False - - def update_state(self, batch_update: BatchUpdate | None): - if not batch_update: - return - - for index, params, _, _ in batch_update.added: - self.req_info[index] = params - nonlocal spy_outputs - # Use setdefault so a resume-triggered added event does not - # reset tokens already collected before the pause. - spy_outputs.setdefault(params.max_tokens, []) - - if SpyreBatchUpdate is not None: - for dense_index, req_id in getattr(batch_update, "resumed", []): - if req_id in self._paused_info: - self.req_info[dense_index] = self._paused_info.pop(req_id) - - if SpyreBatchUpdate is not None: - for dense_index, req_id in getattr(batch_update, "paused", []): - if dense_index in self.req_info: - self._paused_info[req_id] = self.req_info.pop(dense_index) - - if self.req_info: - # Process removed requests. - for index in batch_update.removed: - self.req_info.pop(index, None) - - # Process moved requests, unidirectional move (a->b) and swap - # (a<->b) - for adx, bdx, direct in batch_update.moved: - a_val = self.req_info.pop(adx, None) - b_val = self.req_info.pop(bdx, None) - if a_val is not None: - self.req_info[bdx] = a_val - if direct == MoveDirectionality.SWAP and b_val is not None: - self.req_info[adx] = b_val + class TestSpyLogitsProcessor(SpyLogitsProcessor): + """Test-specific spy processor that uses the shared spy_outputs dict.""" - def apply(self, logits: torch.Tensor) -> torch.Tensor: - if not self.req_info: - return - batch_size = logits.shape[0] - nonlocal spy_outputs - for i in range(batch_size): - params = self.req_info[i] - token_id = logits[i].argmax(-1).reshape(-1).item() - spy_outputs[params.max_tokens].append(token_id) - return logits + def __init__(self, vllm_config, device, is_pin_memory): + super().__init__(vllm_config, device, is_pin_memory, spy_outputs) patch_environment( backend=backend, @@ -148,17 +87,17 @@ def apply(self, logits: torch.Tensor) -> torch.Tensor: revision=model.revision, max_model_len=max_model_len, max_num_seqs=2, - logits_processors=[SpyLogitsProcessor], + logits_processors=[TestSpyLogitsProcessor], max_num_batched_tokens=128, enable_prefix_caching=mode == "pc", ) - prompt = ["1 2 3 4 5 6 7 8 9 " * 10] * 3 + prompt = ["Hello Logits Processors"] * 3 params0 = SamplingParams(max_tokens=5, temperature=0, logprobs=0, ignore_eos=True) params1 = SamplingParams(max_tokens=10, temperature=0, logprobs=0, ignore_eos=True) params2 = SamplingParams(max_tokens=7, temperature=0, logprobs=0, ignore_eos=True) # clear from the warmup - spy_outputs = {} + spy_outputs.clear() params = [params0, params1, params2] outputs = spyre_model.generate(prompt, params) force_engine_shutdown(spyre_model) @@ -186,170 +125,20 @@ def test_logits_processor_advanced( 3. Resumed while others finish 4. Multiple operations happening simultaneously """ - from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput - from vllm.v1.request import Request from tests.v1.worker.mock_model import InstrumentedModelRunner + from vllm.v1.sample.logits_processor.state import LogitsProcessors monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") # Track all state transitions and verify correctness state_log: list[dict] = [] - class StateTrackingLogitsProcessor(LogitsProcessor): - """ - Tracks the complete state of the batch at each update. - Verifies that indices are correct and no overwrites occur. - """ - - def __init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool): - # Map dense_index -> (req_id, token_count) - self.active_requests: dict[int, tuple[str, int]] = {} - # Map req_id -> (last_dense_index, token_count) for paused requests - self.paused_requests: dict[str, tuple[int, int]] = {} - self.step_count = 0 - - def is_argmax_invariant(self) -> bool: - return False - - def update_state(self, batch_update: BatchUpdate | None): - if not batch_update: - return - - self.step_count += 1 - step_info = { - "step": self.step_count, - "batch_size": batch_update.batch_size, - "operations": [], - "active_before": dict(self.active_requests), - "paused_before": dict(self.paused_requests), - } - - # Process added requests - for dense_index, params, prompt_toks, output_toks in batch_update.added: - req_id = f"req_{params.max_tokens}" # Use max_tokens as identifier - token_count = len(output_toks) - - # Verify no overwrite - if dense_index in self.active_requests: - old_req = self.active_requests[dense_index] - step_info["operations"].append( - { - "type": "ERROR_OVERWRITE", - "index": dense_index, - "old_req": old_req, - "new_req": (req_id, token_count), - } - ) - - self.active_requests[dense_index] = (req_id, token_count) - step_info["operations"].append( - {"type": "added", "index": dense_index, "req_id": req_id, "tokens": token_count} - ) - - # Process resumed requests (restore from pause) - if hasattr(batch_update, "resumed"): - for dense_index, req_id in batch_update.resumed: - if req_id in self.paused_requests: - old_index, token_count = self.paused_requests.pop(req_id) - self.active_requests[dense_index] = (req_id, token_count) - step_info["operations"].append( - { - "type": "resumed", - "index": dense_index, - "req_id": req_id, - "old_index": old_index, - "tokens": token_count, - } - ) - - # Process paused requests (save state) - if hasattr(batch_update, "paused"): - for dense_index, req_id in batch_update.paused: - if dense_index in self.active_requests: - req_info = self.active_requests.pop(dense_index) - self.paused_requests[req_id] = (dense_index, req_info[1]) - step_info["operations"].append( - { - "type": "paused", - "index": dense_index, - "req_id": req_id, - "tokens": req_info[1], - } - ) - - # Process removed requests - for dense_index in batch_update.removed: - if dense_index in self.active_requests: - req_info = self.active_requests.pop(dense_index) - step_info["operations"].append( - { - "type": "removed", - "index": dense_index, - "req_id": req_info[0], - "tokens": req_info[1], - } - ) - - # Process moved requests (always swaps like LogitProcessorWrapper) - for src_idx, dst_idx, _ in batch_update.moved: - src_req = self.active_requests.get(src_idx) - dst_req = self.active_requests.get(dst_idx) - - # Always swap both positions (matching LogitProcessorWrapper behavior) - if src_req: - self.active_requests[dst_idx] = src_req - else: - self.active_requests.pop(dst_idx, None) - if dst_req: - self.active_requests[src_idx] = dst_req - else: - self.active_requests.pop(src_idx, None) - - step_info["operations"].append( - { - "type": "moved", - "src": src_idx, - "dst": dst_idx, - "src_req": src_req, - "dst_req": dst_req, - } - ) - - step_info["active_after"] = dict(self.active_requests) - step_info["paused_after"] = dict(self.paused_requests) - - # Verify indices are consecutive immediately after each update - if self.active_requests: - indices = sorted(self.active_requests.keys()) - expected = list(range(len(indices))) - if indices != expected: - step_info["operations"].append( - { - "type": "ERROR_NON_CONSECUTIVE", - "actual_indices": indices, - "expected_indices": expected, - "message": ( - f"Indices are not consecutive: {indices}, expected {expected}" - ), - } - ) - - nonlocal state_log - state_log.append(step_info) - - def apply(self, logits: torch.Tensor) -> torch.Tensor: - # Increment token count for all active requests - for idx in self.active_requests: - req_id, token_count = self.active_requests[idx] - self.active_requests[idx] = (req_id, token_count + 1) - return logits - patch_environment( backend=backend, monkeypatch=monkeypatch, ) - # Build the model runner with our tracking processor + # Build the model runner runner = InstrumentedModelRunner.build( monkeypatch=monkeypatch, enable_prefix_caching=mode == "pc", @@ -359,341 +148,178 @@ def apply(self, logits: torch.Tensor) -> torch.Tensor: max_num_batched_tokens=128, ) - # Replace logits processors with our tracking one - from sendnn_inference.v1.sample.spyre_logits_processor import build_logitsprocs_for_cb - - runner.input_batch.logitsprocs = build_logitsprocs_for_cb( - vllm_config=runner.vllm_config, - device=runner.device, - is_pin_memory=runner.pin_memory, - is_pooling_model=False, - batch_size=4, - custom_logitsprocs=[StateTrackingLogitsProcessor], + # Replace logits processors with our tracking wrapper + tracking_wrapper = StateTrackingLogitsProcessorWrapper( + NoOpLogitsProcessor, + runner.vllm_config, + runner.device, + runner.pin_memory, + 4, + state_log, ) + runner.input_batch.logitsprocs = LogitsProcessors([tracking_wrapper]) - # Helper to create requests - def make_request(req_id: str, prompt_len: int, max_tokens: int) -> Request: - return Request( - request_id=req_id, - prompt_token_ids=[42] * prompt_len, - sampling_params=SamplingParams(max_tokens=max_tokens, temperature=0), - pooling_params=None, - ) - - # Helper to create new request data - def make_new_req_data(req_id: str, prompt_len: int, max_tokens: int) -> NewRequestData: - req = make_request(req_id, prompt_len, max_tokens) - block_ids = list(range(1, (prompt_len + 63) // 64 + 1)) - return NewRequestData.from_request(req, block_ids=(block_ids,)) - - # Helper to create cached request data - def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> CachedRequestData: - cached = CachedRequestData.make_empty() - cached.req_ids = list(req_states.keys()) - cached.num_computed_tokens = [state[0] for state in req_states.values()] - cached.new_block_ids = [None for _ in req_states] # Simplified for test - return cached + # Get reference to our tracking wrapper + processor = tracking_wrapper # Add first request: req 0 - req0 = make_new_req_data("req0", 50, 5) - sched_out = SchedulerOutput( - scheduled_new_reqs=[req0], - scheduled_cached_reqs=CachedRequestData.make_empty(), + execute_step( + runner, + processor, + new_reqs=[("req0", 50, 5)], num_scheduled_tokens={"req0": 50}, - total_num_scheduled_tokens=50, - finished_req_ids=set(), - kv_connector_metadata=None, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], - free_encoder_mm_hashes=[], + expected_active={0: ("req0", 1)}, + expected_paused=set(), ) - runner.execute_model(sched_out) # Decode request 0 - cached = make_cached_req_data({"req0": (50, [])}) - sched_out = SchedulerOutput( - scheduled_new_reqs=[], - scheduled_cached_reqs=cached, + execute_step( + runner, + processor, + cached_reqs=[("req0", 50)], num_scheduled_tokens={"req0": 1}, - total_num_scheduled_tokens=1, - finished_req_ids=set(), - kv_connector_metadata=None, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], - free_encoder_mm_hashes=[], + expected_active={0: ("req0", 2)}, + expected_paused=set(), ) - runner.execute_model(sched_out) - # Add second request (long one, needs three chunks) + # Add request 1 (long one, needs three chunks) # Chunked-prefill 1/3 of request 1 - req1 = make_new_req_data("req1", 266, 10) - sched_out = SchedulerOutput( - scheduled_new_reqs=[req1], - scheduled_cached_reqs=CachedRequestData.make_empty(), + # Chunked-prefill is not added to input_batch unless it is the last chunk + execute_step( + runner, + processor, + new_reqs=[("req1", 266, 10)], num_scheduled_tokens={"req1": 128}, - total_num_scheduled_tokens=128, - finished_req_ids=set(), - kv_connector_metadata=None, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], - free_encoder_mm_hashes=[], + expected_active={}, + expected_paused={"req0"}, ) - runner.execute_model(sched_out) # Decode request 0 - cached = make_cached_req_data({"req0": (51, [])}) - sched_out = SchedulerOutput( - scheduled_new_reqs=[], - scheduled_cached_reqs=cached, + execute_step( + runner, + processor, + cached_reqs=[("req0", 51)], num_scheduled_tokens={"req0": 1}, - total_num_scheduled_tokens=1, - finished_req_ids=set(), - kv_connector_metadata=None, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], - free_encoder_mm_hashes=[], + expected_active={0: ("req0", 3)}, + expected_paused=set(), ) - runner.execute_model(sched_out) # Chunked-prefill 2/3 of request 1 - cached = make_cached_req_data({"req1": (128, [])}) - sched_out = SchedulerOutput( - scheduled_new_reqs=[], - scheduled_cached_reqs=cached, + execute_step( + runner, + processor, + cached_reqs=[("req1", 128)], num_scheduled_tokens={"req1": 128}, - total_num_scheduled_tokens=128, - finished_req_ids=set(), - kv_connector_metadata=None, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], - free_encoder_mm_hashes=[], + expected_active={}, + expected_paused={"req0"}, ) - runner.execute_model(sched_out) # Decode request 0 - cached = make_cached_req_data({"req0": (52, [])}) - sched_out = SchedulerOutput( - scheduled_new_reqs=[], - scheduled_cached_reqs=cached, + execute_step( + runner, + processor, + cached_reqs=[("req0", 52)], num_scheduled_tokens={"req0": 1}, - total_num_scheduled_tokens=1, - finished_req_ids=set(), - kv_connector_metadata=None, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], - free_encoder_mm_hashes=[], + expected_active={0: ("req0", 4)}, + expected_paused=set(), ) - runner.execute_model(sched_out) # Chunked-prefill 3/3 of request 1 - cached = make_cached_req_data({"req1": (256, [])}) - sched_out = SchedulerOutput( - scheduled_new_reqs=[], - scheduled_cached_reqs=cached, + execute_step( + runner, + processor, + cached_reqs=[("req1", 256)], num_scheduled_tokens={"req1": 138}, - total_num_scheduled_tokens=138, - finished_req_ids=set(), - kv_connector_metadata=None, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], - free_encoder_mm_hashes=[], + expected_active={0: ("req1", 1)}, + expected_paused={"req0"}, ) - runner.execute_model(sched_out) # Decode requests 0 and 1 - cached = make_cached_req_data({"req0": (53, []), "req1": (266, [])}) - sched_out = SchedulerOutput( - scheduled_new_reqs=[], - scheduled_cached_reqs=cached, + execute_step( + runner, + processor, + cached_reqs=[("req0", 53), ("req1", 266)], num_scheduled_tokens={"req0": 1, "req1": 1}, - total_num_scheduled_tokens=2, - finished_req_ids=set(), - kv_connector_metadata=None, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], - free_encoder_mm_hashes=[], + expected_active={0: ("req1", 2), 1: ("req0", 5)}, + expected_paused=set(), ) - runner.execute_model(sched_out) # Decode request 0, pause request 1 - cached = make_cached_req_data({"req0": (54, [])}) - sched_out = SchedulerOutput( - scheduled_new_reqs=[], - scheduled_cached_reqs=cached, + execute_step( + runner, + processor, + cached_reqs=[("req0", 54)], num_scheduled_tokens={"req0": 1}, - total_num_scheduled_tokens=1, - finished_req_ids=set(), - kv_connector_metadata=None, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], - free_encoder_mm_hashes=[], + expected_active={0: ("req0", 6)}, + expected_paused={"req1"}, ) - runner.execute_model(sched_out) - # Now finish req1, pause req0, and add req2 - req2 = make_new_req_data("req2", 50, 7) - sched_out = SchedulerOutput( - scheduled_new_reqs=[req2], - scheduled_cached_reqs=CachedRequestData.make_empty(), + # Finish req1, pause req0, and add req2 + execute_step( + runner, + processor, + new_reqs=[("req2", 50, 7)], num_scheduled_tokens={"req2": 50}, - total_num_scheduled_tokens=50, finished_req_ids={"req1"}, - kv_connector_metadata=None, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], - free_encoder_mm_hashes=[], + expected_active={0: ("req2", 1)}, + expected_paused={"req0"}, ) - runner.execute_model(sched_out) - - # Verify no overwrites occurred - for step in state_log: - for op in step["operations"]: - assert op["type"] != "ERROR_OVERWRITE", ( - f"Overwrite detected at step {step['step']}: {op}" - ) # Clean up: finish remaining requests - sched_out = SchedulerOutput( - scheduled_new_reqs=[], - scheduled_cached_reqs=CachedRequestData.make_empty(), - num_scheduled_tokens={}, - total_num_scheduled_tokens=0, - finished_req_ids={"req1", "req2"}, - kv_connector_metadata=None, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], - free_encoder_mm_hashes=[], + execute_step( + runner, + processor, + finished_req_ids={"req0", "req2"}, + expected_active={}, + expected_paused=set(), ) - runner.execute_model(sched_out) # Start fresh with req3 - req3 = make_new_req_data("req3", 50, 8) - sched_out = SchedulerOutput( - scheduled_new_reqs=[req3], - scheduled_cached_reqs=CachedRequestData.make_empty(), + execute_step( + runner, + processor, + new_reqs=[("req3", 50, 8)], num_scheduled_tokens={"req3": 50}, - total_num_scheduled_tokens=50, - finished_req_ids=set(), - kv_connector_metadata=None, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], - free_encoder_mm_hashes=[], - ) - runner.execute_model(sched_out) - - # Decode request 3 - cached = make_cached_req_data({"req3": (50, [])}) - sched_out = SchedulerOutput( - scheduled_new_reqs=[], - scheduled_cached_reqs=cached, - num_scheduled_tokens={"req3": 1}, - total_num_scheduled_tokens=1, - finished_req_ids=set(), - kv_connector_metadata=None, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], - free_encoder_mm_hashes=[], + expected_active={0: ("req3", 1)}, + expected_paused=set(), ) - runner.execute_model(sched_out) # Prefill request 4 - req4 = make_new_req_data("req4", 50, 12) - sched_out = SchedulerOutput( - scheduled_new_reqs=[req4], - scheduled_cached_reqs=CachedRequestData.make_empty(), + execute_step( + runner, + processor, + new_reqs=[("req4", 50, 12)], num_scheduled_tokens={"req4": 50}, - total_num_scheduled_tokens=50, - finished_req_ids=set(), - kv_connector_metadata=None, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], - free_encoder_mm_hashes=[], + expected_active={0: ("req4", 1)}, + expected_paused={"req3"}, ) - runner.execute_model(sched_out) - - # Simulate pause of req3 (scheduler removes it temporarily) - # In real scenario, scheduler would not include req3 in cached_reqs - cached = make_cached_req_data({"req4": (50, [])}) - sched_out = SchedulerOutput( - scheduled_new_reqs=[], - scheduled_cached_reqs=cached, + + # Decode request 4, keep request 3 paused + execute_step( + runner, + processor, + cached_reqs=[("req4", 50)], num_scheduled_tokens={"req4": 1}, - total_num_scheduled_tokens=1, - finished_req_ids=set(), - kv_connector_metadata=None, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], - free_encoder_mm_hashes=[], + expected_active={0: ("req4", 2)}, + expected_paused={"req3"}, ) - runner.execute_model(sched_out) # Resume req3 and finish req4 simultaneously - cached = make_cached_req_data({"req3": (51, [])}) - sched_out = SchedulerOutput( - scheduled_new_reqs=[], - scheduled_cached_reqs=cached, + execute_step( + runner, + processor, + cached_reqs=[("req3", 51)], num_scheduled_tokens={"req3": 1}, - total_num_scheduled_tokens=1, finished_req_ids={"req4"}, - kv_connector_metadata=None, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], - free_encoder_mm_hashes=[], + expected_active={0: ("req3", 2)}, + expected_paused=set(), + ) + + # Finish request 3 + execute_step( + runner, + processor, + finished_req_ids={"req3"}, + expected_active={}, + expected_paused=set(), ) - runner.execute_model(sched_out) - - # Final verification - print("\n=== State Log Summary ===") - for step in state_log: - print(f"\nStep {step['step']}:") - print(f" Batch size: {step['batch_size']}") - print(f" Operations: {len(step['operations'])}") - for op in step["operations"]: - print(f" - {op['type']}: {op}") - print(f" Active after: {step['active_after']}") - print(f" Paused after: {step['paused_after']}") - - # Verify no errors occurred during state transitions - errors = [ - (step["step"], op) - for step in state_log - for op in step["operations"] - if op["type"].startswith("ERROR") - ] - - if errors: - print("\n=== ERRORS DETECTED ===") - for step_num, error_op in errors: - print(f"Step {step_num}: {error_op['type']}") - print(f" Details: {error_op}") - - assert len(errors) == 0, f"Found {len(errors)} errors in state transitions: {errors}" - - # Additional verification: all steps should have consecutive indices - for step in state_log: - if step["active_after"]: - indices = sorted(step["active_after"].keys()) - expected = list(range(len(indices))) - assert indices == expected, ( - f"Step {step['step']}: Non-contiguous indices {indices}, expected {expected}" - ) - - print("\n=== Test Passed ===") - print(f"Total steps: {len(state_log)}") - print(f"Total operations: {sum(len(s['operations']) for s in state_log)}") diff --git a/tests/logits_processor_utils.py b/tests/logits_processor_utils.py new file mode 100644 index 000000000..8c96509e1 --- /dev/null +++ b/tests/logits_processor_utils.py @@ -0,0 +1,383 @@ +"""Utility functions and classes for logits processor tests.""" + +import torch +from vllm import SamplingParams +from vllm.config import VllmConfig +from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput +from vllm.v1.request import Request +from vllm.v1.sample.logits_processor import BatchUpdate, LogitsProcessor, MoveDirectionality + +from sendnn_inference.v1.sample.spyre_logits_processor import ( + LogitProcessorWrapper, +) + + +class DummyLogitsProcessor(LogitsProcessor): + """A simple dummy logits processor for testing registration.""" + + def __init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool): + # Required to register LogitsProcessor + pass + + def is_argmax_invariant(self) -> bool: + return False + + def update_state(self, batch_update: BatchUpdate | None): + # Required to register LogitsProcessor + pass + + def apply(self, logits: torch.Tensor) -> torch.Tensor: + return logits + + +class SpyLogitsProcessor(LogitsProcessor): + """ + A logits processor that collects generated tokens for verification. + + This processor tracks the state of requests and collects the tokens + generated by each request, storing them in the provided spy_outputs dict. + """ + + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + is_pin_memory: bool, + spy_outputs: dict[int, list[int]], + ): + self.req_info: dict[int, SamplingParams] = {} + # Saved state for paused requests, keyed by req_id. + self._paused_info: dict[str, SamplingParams] = {} + self.spy_outputs = spy_outputs + + def is_argmax_invariant(self) -> bool: + return False + + def update_state(self, batch_update: BatchUpdate | None): + if not batch_update: + return + + for index, params, _, _ in batch_update.added: + self.req_info[index] = params + # Use setdefault so a resume-triggered added event does not + # reset tokens already collected before the pause. + self.spy_outputs.setdefault(params.max_tokens, []) + + if self.req_info: + # Process removed requests. + for index in batch_update.removed: + self.req_info.pop(index, None) + + # Process moved requests, unidirectional move (a->b) and swap + # (a<->b) + for adx, bdx, direct in batch_update.moved: + a_val = self.req_info.pop(adx, None) + b_val = self.req_info.pop(bdx, None) + if a_val is not None: + self.req_info[bdx] = a_val + if direct == MoveDirectionality.SWAP and b_val is not None: + self.req_info[adx] = b_val + + def apply(self, logits: torch.Tensor) -> torch.Tensor: + if not self.req_info: + return logits + batch_size = logits.shape[0] + for i in range(batch_size): + params = self.req_info[i] + token_id = logits[i].argmax(-1).reshape(-1).item() + self.spy_outputs[params.max_tokens].append(token_id) + return logits + + +class NoOpLogitsProcessor(LogitsProcessor): + """A no-op logits processor that does nothing.""" + + def __init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool): + pass + + def is_argmax_invariant(self) -> bool: + return False + + def update_state(self, batch_update: BatchUpdate | None): + pass + + def apply(self, logits: torch.Tensor) -> torch.Tensor: + return logits + + +class StateTrackingLogitsProcessorWrapper(LogitProcessorWrapper): + """ + Extends LogitProcessorWrapper to track batch-level state for testing. + + This wrapper tracks all state transitions including adds, removes, pauses, + resumes, and moves. It maintains a log of all operations and verifies that + indices remain consecutive after each update. + """ + + def __init__( + self, + logit_processor: type[LogitsProcessor], + vllm_config: VllmConfig, + device: torch.device, + is_pin_memory: bool, + batch_size: int, + state_log: list[dict], + ): + super().__init__(logit_processor, vllm_config, device, is_pin_memory, batch_size) + # Map dense_index -> (req_id, token_count) + self.active_requests: dict[int, tuple[str, int]] = {} + # Map req_id -> (last_dense_index, token_count) for paused requests + self.paused_requests: dict[str, tuple[int, int]] = {} + self.step_count = 0 + self.request_counter = 0 # Counter for number of requests seen + self.state_log = state_log + + def update_state(self, batch_update: BatchUpdate | None) -> None: + # Call parent to handle the actual logits processor state + super().update_state(batch_update) + + # Track our own state for testing + if not batch_update: + return + + self.step_count += 1 + step_info = { + "step": self.step_count, + "batch_size": batch_update.batch_size, + "operations": [], + "active_before": dict(self.active_requests), + "paused_before": dict(self.paused_requests), + } + + # Process added requests + for dense_index, params, prompt_toks, output_toks in batch_update.added: + req_id = f"req{self.request_counter}" # Use request counter as identifier + self.request_counter += 1 + token_count = len(output_toks) + + # Verify no overwrite - raise error immediately + if dense_index in self.active_requests: + old_req = self.active_requests[dense_index] + raise AssertionError( + f"Overwrite detected at step {self.step_count}: " + f"index {dense_index} already has {old_req}, " + f"attempting to add {(req_id, token_count)}" + ) + + self.active_requests[dense_index] = (req_id, token_count) + step_info["operations"].append( + {"type": "added", "index": dense_index, "req_id": req_id, "tokens": token_count} + ) + + # Process resumed requests (restore from pause) + if hasattr(batch_update, "resumed"): + for dense_index, req_id in batch_update.resumed: + if req_id in self.paused_requests: + old_index, token_count = self.paused_requests.pop(req_id) + self.active_requests[dense_index] = (req_id, token_count) + step_info["operations"].append( + { + "type": "resumed", + "index": dense_index, + "req_id": req_id, + "old_index": old_index, + "tokens": token_count, + } + ) + + # Process paused requests (save state) + if hasattr(batch_update, "paused"): + for dense_index, req_id in batch_update.paused: + if dense_index in self.active_requests: + req_info = self.active_requests.pop(dense_index) + self.paused_requests[req_id] = (dense_index, req_info[1]) + step_info["operations"].append( + { + "type": "paused", + "index": dense_index, + "req_id": req_id, + "tokens": req_info[1], + } + ) + + # Process finished paused requests (remove from paused state) + if hasattr(batch_update, "finished_paused"): + for req_id in batch_update.finished_paused: + if req_id in self.paused_requests: + paused_info = self.paused_requests.pop(req_id) + step_info["operations"].append( + { + "type": "finished_paused", + "req_id": req_id, + "last_index": paused_info[0], + "tokens": paused_info[1], + } + ) + + # Process removed requests + for dense_index in batch_update.removed: + if dense_index in self.active_requests: + req_info = self.active_requests.pop(dense_index) + step_info["operations"].append( + { + "type": "removed", + "index": dense_index, + "req_id": req_info[0], + "tokens": req_info[1], + } + ) + + # Process moved requests (always swaps like LogitProcessorWrapper) + for src_idx, dst_idx, _ in batch_update.moved: + src_req = self.active_requests.get(src_idx) + dst_req = self.active_requests.get(dst_idx) + + # Always swap both positions (matching LogitProcessorWrapper behavior) + if src_req: + self.active_requests[dst_idx] = src_req + else: + self.active_requests.pop(dst_idx, None) + if dst_req: + self.active_requests[src_idx] = dst_req + else: + self.active_requests.pop(src_idx, None) + + step_info["operations"].append( + { + "type": "moved", + "src": src_idx, + "dst": dst_idx, + "src_req": src_req, + "dst_req": dst_req, + } + ) + + step_info["active_after"] = dict(self.active_requests) + step_info["paused_after"] = dict(self.paused_requests) + + # Verify indices are consecutive immediately after each update + if self.active_requests: + indices = sorted(self.active_requests.keys()) + expected = list(range(len(indices))) + if indices != expected: + step_info["operations"].append( + { + "type": "ERROR_NON_CONSECUTIVE", + "actual_indices": indices, + "expected_indices": expected, + "message": (f"Indices are not consecutive: {indices}, expected {expected}"), + } + ) + + self.state_log.append(step_info) + + def apply(self, logits: torch.Tensor) -> torch.Tensor: + # Increment token count for all active requests + for idx in self.active_requests: + req_id, token_count = self.active_requests[idx] + self.active_requests[idx] = (req_id, token_count + 1) + return super().apply(logits) + + +def make_request(req_id: str, prompt_len: int, max_tokens: int) -> Request: + """Create a Request object for testing.""" + return Request( + request_id=req_id, + prompt_token_ids=[42] * prompt_len, + sampling_params=SamplingParams(max_tokens=max_tokens, temperature=0), + pooling_params=None, + ) + + +def make_new_req_data(req_id: str, prompt_len: int, max_tokens: int) -> NewRequestData: + """Create NewRequestData for testing.""" + req = make_request(req_id, prompt_len, max_tokens) + block_ids = list(range(1, (prompt_len + 63) // 64 + 1)) + return NewRequestData.from_request(req, block_ids=(block_ids,)) + + +def make_cached_req_data(req_states: dict[str, tuple[int, list[int]]]) -> CachedRequestData: + """Create CachedRequestData for testing.""" + cached = CachedRequestData.make_empty() + cached.req_ids = list(req_states.keys()) + cached.num_computed_tokens = [state[0] for state in req_states.values()] + cached.new_block_ids = [None for _ in req_states] # Simplified for test + return cached + + +def execute_step( + runner, + processor: StateTrackingLogitsProcessorWrapper, + new_reqs: list[tuple[str, int, int]] | None = None, + cached_reqs: list[tuple[str, int]] | None = None, + num_scheduled_tokens: dict[str, int] | None = None, + finished_req_ids: set[str] | None = None, + expected_active: dict[int, tuple[str, int]] | None = None, + expected_paused: set[str] | None = None, +): + """ + Execute a scheduler step and verify the resulting state. + + Args: + runner: The model runner to execute the step on + processor: The state tracking processor to verify + new_reqs: List of new requests as (req_id, prompt_len, max_tokens) + cached_reqs: List of cached requests as (req_id, num_computed_tokens) + num_scheduled_tokens: Dict mapping req_id to number of tokens scheduled + finished_req_ids: Set of request IDs that finished in this step + expected_active: Expected active requests mapping dense_index to (req_id, token_count) + expected_paused: Expected set of paused request IDs + """ + # Build new request data + scheduled_new_reqs = [] + if new_reqs: + for req_id, prompt_len, max_tokens in new_reqs: + scheduled_new_reqs.append(make_new_req_data(req_id, prompt_len, max_tokens)) + + # Build cached request data + if cached_reqs: + req_states = {req_id: (num_computed, []) for req_id, num_computed in cached_reqs} + scheduled_cached_reqs = make_cached_req_data(req_states) + else: + scheduled_cached_reqs = CachedRequestData.make_empty() + + # Build scheduler output + sched_out = SchedulerOutput( + scheduled_new_reqs=scheduled_new_reqs, + scheduled_cached_reqs=scheduled_cached_reqs, + num_scheduled_tokens=num_scheduled_tokens or {}, + total_num_scheduled_tokens=sum((num_scheduled_tokens or {}).values()), + finished_req_ids=finished_req_ids or set(), + kv_connector_metadata=None, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + free_encoder_mm_hashes=[], + ) + + # Execute the step + runner.execute_model(sched_out) + + # Verify expected state if provided + if expected_active is not None: + for dense_idx, (expected_req_id, expected_token_count) in expected_active.items(): + assert dense_idx in processor.active_requests, ( + f"Expected index {dense_idx} to be active" + ) + actual_req_id, actual_token_count = processor.active_requests[dense_idx] + assert actual_req_id == expected_req_id, ( + f"Expected {expected_req_id} at index {dense_idx}, got {actual_req_id}" + ) + assert actual_token_count == expected_token_count, ( + f"Expected {expected_token_count} tokens for {expected_req_id} " + f"at index {dense_idx}, got {actual_token_count}" + ) + assert len(processor.active_requests) == len(expected_active), ( + f"Expected {len(expected_active)} active requests, got {len(processor.active_requests)}" + ) + + if expected_paused is not None: + actual_paused = set(processor.paused_requests.keys()) + assert actual_paused == expected_paused, ( + f"Expected paused requests {expected_paused}, got {actual_paused}" + ) From 75eb053509b2099aadb0a7df1c5cffd878b915e8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 2 Jun 2026 16:06:14 +0000 Subject: [PATCH 025/106] logits_processor advanced test: also check the output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- tests/e2e/test_logits_processors.py | 87 ++++++++++++++++++++++++++--- tests/logits_processor_utils.py | 22 +++++++- 2 files changed, 100 insertions(+), 9 deletions(-) diff --git a/tests/e2e/test_logits_processors.py b/tests/e2e/test_logits_processors.py index a549a7495..c5bfc8dfd 100644 --- a/tests/e2e/test_logits_processors.py +++ b/tests/e2e/test_logits_processors.py @@ -3,7 +3,6 @@ from llm_cache_util import force_engine_shutdown from logits_processor_utils import ( DummyLogitsProcessor, - NoOpLogitsProcessor, SpyLogitsProcessor, StateTrackingLogitsProcessorWrapper, execute_step, @@ -118,12 +117,15 @@ def test_logits_processor_advanced( - Pausing and resuming requests - Verifying correct index management - Ensuring no state overwrites occur + - Verifying that the spy logits processor produces correct tokens This test simulates various scheduler scenarios where requests can be: 1. Added and finished in the same step 2. Paused and resumed 3. Resumed while others finish 4. Multiple operations happening simultaneously + + Uses SpyLogitsProcessor as the inner processor to verify token generation. """ from tests.v1.worker.mock_model import InstrumentedModelRunner from vllm.v1.sample.logits_processor.state import LogitsProcessors @@ -133,6 +135,12 @@ def test_logits_processor_advanced( # Track all state transitions and verify correctness state_log: list[dict] = [] + # Track spy outputs to verify token generation + spy_outputs: dict[int, list[int]] = {} + + # Track actual generated outputs from model runner + actual_outputs: dict[str, list[int]] = {} + patch_environment( backend=backend, monkeypatch=monkeypatch, @@ -148,9 +156,13 @@ def test_logits_processor_advanced( max_num_batched_tokens=128, ) - # Replace logits processors with our tracking wrapper + # Create a SpyLogitsProcessor factory that uses our spy_outputs dict + def create_spy_processor(vllm_config, device, is_pin_memory): + return SpyLogitsProcessor(vllm_config, device, is_pin_memory, spy_outputs) + + # Replace logits processors with our tracking wrapper that wraps SpyLogitsProcessor tracking_wrapper = StateTrackingLogitsProcessorWrapper( - NoOpLogitsProcessor, + create_spy_processor, runner.vllm_config, runner.device, runner.pin_memory, @@ -166,10 +178,11 @@ def test_logits_processor_advanced( execute_step( runner, processor, - new_reqs=[("req0", 50, 5)], + new_reqs=[("req0", 50, 10)], num_scheduled_tokens={"req0": 50}, expected_active={0: ("req0", 1)}, expected_paused=set(), + actual_outputs=actual_outputs, ) # Decode request 0 @@ -180,6 +193,7 @@ def test_logits_processor_advanced( num_scheduled_tokens={"req0": 1}, expected_active={0: ("req0", 2)}, expected_paused=set(), + actual_outputs=actual_outputs, ) # Add request 1 (long one, needs three chunks) @@ -188,10 +202,11 @@ def test_logits_processor_advanced( execute_step( runner, processor, - new_reqs=[("req1", 266, 10)], + new_reqs=[("req1", 266, 11)], num_scheduled_tokens={"req1": 128}, expected_active={}, expected_paused={"req0"}, + actual_outputs=actual_outputs, ) # Decode request 0 @@ -202,6 +217,7 @@ def test_logits_processor_advanced( num_scheduled_tokens={"req0": 1}, expected_active={0: ("req0", 3)}, expected_paused=set(), + actual_outputs=actual_outputs, ) # Chunked-prefill 2/3 of request 1 @@ -212,6 +228,7 @@ def test_logits_processor_advanced( num_scheduled_tokens={"req1": 128}, expected_active={}, expected_paused={"req0"}, + actual_outputs=actual_outputs, ) # Decode request 0 @@ -222,6 +239,7 @@ def test_logits_processor_advanced( num_scheduled_tokens={"req0": 1}, expected_active={0: ("req0", 4)}, expected_paused=set(), + actual_outputs=actual_outputs, ) # Chunked-prefill 3/3 of request 1 @@ -232,6 +250,7 @@ def test_logits_processor_advanced( num_scheduled_tokens={"req1": 138}, expected_active={0: ("req1", 1)}, expected_paused={"req0"}, + actual_outputs=actual_outputs, ) # Decode requests 0 and 1 @@ -242,6 +261,7 @@ def test_logits_processor_advanced( num_scheduled_tokens={"req0": 1, "req1": 1}, expected_active={0: ("req1", 2), 1: ("req0", 5)}, expected_paused=set(), + actual_outputs=actual_outputs, ) # Decode request 0, pause request 1 @@ -252,17 +272,19 @@ def test_logits_processor_advanced( num_scheduled_tokens={"req0": 1}, expected_active={0: ("req0", 6)}, expected_paused={"req1"}, + actual_outputs=actual_outputs, ) # Finish req1, pause req0, and add req2 execute_step( runner, processor, - new_reqs=[("req2", 50, 7)], + new_reqs=[("req2", 50, 12)], num_scheduled_tokens={"req2": 50}, finished_req_ids={"req1"}, expected_active={0: ("req2", 1)}, expected_paused={"req0"}, + actual_outputs=actual_outputs, ) # Clean up: finish remaining requests @@ -272,26 +294,29 @@ def test_logits_processor_advanced( finished_req_ids={"req0", "req2"}, expected_active={}, expected_paused=set(), + actual_outputs=actual_outputs, ) # Start fresh with req3 execute_step( runner, processor, - new_reqs=[("req3", 50, 8)], + new_reqs=[("req3", 50, 13)], num_scheduled_tokens={"req3": 50}, expected_active={0: ("req3", 1)}, expected_paused=set(), + actual_outputs=actual_outputs, ) # Prefill request 4 execute_step( runner, processor, - new_reqs=[("req4", 50, 12)], + new_reqs=[("req4", 50, 14)], num_scheduled_tokens={"req4": 50}, expected_active={0: ("req4", 1)}, expected_paused={"req3"}, + actual_outputs=actual_outputs, ) # Decode request 4, keep request 3 paused @@ -302,6 +327,7 @@ def test_logits_processor_advanced( num_scheduled_tokens={"req4": 1}, expected_active={0: ("req4", 2)}, expected_paused={"req3"}, + actual_outputs=actual_outputs, ) # Resume req3 and finish req4 simultaneously @@ -313,6 +339,7 @@ def test_logits_processor_advanced( finished_req_ids={"req4"}, expected_active={0: ("req3", 2)}, expected_paused=set(), + actual_outputs=actual_outputs, ) # Finish request 3 @@ -322,4 +349,48 @@ def test_logits_processor_advanced( finished_req_ids={"req3"}, expected_active={}, expected_paused=set(), + actual_outputs=actual_outputs, + ) + + # Verify that actual outputs match spy outputs for each request + # req0: max_tokens=10, req1: max_tokens=11, req2: max_tokens=12, + # req3: max_tokens=13, req4: max_tokens=14 + assert "req0" in actual_outputs, "Expected actual_outputs to contain tokens for req0" + assert len(actual_outputs["req0"]) == 6, ( + f"Expected 6 tokens for req0, got {len(actual_outputs['req0'])}" + ) + assert actual_outputs["req0"] == spy_outputs[10], ( + f"Token mismatch for req0: {actual_outputs['req0']} != {spy_outputs[10]}" + ) + + assert "req1" in actual_outputs, "Expected actual_outputs to contain tokens for req1" + assert len(actual_outputs["req1"]) == 2, ( + f"Expected 2 tokens for req1, got {len(actual_outputs['req1'])}" + ) + assert actual_outputs["req1"] == spy_outputs[11], ( + f"Token mismatch for req1: {actual_outputs['req1']} != {spy_outputs[11]}" + ) + + assert "req2" in actual_outputs, "Expected actual_outputs to contain tokens for req2" + assert len(actual_outputs["req2"]) == 1, ( + f"Expected 1 tokens for req2, got {len(actual_outputs['req2'])}" + ) + assert actual_outputs["req2"] == spy_outputs[12], ( + f"Token mismatch for req2: {actual_outputs['req2']} != {spy_outputs[12]}" + ) + + assert "req3" in actual_outputs, "Expected actual_outputs to contain tokens for req3" + assert len(actual_outputs["req3"]) == 2, ( + f"Expected 2 tokens for req3, got {len(actual_outputs['req3'])}" + ) + assert actual_outputs["req3"] == spy_outputs[13], ( + f"Token mismatch for req3: {actual_outputs['req3']} != {spy_outputs[13]}" + ) + + assert "req4" in actual_outputs, "Expected actual_outputs to contain tokens for req4" + assert len(actual_outputs["req4"]) == 2, ( + f"Expected 2 tokens for req4, got {len(actual_outputs['req4'])}" + ) + assert actual_outputs["req4"] == spy_outputs[14], ( + f"Token mismatch for req4: {actual_outputs['req4']} != {spy_outputs[14]}" ) diff --git a/tests/logits_processor_utils.py b/tests/logits_processor_utils.py index 8c96509e1..f62df8d55 100644 --- a/tests/logits_processor_utils.py +++ b/tests/logits_processor_utils.py @@ -314,6 +314,7 @@ def execute_step( finished_req_ids: set[str] | None = None, expected_active: dict[int, tuple[str, int]] | None = None, expected_paused: set[str] | None = None, + actual_outputs: dict[str, list[int]] | None = None, ): """ Execute a scheduler step and verify the resulting state. @@ -327,6 +328,7 @@ def execute_step( finished_req_ids: Set of request IDs that finished in this step expected_active: Expected active requests mapping dense_index to (req_id, token_count) expected_paused: Expected set of paused request IDs + actual_outputs: Optional dict to collect generated tokens, keyed by request_id """ # Build new request data scheduled_new_reqs = [] @@ -356,7 +358,23 @@ def execute_step( ) # Execute the step - runner.execute_model(sched_out) + output = runner.execute_model(sched_out) + + # Collect generated tokens into actual_outputs if provided + if actual_outputs is not None and output.sampled_token_ids: + # Map dense indices to request IDs and collect tokens + for dense_idx, token_ids in enumerate(output.sampled_token_ids): + if dense_idx in processor.active_requests: + req_id, token_count = processor.active_requests[dense_idx] + # Extract the actual token (first element of token_ids list) + actual_token = token_ids[0] if isinstance(token_ids, list) else token_ids + + # Initialize list for this request if not exists + if req_id not in actual_outputs: + actual_outputs[req_id] = [] + + # Append the generated token + actual_outputs[req_id].append(actual_token) # Verify expected state if provided if expected_active is not None: @@ -381,3 +399,5 @@ def execute_step( assert actual_paused == expected_paused, ( f"Expected paused requests {expected_paused}, got {actual_paused}" ) + + return output From 0530cab18ef6f50af286739f6369a510c4f9aedb Mon Sep 17 00:00:00 2001 From: Max de Bayser Date: Tue, 2 Jun 2026 14:59:42 -0400 Subject: [PATCH 026/106] remove superseded design idea for block reservation Signed-off-by: Max de Bayser --- sendnn_inference/v1/core/scheduler.py | 37 --------------------------- 1 file changed, 37 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 5d460edc8..ceebc8109 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -257,43 +257,6 @@ def adjust_computed_tokens( # Otherwise just account for the left padding return computed_tokens - left_padding - def get_required_blocks(self, request: Request) -> tuple[int, int, bool]: - assert request.prompt_token_ids is not None - assert ( - request.sampling_params is not None and request.sampling_params.max_tokens is not None - ) - max_tokens = len(request.prompt_token_ids) + request.sampling_params.max_tokens - max_tokens = min(self.max_model_len, max_tokens) - - total_blocks = math.ceil(max_tokens / self.block_size) - - block_ids_per_kv_cache_group = self.kv_cache_manager.get_block_ids(request.request_id) - assert len(block_ids_per_kv_cache_group) == 1 - used_blocks = len(block_ids_per_kv_cache_group[0]) - - total_tokens = request.num_tokens - # the request will get a new block in the next iteration - needs_new_block_now = total_tokens < max_tokens and total_tokens % 64 == 0 - - return total_blocks, used_blocks, needs_new_block_now - - def get_blocks_required_for_decode_batch(self, before_allocation: bool = True) -> int: - """ - Returns the number of blocks that the current decode batch needs to - finish all requests. Reducing the number of available blocks below - this number will cause deadlocks. - """ - # Warning, depending on when this function is called, self.running - # might be out of sync with the worker's input_batch - required_blocks = 0 - for request in self.running: - total_blocks, used_blocks, needs_new_block_now = self.get_required_blocks(request) - required_blocks += ( - total_blocks - used_blocks + int(before_allocation and needs_new_block_now) - ) - - return required_blocks - def schedule(self) -> "SchedulerOutput": """ The chunked prefill scheduling policy is enforced in this method, then From 84cf5db67754a724e74789f17b45ba37b4b148bc Mon Sep 17 00:00:00 2001 From: Max de Bayser Date: Tue, 2 Jun 2026 21:23:15 -0400 Subject: [PATCH 027/106] Fix input batch removal of requests It seems that in some cases the request id that arrives at the input_batch for removal has never been seen by the input batch before. This could happen because of the cancellation of a request that was never scheduled. Also, it seems that in the ChunkedPrefillModelRunner we somehow forgot to remove requests from self.requests, so that "warmup-0" and other requests were accumulating forever. Signed-off-by: Max de Bayser --- sendnn_inference/v1/sample/spyre_logits_processor.py | 8 +++++--- sendnn_inference/v1/worker/spyre_model_runner.py | 1 + 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/sendnn_inference/v1/sample/spyre_logits_processor.py b/sendnn_inference/v1/sample/spyre_logits_processor.py index 5417da3b1..3b286232b 100644 --- a/sendnn_inference/v1/sample/spyre_logits_processor.py +++ b/sendnn_inference/v1/sample/spyre_logits_processor.py @@ -161,11 +161,13 @@ def update_state(self, batch_update: BatchUpdate | None) -> None: for index, req_id in getattr(batch_update, "paused", []): self._saved[req_id] = self.logitprocs[index] - self.logitprocs[index] = None + self.logitprocs[index] = self._factory() for req_id in getattr(batch_update, "finished_paused", []): - assert req_id in self._saved - self._saved.pop(req_id) + # Max: I think we can't assume that the request will + # be here because it could be a cancelled request that + # never made it into the batch. + self._saved.pop(req_id, 0) for adx, bdx, _ in batch_update.moved: update_called[adx], update_called[bdx] = update_called[bdx], update_called[adx] diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index 22f131c5b..1954eb3f3 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -1484,6 +1484,7 @@ def _update_batch(self, scheduler_output: SchedulerOutput): if scheduler_output.finished_req_ids: for req_id in scheduler_output.finished_req_ids: self.input_batch.remove_request(req_id) + self.requests.pop(req_id, None) # Clean up paused tracking for finished requests self.paused_req_ids.discard(req_id) # TODO: Processing multiple removals at once can break alignment From 4bd1c833da676af45450089334c3af555902a307 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 3 Jun 2026 12:14:51 +0000 Subject: [PATCH 028/106] cleanup test code MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- tests/logits_processor_utils.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/tests/logits_processor_utils.py b/tests/logits_processor_utils.py index f62df8d55..e572ef300 100644 --- a/tests/logits_processor_utils.py +++ b/tests/logits_processor_utils.py @@ -89,22 +89,6 @@ def apply(self, logits: torch.Tensor) -> torch.Tensor: return logits -class NoOpLogitsProcessor(LogitsProcessor): - """A no-op logits processor that does nothing.""" - - def __init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool): - pass - - def is_argmax_invariant(self) -> bool: - return False - - def update_state(self, batch_update: BatchUpdate | None): - pass - - def apply(self, logits: torch.Tensor) -> torch.Tensor: - return logits - - class StateTrackingLogitsProcessorWrapper(LogitProcessorWrapper): """ Extends LogitProcessorWrapper to track batch-level state for testing. From 9d5f9d00462852405af48587b4aa1045e13759ed Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 3 Jun 2026 13:00:55 +0000 Subject: [PATCH 029/106] factor out common part MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../v1/worker/spyre_input_batch.py | 157 ++++++------------ 1 file changed, 49 insertions(+), 108 deletions(-) diff --git a/sendnn_inference/v1/worker/spyre_input_batch.py b/sendnn_inference/v1/worker/spyre_input_batch.py index 60831cb0d..28c2f0b25 100644 --- a/sendnn_inference/v1/worker/spyre_input_batch.py +++ b/sendnn_inference/v1/worker/spyre_input_batch.py @@ -406,79 +406,9 @@ def add_request( end_idx = start_idx + len(request.output_token_ids) self.token_ids_cpu[req_index, start_idx:end_idx] = request.output_token_ids - sampling_params = request.sampling_params - if sampling_params.sampling_type == SamplingType.GREEDY: - # Avoid later division by zero. - self.temperature_cpu[req_index] = -1.0 - self.greedy_reqs.add(req_id) - else: - self.temperature_cpu[req_index] = sampling_params.temperature - self.random_reqs.add(req_id) - - self.top_p_cpu[req_index] = sampling_params.top_p - if sampling_params.top_p < 1: - self.top_p_reqs.add(req_id) - top_k = sampling_params.top_k - if 0 < top_k < self.vocab_size: - self.top_k_reqs.add(req_id) - else: - top_k = self.vocab_size - self.top_k_cpu[req_index] = top_k - self.frequency_penalties_cpu[req_index] = sampling_params.frequency_penalty - if sampling_params.frequency_penalty != 0.0: - self.frequency_penalties_reqs.add(req_id) - self.presence_penalties_cpu[req_index] = sampling_params.presence_penalty - if sampling_params.presence_penalty != 0.0: - self.presence_penalties_reqs.add(req_id) - self.repetition_penalties_cpu[req_index] = sampling_params.repetition_penalty - if sampling_params.repetition_penalty != 1.0: - self.repetition_penalties_reqs.add(req_id) - - # NOTE(woosuk): self.generators should not include the requests that - # do not have their own generator. - if request.generator is not None: - self.generators[req_index] = request.generator - - if sampling_params.logprobs is not None: - self.num_logprobs[req_id] = sampling_params.logprobs - - if sampling_params.allowed_token_ids: - self.has_allowed_token_ids.add(req_id) - if self.allowed_token_ids_mask is None: - # Lazy allocation for this tensor, which can be large. - self.allowed_token_ids_mask = torch.zeros( - self.max_num_reqs, self.vocab_size, dtype=torch.bool, device=self.device - ) - self.allowed_token_ids_mask[req_index][sampling_params.allowed_token_ids] = True - - if sampling_params.bad_words_token_ids: - self.bad_words_token_ids[req_index] = sampling_params.bad_words_token_ids + self._register_sampling_params(req_id, req_index, request) return req_index - def clear_requests(self): - """ - Clear the batch, mostly used by static batching - """ - super().clear_requests() - self.req_indices_mask.fill_(False) - self.req_output_token_ids = [] - - self.greedy_reqs = set() - self.random_reqs = set() - self.top_p_reqs = set() - self.top_k_reqs = set() - self.frequency_penalties_reqs = set() - self.presence_penalties_reqs = set() - self.repetition_penalties_reqs = set() - self.generators = {} - self.num_logprobs = {} - - self.has_allowed_token_ids = set() - if self.allowed_token_ids_mask is not None: - self.allowed_token_ids_mask.fill_(False) - - self.batch_update_builder.get_and_reset(0) - def remove_request(self, req_id: str): """ Free a slot of a request from the batch @@ -518,26 +448,8 @@ def remove_request(self, req_id: str): (tmp_dense, tmp_dense + 1, MoveDirectionality.UNIDIRECTIONAL) ) - # Remove the references self.req_output_token_ids.pop(dense_index) - - self.greedy_reqs.discard(req_id) - self.random_reqs.discard(req_id) - self.top_p_reqs.discard(req_id) - self.top_k_reqs.discard(req_id) - - self.frequency_penalties_reqs.discard(req_id) - self.presence_penalties_reqs.discard(req_id) - self.repetition_penalties_reqs.discard(req_id) - self.generators.pop(req_index, None) - self.num_logprobs.pop(req_id, None) - - self.has_allowed_token_ids.discard(req_id) - - if self.allowed_token_ids_mask is not None: - self.allowed_token_ids_mask[req_index].fill_(False) - - self.bad_words_token_ids.pop(req_index, None) + self._unregister_sampling_params(req_id, req_index) def pause_request(self, req_id: str) -> None: """Temporarily remove a request from the active batch. @@ -572,24 +484,7 @@ def pause_request(self, req_id: str) -> None: ) self.req_output_token_ids.pop(dense_index) - - self.greedy_reqs.discard(req_id) - self.random_reqs.discard(req_id) - self.top_p_reqs.discard(req_id) - self.top_k_reqs.discard(req_id) - - self.frequency_penalties_reqs.discard(req_id) - self.presence_penalties_reqs.discard(req_id) - self.repetition_penalties_reqs.discard(req_id) - self.generators.pop(req_index, None) - self.num_logprobs.pop(req_id, None) - - self.has_allowed_token_ids.discard(req_id) - - if self.allowed_token_ids_mask is not None: - self.allowed_token_ids_mask[req_index].fill_(False) - - self.bad_words_token_ids.pop(req_index, None) + self._unregister_sampling_params(req_id, req_index) def resume_request(self, req_id: str, request: "SamplingRequestState") -> None: """Restore a previously paused request to the active batch. @@ -635,6 +530,12 @@ def resume_request(self, req_id: str, request: "SamplingRequestState") -> None: ) tmp_dense -= 1 + self._register_sampling_params(req_id, req_index, request) + + def _register_sampling_params( + self, req_id: str, req_index: int, request: "SamplingRequestState" + ) -> None: + """Write all sampling parameter fields for a newly-occupied slot.""" sampling_params = request.sampling_params if sampling_params.sampling_type == SamplingType.GREEDY: self.temperature_cpu[req_index] = -1.0 @@ -679,6 +580,46 @@ def resume_request(self, req_id: str, request: "SamplingRequestState") -> None: if sampling_params.bad_words_token_ids: self.bad_words_token_ids[req_index] = sampling_params.bad_words_token_ids + def _unregister_sampling_params(self, req_id: str, req_index: int) -> None: + """Clear all per-request fields when vacating a slot.""" + self.greedy_reqs.discard(req_id) + self.random_reqs.discard(req_id) + self.top_p_reqs.discard(req_id) + self.top_k_reqs.discard(req_id) + self.frequency_penalties_reqs.discard(req_id) + self.presence_penalties_reqs.discard(req_id) + self.repetition_penalties_reqs.discard(req_id) + self.generators.pop(req_index, None) + self.num_logprobs.pop(req_id, None) + self.has_allowed_token_ids.discard(req_id) + if self.allowed_token_ids_mask is not None: + self.allowed_token_ids_mask[req_index].fill_(False) + self.bad_words_token_ids.pop(req_index, None) + + def clear_requests(self): + """ + Clear the batch, mostly used by static batching + """ + super().clear_requests() + self.req_indices_mask.fill_(False) + self.req_output_token_ids = [] + + self.greedy_reqs = set() + self.random_reqs = set() + self.top_p_reqs = set() + self.top_k_reqs = set() + self.frequency_penalties_reqs = set() + self.presence_penalties_reqs = set() + self.repetition_penalties_reqs = set() + self.generators = {} + self.num_logprobs = {} + + self.has_allowed_token_ids = set() + if self.allowed_token_ids_mask is not None: + self.allowed_token_ids_mask.fill_(False) + + self.batch_update_builder.get_and_reset(0) + def refresh_metadata(self): """Apply batch updates, reset input batch at end of step From fa6de4af0c4e83622a3eb7fde1f0f3e4301641dc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Thu, 4 Jun 2026 08:46:59 +0000 Subject: [PATCH 030/106] predict next tkv using num_computed_tokens instead of allocated blocks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 25 +++---------------- .../v1/worker/spyre_model_runner.py | 2 -- 2 files changed, 4 insertions(+), 23 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index ceebc8109..b35009b83 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -620,32 +620,15 @@ def predict_next_decode_tkv(self, running_requests: list[Request]) -> int: # Step 1: Find the maximum number of blocks across all requests # Account for requests that will need a new block after the next token max_n_blocks = 0 + num_blocks_per_req: list[int] = [] for request in running_requests: - block_ids_per_kv_cache_group = self.kv_cache_manager.get_block_ids(request.request_id) - assert len(block_ids_per_kv_cache_group) == 1 - num_blocks = len(block_ids_per_kv_cache_group[0]) - - # Check if the next token will require a new block - next_token_count = request.num_computed_tokens + 1 - if next_token_count % self.block_size == 1: - # The next token will fill the current block and require a new one - num_blocks += 1 - + num_blocks = math.ceil((request.num_computed_tokens + 1) / self.block_size) + num_blocks_per_req.append(num_blocks) max_n_blocks = max(max_n_blocks, num_blocks) # Step 2: Calculate TKV for each request and find the maximum max_tkv = 0 - for request in running_requests: - # Get the number of blocks for this request - block_ids_per_kv_cache_group = self.kv_cache_manager.get_block_ids(request.request_id) - num_blocks = len(block_ids_per_kv_cache_group[0]) - - # Check if the next token will require a new block - next_token_count = request.num_computed_tokens + 1 - if next_token_count % self.block_size == 1: - # The next token will fill the current block and require a new one - num_blocks += 1 - + for request, num_blocks in zip(running_requests, num_blocks_per_req): # Calculate left padding blocks needed for alignment left_pad_blocks_count = max_n_blocks - num_blocks left_padding = left_pad_blocks_count * self.block_size diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index 1954eb3f3..513770206 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -1450,7 +1450,6 @@ def _update_batch(self, scheduler_output: SchedulerOutput): for req_id in paused_req_ids: # Only pause if it's not a finished request (finished requests are handled separately) if req_id not in (scheduler_output.finished_req_ids or []): - logger.info("Pausing request %s from input_batch", req_id) self.input_batch.pause_request(req_id) self.paused_req_ids.add(req_id) self.input_batch.refresh_metadata() @@ -1462,7 +1461,6 @@ def _update_batch(self, scheduler_output: SchedulerOutput): for req_id in restored_req_ids: # Only restore requests that were previously paused if req_id in self.paused_req_ids and req_id in self.requests: - logger.info("Restoring paused request %s to input_batch", req_id) req_state = self.requests[req_id] self.input_batch.resume_request(req_id, req_state) self.paused_req_ids.discard(req_id) From 24691b99f6eab21190f442bef9edcd11b5c1ef64 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Thu, 4 Jun 2026 11:10:14 +0200 Subject: [PATCH 031/106] don't schedule new requests if there are paused ones MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index b35009b83..971527281 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -400,6 +400,10 @@ def can_schedule_prefill(self, request: Request) -> bool: if len(self.running) + len(self.waiting) == 0: return True + # Paused request have the priority and will be resumed if the tkv_batch limit allows it + if self.paused_decoding_requests: + return False + if not self._has_scheduling_priority(request): return False @@ -442,9 +446,9 @@ def _satisfies_first_chunk_constraints(self, request: Request) -> bool: """First chunked prefill can be scheduled only if there is space in the input batch (cond1) and in the prefill batch (cond2).""" - # TODO theoretically we could already do a chunked prefill even - # if the decode batch is full, but the current implementation of input - # batch doesn't allow to do so. + # NOTE: We could already do a chunked prefill even if the decode batch + # is full, this could potentially increase the ITL of the request + # if it then request doesn't satisfy the volumetric constraint num_running = len(self.running) cond1 = num_running + len(self.waiting) < self.max_num_running_reqs From 5395320039989faeee45807dc0f93c48c6541b76 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Thu, 4 Jun 2026 13:13:54 +0200 Subject: [PATCH 032/106] comments and docstrings cleaning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 25 ++++++++----------- .../v1/worker/spyre_model_runner.py | 1 + tests/e2e/test_spyre_cp_scheduler_steps.py | 11 +++----- tests/e2e/test_spyre_pc_scheduler_steps.py | 4 +-- 4 files changed, 17 insertions(+), 24 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 971527281..24414c97d 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -170,17 +170,13 @@ class ChunkedPrefillSpyreScheduler(SpyreScheduler): Note: all the remaining constraints need to be satisfied at the time of scheduling the last chunk of a chunked prefill - - Max model length constraint: the number of requested tokens must fit - between the maximum TKV of all the running requests and the end of - the model's context + - Volumetric constraint: the product of batch_size and current TKV + must not exceed `VLLM_DT_MAX_BATCH_TKV_LIMIT` when adding a new + request. See `check_batch_tkv_limit()` method for details. - - Volumetric constraint: the total "surface" defined by the running - requests should never exceed `VLLM_DT_MAX_BATCH_TKV_LIMIT`. See - `check_batch_tkv_limit()` method for details. - - - The surface defined by the maximum TKV of - all the running requests and the number of running requests must - not exceed the limit defined by `VLLM_DT_MAX_BATCH_TKV_LIMIT` + - Decode pausing: requests may be temporarily paused from decoding + when the batch TKV limit would be exceeded in the next decode step. + Paused requests are resumed when capacity becomes available. """ def __init__(self, *args, **kwargs) -> None: @@ -477,9 +473,8 @@ def _satisfies_last_chunk_constraints(self, request: Request) -> bool: n_blocks = math.floor(max(self.tkv, prompt_len) / self.block_size) new_req_tkv = n_blocks * self.block_size + prompt_len % self.block_size - # check that batch size x tkv is smaller than the max supported number - # Note: using max_tkv is a conservative upper bound here. For the - # optimal check we need model runner to return per sequence tkvs + # check that adding the new request to the decode batch still have + # batch size x tkv value being smaller than the accepted limit cond2 = lambda: self.check_batch_tkv_limit( request=request, new_req_tkv=new_req_tkv, @@ -508,8 +503,8 @@ def _has_scheduling_priority(self, request): def check_batch_tkv_limit(self, request: Request, new_req_tkv: int, running) -> bool: """ - Check whether adding a new sequence to the decode batch would violate - Spyre's maximum batch volume constraint for chunked prefill. + Check whether adding a new sequence to the decode batch would immediately + violate Spyre's maximum batch volume constraint for chunked prefill. In Spyre, the product of `batch_size` and the current `tkv` (tokens-per-sequence) must not exceed the limit defined by diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index 513770206..3ad1b757e 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -1431,6 +1431,7 @@ def update_states(self, scheduler_output: SchedulerOutput): def _update_batch(self, scheduler_output: SchedulerOutput): """Updates the states for the in progress batch + - Synchronizes input_batch with scheduler output (handles pause/resume) - Bumps the count of computed tokens for each request - Updates the KV cache metadata for each request - Safely removes finished requests from the batch diff --git a/tests/e2e/test_spyre_cp_scheduler_steps.py b/tests/e2e/test_spyre_cp_scheduler_steps.py index 0c5279955..0572912d8 100644 --- a/tests/e2e/test_spyre_cp_scheduler_steps.py +++ b/tests/e2e/test_spyre_cp_scheduler_steps.py @@ -491,7 +491,7 @@ def test_cp_prefill_interleave1( { # Chunk 0 of request 1 prefill "step": 3, - "tkv": 512, + "tkv": 512, # prompt len of req1 "waiting": [], "running": ["1", "0"], "request_outputs": [], @@ -508,9 +508,8 @@ def test_cp_prefill_interleave1( }, { # Chunk 1 of request 1 prefill - # tkv of decode batch (tkv not updated until last chunk) "step": 5, - "tkv": 512, + "tkv": 512, # prompt len of request "waiting": [], "running": ["1", "0"], "request_outputs": [], @@ -527,9 +526,8 @@ def test_cp_prefill_interleave1( }, { # Chunk 2 of request 1 prefill - # tkv of decode batch (tkv not updated until last chunk) "step": 7, - "tkv": 512, + "tkv": 512, # prompt len of request "waiting": [], "running": ["1", "0"], "request_outputs": [], @@ -547,9 +545,8 @@ def test_cp_prefill_interleave1( { # Chunk 3 of request 1 prefill. # First token is generated - # tkv updated for last chunk "step": 9, - "tkv": 512, + "tkv": 512, # prompt len of request "waiting": [], "running": ["1", "0"], "request_outputs": ["1"], diff --git a/tests/e2e/test_spyre_pc_scheduler_steps.py b/tests/e2e/test_spyre_pc_scheduler_steps.py index 624522884..e2422da91 100644 --- a/tests/e2e/test_spyre_pc_scheduler_steps.py +++ b/tests/e2e/test_spyre_pc_scheduler_steps.py @@ -1673,9 +1673,9 @@ def test_first_chunk_partial_match( "n_prefix_hits": 0, "block_tables": {"0": [1]}, }, - { # prefill seq 1. This step was crashing before + { # prefill seq 1 "step": 2, - "tkv": 320, + "tkv": 320, # prompt len of request 1 "waiting": [], "running": ["1", "0"], "request_outputs": [], From b39f262e44673f67c781027e98d758015b32b789 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 5 Jun 2026 18:47:09 +0200 Subject: [PATCH 033/106] log preemption MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 24414c97d..9a5bdeb62 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -567,6 +567,8 @@ def _handle_decode_requests_pausing(self) -> None: # the remaining batch fits within constraints while not self._can_decode_all_requests(decoding_requests): had_to_remove = True + + # TODO we should test different removal logics: longest request, optimize padding # Remove the request with the fewest decoded tokens # Decoded tokens = num_computed_tokens - num_prompt_tokens request_to_remove = min( @@ -575,6 +577,7 @@ def _handle_decode_requests_pausing(self) -> None: decoding_requests.remove(request_to_remove) self.running.remove(request_to_remove) self.paused_decoding_requests.append(request_to_remove) + logger.info("Request %s paused due to batch TKV limit ", request_to_remove.request_id) # It shouldn't be possible to remove all requests if we started with some assert not initial_had_requests or len(decoding_requests) > 0 @@ -592,6 +595,10 @@ def _handle_decode_requests_pausing(self) -> None: self.paused_decoding_requests.pop(0) self.running.append(request_to_add) decoding_requests.append(request_to_add) + logger.info( + "Request %s resumed (batch TKV capacity available).", + request_to_add.request_id, + ) else: # Can't add any more requests break From 82167b5fa924cbb88b50fd48d4a35ff43806b849 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 9 Jun 2026 20:41:32 +0200 Subject: [PATCH 034/106] custom bench serve: first version MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- pyproject.toml | 3 + sendnn_inference/benchmarks/__init__.py | 0 .../benchmarks/spyre_bench_serve.py | 136 ++++++++++++++++++ .../benchmarks/spyre_request_func.py | 127 ++++++++++++++++ sendnn_inference/envs.py | 8 ++ sendnn_inference/platform.py | 5 +- sendnn_inference/v1/core/scheduler.py | 26 ++++ sendnn_inference/v1/metrics/__init__.py | 8 +- sendnn_inference/v1/metrics/patch_serving.py | 61 ++++++++ sendnn_inference/v1/metrics/stats_logger.py | 77 ++++++++++ .../v1/worker/spyre_model_runner.py | 14 +- 11 files changed, 462 insertions(+), 3 deletions(-) create mode 100644 sendnn_inference/benchmarks/__init__.py create mode 100644 sendnn_inference/benchmarks/spyre_bench_serve.py create mode 100644 sendnn_inference/benchmarks/spyre_request_func.py create mode 100644 sendnn_inference/v1/metrics/patch_serving.py diff --git a/pyproject.toml b/pyproject.toml index b0325dca6..a63cde89d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,6 +27,9 @@ dynamic = ["version"] [project.entry-points."vllm.platform_plugins"] sendnn_inference = "sendnn_inference:register" +[project.scripts] +sendnn-bench = "sendnn_inference.benchmarks.spyre_bench_serve:main" + [tool.setuptools.packages.find] where = ["."] # list of folders that contain the packages (["."] by default) include = ["sendnn_inference*"] # package names should match these glob patterns (["*"] by default) diff --git a/sendnn_inference/benchmarks/__init__.py b/sendnn_inference/benchmarks/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py new file mode 100644 index 000000000..b53cc0aac --- /dev/null +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +"""sendnn-bench serve — vllm bench serve extended with Spyre per-request metrics. + +Usage: + sendnn-bench serve --host localhost --port 8000 --model \\ + --dataset-name random --num-prompts 20 --request-rate 2 + +Env var: + SENDNN_INFERENCE_BENCH_METRICS_ENABLED=1 (must also be set on the server) +""" + +import argparse +import asyncio +from typing import Any + +import numpy as np + +from vllm.benchmarks.lib.endpoint_request_func import ( + ASYNC_REQUEST_FUNCS, + RequestFuncInput, +) +from vllm.benchmarks.serve import add_cli_args, main_async + +from sendnn_inference.benchmarks.spyre_request_func import async_request_spyre_chat + +_BACKEND_NAME = "spyre-chat" + +# Shared accumulator — populated by the wrapper below during the benchmark run. +_spyre_metrics_collected: list[dict[str, Any]] = [] + + +def _make_collecting_func(): + """Return a wrapper around async_request_spyre_chat that accumulates + custom_metrics_dict into _spyre_metrics_collected.""" + + async def _wrapper( + request_func_input: RequestFuncInput, + session, + pbar=None, + ): + output = await async_request_spyre_chat(request_func_input, session, pbar) + if output.success and output.custom_metrics_dict: + _spyre_metrics_collected.append(output.custom_metrics_dict) + return output + + return _wrapper + + +def _register_backend() -> None: + from vllm.benchmarks.lib.endpoint_request_func import OPENAI_COMPATIBLE_BACKENDS + + ASYNC_REQUEST_FUNCS[_BACKEND_NAME] = _make_collecting_func() + # Register as an OpenAI-compatible backend so that vllm's main_async + # enables ignore_eos for random datasets and allows sampling parameters. + if _BACKEND_NAME not in OPENAI_COMPATIBLE_BACKENDS: + OPENAI_COMPATIBLE_BACKENDS.append(_BACKEND_NAME) + + +def _build_parser() -> argparse.ArgumentParser: + """Build an arg parser based on vllm's but with spyre-chat as default backend.""" + parser = argparse.ArgumentParser( + description="Spyre-extended vllm bench serve", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + # Register our backend before add_cli_args so it appears in --backend choices. + _register_backend() + add_cli_args(parser) + + # Override the default so --backend doesn't need to be specified explicitly. + for action in parser._actions: + if action.dest == "backend": + action.default = _BACKEND_NAME + break + + return parser + + +def _print_spyre_section( + metrics_list: list[dict[str, Any]], + selected_percentiles: list[float], +) -> None: + """Print Spyre-specific metrics in vllm bench serve format.""" + if not metrics_list: + return + + queue_times_ms = [m["queued_time_s"] * 1000 for m in metrics_list if "queued_time_s" in m] + num_chunks_list = [ + m["num_chunked_prefills"] for m in metrics_list if "num_chunked_prefills" in m + ] + chunk_lats_ms = [ + lat * 1000 for m in metrics_list for lat in m.get("chunk_prefill_latencies_s", []) + ] + + def _section(header: str, values: list[float], label: str) -> None: + if not values: + return + arr = np.array(values) + print("{s:{c}^{n}}".format(s=f" {header} ", n=50, c="-")) + print("{:<40} {:<10.2f}".format(f"Mean {label}:", float(np.mean(arr)))) + print("{:<40} {:<10.2f}".format(f"Median {label}:", float(np.median(arr)))) + for p in selected_percentiles: + print( + "{:<40} {:<10.2f}".format( + f"P{int(p) if int(p) == p else p} {label}:", + float(np.percentile(arr, p)), + ) + ) + + _section("Queue Wait Time", queue_times_ms, "Queue Wait Time (ms)") + _section("Chunked Prefill Count", num_chunks_list, "Num Chunked Prefills") + _section("Chunked Prefill Latency", chunk_lats_ms, "Chunk Prefill Latency (ms)") + + print("=" * 50) + + +def main() -> None: + parser = _build_parser() + args = parser.parse_args() + + # Force chat endpoint and our backend. + args.backend = _BACKEND_NAME + if not hasattr(args, "endpoint") or args.endpoint == "/v1/completions": + args.endpoint = "/v1/chat/completions" + + selected_percentiles = [float(p) for p in args.metric_percentiles.split(",")] + + _spyre_metrics_collected.clear() + + asyncio.run(main_async(args)) + + _print_spyre_section(_spyre_metrics_collected, selected_percentiles) + + +if __name__ == "__main__": + main() diff --git a/sendnn_inference/benchmarks/spyre_request_func.py b/sendnn_inference/benchmarks/spyre_request_func.py new file mode 100644 index 000000000..519f9a9c8 --- /dev/null +++ b/sendnn_inference/benchmarks/spyre_request_func.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Custom request function that captures Spyre per-request metrics injected +into the final SSE usage chunk by the server-side serving patch.""" + +import json +import sys +import traceback +import time +from dataclasses import dataclass, field +from typing import Any, Literal + +import aiohttp +from tqdm import tqdm + +from vllm.benchmarks.lib.endpoint_request_func import ( + RequestFuncInput, + RequestFuncOutput, + StreamedResponseHandler, + _get_chat_content, + _get_headers, + _update_headers_common, + _update_payload_common, + _validate_api_url, +) + + +@dataclass +class SpyreRequestFuncOutput(RequestFuncOutput): + """Extends RequestFuncOutput with Spyre-specific per-request metrics.""" + + custom_metrics_dict: dict[str, Any] = field(default_factory=dict) + + +async def async_request_spyre_chat( + request_func_input: RequestFuncInput, + session: aiohttp.ClientSession, + pbar: tqdm | None = None, + mm_position: Literal["first", "last"] = "last", +) -> SpyreRequestFuncOutput: + """Chat completions request function that additionally parses the + ``spyre_metrics`` field injected into the final SSE usage chunk.""" + + api_url = request_func_input.api_url + _validate_api_url(api_url, "OpenAI Chat Completions API", "chat/completions") + + content = _get_chat_content(request_func_input, mm_position=mm_position) + + payload = { + "model": ( + request_func_input.model_name + if request_func_input.model_name + else request_func_input.model + ), + "messages": [ + {"role": "user", "content": content}, + ], + "max_completion_tokens": request_func_input.output_len, + "stream": True, + "stream_options": { + "include_usage": True, + }, + } + _update_payload_common(payload, request_func_input) + + headers = _get_headers("application/json") + _update_headers_common(headers, request_func_input) + + output = SpyreRequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + ttft = 0.0 + st = time.perf_counter() + output.start_time = st + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload, headers=headers) as response: + if response.status == 200: + handler = StreamedResponseHandler() + async for chunk_bytes in response.content.iter_any(): + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + messages = handler.add_chunk(chunk_bytes) + for message in messages: + if message.startswith(":"): + continue + + chunk = message.removeprefix("data: ") + + if chunk != "[DONE]": + timestamp = time.perf_counter() + data = json.loads(chunk) + + if choices := data.get("choices"): + content_delta = choices[0]["delta"].get("content") + if ttft == 0.0: + ttft = timestamp - st + output.ttft = ttft + else: + output.itl.append(timestamp - most_recent_timestamp) + generated_text += content_delta or "" + elif usage := data.get("usage"): + output.output_tokens = usage.get("completion_tokens") + if (pt := usage.get("prompt_tokens")) is not None: + output.prompt_len = pt + # Parse Spyre-specific metrics from the same chunk + if spyre_metrics := data.get("spyre_metrics"): + output.custom_metrics_dict = spyre_metrics + + most_recent_timestamp = timestamp + + output.generated_text = generated_text + output.success = True + output.latency = most_recent_timestamp - st + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output diff --git a/sendnn_inference/envs.py b/sendnn_inference/envs.py index 934aa758f..676ea64c4 100644 --- a/sendnn_inference/envs.py +++ b/sendnn_inference/envs.py @@ -26,6 +26,7 @@ SENDNN_INFERENCE_MODEL_CONFIG_FILE: str | None = None SENDNN_INFERENCE_CPU_MM_DTYPE: torch.dtype = torch.float16 SENDNN_INFERENCE_MM_DEVICE: str = "auto" + SENDNN_INFERENCE_BENCH_METRICS_ENABLED: bool = False logger = init_logger(__name__) @@ -171,6 +172,13 @@ def clear_env_cache(): "SENDNN_INFERENCE_MM_DEVICE": lambda: parse_mm_device( os.getenv("SENDNN_INFERENCE_MM_DEVICE", "auto") ), + # Enable collection of per-request Spyre-specific benchmark metrics + # (queue wait time, chunked prefill count and latencies). Only needed + # when running `sendnn-bench serve`. Disabled by default to avoid + # overhead in production deployments. + "SENDNN_INFERENCE_BENCH_METRICS_ENABLED": lambda: bool( + int(os.environ.get("SENDNN_INFERENCE_BENCH_METRICS_ENABLED", "0")) + ), } # --8<-- [end:env-vars-definition] diff --git a/sendnn_inference/platform.py b/sendnn_inference/platform.py index dbedb0d6f..012d31f74 100644 --- a/sendnn_inference/platform.py +++ b/sendnn_inference/platform.py @@ -215,10 +215,13 @@ def get_total_spyre_blocks(cls, vllm_config: VllmConfig) -> int: @classmethod def check_and_update_config(cls, vllm_config: VllmConfig) -> None: # 🌶️🌶️🌶️ Patch in our perf logger before the engine is created - from sendnn_inference.v1.metrics import patch_async_llm_stat_loggers + from sendnn_inference.v1.metrics import patch_async_llm_stat_loggers, patch_serving patch_async_llm_stat_loggers() + if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED: + patch_serving() + # In case vllm passes a default vllm_config to us. # This happens when get_current_vllm_config is called # without setting the vllm config through diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 9b1cf723f..56d4f3a07 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -203,6 +203,15 @@ def __init__(self, *args, **kwargs) -> None: self.block_size = SpyrePlatform.get_block_size() self.max_batch_tkv_limit = SpyrePlatform.get_max_batch_tkv_limit() + # Per-request chunk prefill latency accumulator. + # Only populated when SENDNN_INFERENCE_BENCH_METRICS_ENABLED is set. + self._chunk_latencies: dict[str, list[float]] = {} + + if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED: + from sendnn_inference.v1.metrics.stats_logger import register_scheduler + + register_scheduler(self) + assert self.max_batch_tkv_limit != -1, ( "Expecting the env var VLLM_DT_MAX_BATCH_TKV_LIMIT to be set in platform.py" ) @@ -229,6 +238,13 @@ def update_from_output(self, scheduler_output, model_runner_output): prefix_cache_len=prefix_cache_len, ) + # Accumulate per-chunk timings before removing completed prefills + if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED: + for req in self.ongoing_prefills: + t = model_runner_output.chunk_prefill_time_s.get(req.request_id) + if t is not None: + self._chunk_latencies.setdefault(req.request_id, []).append(t) + # Remove completed prefills self.ongoing_prefills = [ req for req in self.ongoing_prefills if req.num_computed_tokens < req.num_prompt_tokens @@ -237,6 +253,16 @@ def update_from_output(self, scheduler_output, model_runner_output): self.tkv = model_runner_output.tkv return super(SpyreScheduler, self).update_from_output(scheduler_output, model_runner_output) + def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: + """Return and clear accumulated chunk timing for a finished request.""" + lats = self._chunk_latencies.pop(req_id, None) + if lats is None: + return None + return { + "num_chunked_prefills": len(lats), + "chunk_prefill_latencies_s": lats, + } + def adjust_computed_tokens( self, computed_tokens: int, left_padding: int, prefix_cache_len: int ) -> int: diff --git a/sendnn_inference/v1/metrics/__init__.py b/sendnn_inference/v1/metrics/__init__.py index 5b25355e3..af1d6f4b0 100644 --- a/sendnn_inference/v1/metrics/__init__.py +++ b/sendnn_inference/v1/metrics/__init__.py @@ -1,3 +1,9 @@ +from .patch_serving import patch_serving from .stats_logger import FileStatLogger, file_stat_logger_factory, patch_async_llm_stat_loggers -__all__ = ["patch_async_llm_stat_loggers", "file_stat_logger_factory", "FileStatLogger"] +__all__ = [ + "patch_async_llm_stat_loggers", + "patch_serving", + "file_stat_logger_factory", + "FileStatLogger", +] diff --git a/sendnn_inference/v1/metrics/patch_serving.py b/sendnn_inference/v1/metrics/patch_serving.py new file mode 100644 index 000000000..fcab9c8e5 --- /dev/null +++ b/sendnn_inference/v1/metrics/patch_serving.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Patch OpenAIServingChat to inject Spyre per-request metrics into the final +SSE usage chunk when SENDNN_INFERENCE_BENCH_METRICS_ENABLED is set. + +The final streaming chunk (empty choices, populated usage) is the natural +carrier for per-request metadata because the bench client already parses it. +We intercept only that one chunk per request (one json.loads + json.dumps), +so overhead is negligible. +""" + +import dataclasses +import json + +from vllm.logger import init_logger + +from sendnn_inference.v1.metrics.stats_logger import get_registry + +logger = init_logger(__name__) + +_patched = False + + +def patch_serving() -> None: + """Wrap OpenAIServingChat.chat_completion_stream_generator to inject + spyre_metrics into the final SSE usage chunk. Idempotent.""" + global _patched + if _patched: + return + + try: + from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat + except ImportError: + logger.warning("Could not import OpenAIServingChat — serving patch skipped") + return + + _original = OpenAIServingChat.chat_completion_stream_generator + + async def _patched_generator(self, request, result_generator, request_id, *args, **kwargs): + registry = get_registry() + async for chunk in _original(self, request, result_generator, request_id, *args, **kwargs): + if ( + registry is not None + and isinstance(chunk, str) + and '"usage"' in chunk + and '"choices":[]' in chunk + ): + try: + prefix = "data: " + data_str = chunk.removeprefix(prefix).rstrip("\n") + data = json.loads(data_str) + metrics = registry.get_and_clear(request_id) + if metrics: + data["spyre_metrics"] = dataclasses.asdict(metrics) + chunk = f"{prefix}{json.dumps(data)}\n\n" + except (json.JSONDecodeError, KeyError, TypeError): + pass # yield original chunk unchanged on any error + yield chunk + + OpenAIServingChat.chat_completion_stream_generator = _patched_generator # ty: ignore[invalid-assignment] + _patched = True + logger.debug("Spyre serving patch applied: spyre_metrics will be injected in final SSE chunk") diff --git a/sendnn_inference/v1/metrics/stats_logger.py b/sendnn_inference/v1/metrics/stats_logger.py index 239706422..7c5967b81 100644 --- a/sendnn_inference/v1/metrics/stats_logger.py +++ b/sendnn_inference/v1/metrics/stats_logger.py @@ -1,9 +1,11 @@ import dataclasses import json +import threading import time from datetime import datetime from functools import wraps from pathlib import Path +from typing import Any from vllm.config import VllmConfig from vllm.logger import init_logger @@ -21,6 +23,59 @@ logger = init_logger(__name__) +# --------------------------------------------------------------------------- +# Bench metrics registry — only active when SENDNN_INFERENCE_BENCH_METRICS_ENABLED +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class SpyreRequestMetrics: + """Per-request Spyre-specific metrics surfaced to benchmark clients.""" + + request_id: str + queued_time_s: float + num_chunked_prefills: int + chunk_prefill_latencies_s: list[float] + + +class SpyreMetricsRegistry: + """Thread-safe store of per-request SpyreRequestMetrics, cleared on read.""" + + def __init__(self) -> None: + self._store: dict[str, SpyreRequestMetrics] = {} + self._lock = threading.Lock() + + def put(self, metrics: SpyreRequestMetrics) -> None: + with self._lock: + self._store[metrics.request_id] = metrics + + def get_and_clear(self, request_id: str) -> SpyreRequestMetrics | None: + with self._lock: + return self._store.pop(request_id, None) + + +_REGISTRY: SpyreMetricsRegistry | None = None +_SCHEDULER: Any = None # set by ChunkedPrefillSpyreScheduler.__init__ + + +def enable_registry() -> SpyreMetricsRegistry: + global _REGISTRY + _REGISTRY = SpyreMetricsRegistry() + return _REGISTRY + + +def get_registry() -> SpyreMetricsRegistry | None: + return _REGISTRY + + +def register_scheduler(scheduler: Any) -> None: + """Called by ChunkedPrefillSpyreScheduler at init time.""" + global _SCHEDULER, _REGISTRY + _SCHEDULER = scheduler + if _REGISTRY is None: + _REGISTRY = SpyreMetricsRegistry() + + @dataclasses.dataclass class PerfRecord: """A record for request_metrics.jsonl. @@ -138,6 +193,28 @@ def record( ) records_to_write.append(record.to_json()) + # Feed the bench metrics registry when enabled + if ( + envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED + and _REGISTRY is not None + and r.request_id + ): + chunk_stats = ( + _SCHEDULER.get_and_clear_chunk_stats(r.request_id) if _SCHEDULER else None + ) + _REGISTRY.put( + SpyreRequestMetrics( + request_id=r.request_id, + queued_time_s=r.queued_time, + num_chunked_prefills=chunk_stats["num_chunked_prefills"] + if chunk_stats + else 0, + chunk_prefill_latencies_s=chunk_stats["chunk_prefill_latencies_s"] + if chunk_stats + else [], + ) + ) + self.open_file_pointer.write("\n".join(records_to_write) + "\n") self.open_file_pointer.flush() diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index d7accd3b2..7e7721f92 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -98,6 +98,9 @@ class SpyreModelRunnerOutput(ModelRunnerOutput): # available than the number of scheduled tokens. In that case, the scheduler # needs to update its state to reflect the correct number of computed tokens prefix_cache_hit_len: dict[str, int] = field(default_factory=dict) + # Per-chunk prefill wall-clock time in seconds, keyed by request_id. + # Only populated when SENDNN_INFERENCE_BENCH_METRICS_ENABLED is set. + chunk_prefill_time_s: dict[str, float] = field(default_factory=dict) InputBatchT = TypeVar("InputBatchT", bound=BaseInputBatch) @@ -1567,7 +1570,12 @@ def execute_model( t1 = time.time() - t0 logger.debug("t_forward_pass: %.2fms [prefill single chunk][batch size 1]", (t1 * 1000)) - return self.prefill_output() + output = self.prefill_output() + if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED: + req_ids = list(scheduler_output.num_scheduled_tokens) + if req_ids: + output.chunk_prefill_time_s[req_ids[0]] = t1 + return output # Apply grammar bitmask for structured output requests. self.apply_grammar_bitmask( @@ -1611,6 +1619,10 @@ def execute_model( return self.get_empty_output() model_output = self.sampled_output(output, is_prefill) + if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED and is_prefill: + req_ids = list(scheduler_output.num_scheduled_tokens) + if req_ids: + model_output.chunk_prefill_time_s[req_ids[0]] = t1 return model_output def prefill_output(self) -> SpyreModelRunnerOutput: From b523729155ea31b50037e4db27fa3c56a85970bb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 9 Jun 2026 21:35:37 +0200 Subject: [PATCH 035/106] allow to use sendnn-bench serve MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/benchmarks/spyre_bench_serve.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index b53cc0aac..fe2638da6 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -115,8 +115,15 @@ def _section(header: str, values: list[float], label: str) -> None: def main() -> None: + import sys + + # Allow `sendnn-bench serve ` as an alias (the word "serve" is ignored). + argv = sys.argv[1:] + if argv and argv[0] == "serve": + argv = argv[1:] + parser = _build_parser() - args = parser.parse_args() + args = parser.parse_args(argv) # Force chat endpoint and our backend. args.backend = _BACKEND_NAME From e850da505a1055c23d870c5c4f61c4d575087c0b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 9 Jun 2026 21:48:42 +0200 Subject: [PATCH 036/106] add some debug prints MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_request_func.py | 7 +++- sendnn_inference/v1/metrics/patch_serving.py | 41 +++++++++++-------- sendnn_inference/v1/metrics/stats_logger.py | 26 +++++++----- 3 files changed, 46 insertions(+), 28 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_request_func.py b/sendnn_inference/benchmarks/spyre_request_func.py index 519f9a9c8..fcd195735 100644 --- a/sendnn_inference/benchmarks/spyre_request_func.py +++ b/sendnn_inference/benchmarks/spyre_request_func.py @@ -106,7 +106,12 @@ async def async_request_spyre_chat( if (pt := usage.get("prompt_tokens")) is not None: output.prompt_len = pt # Parse Spyre-specific metrics from the same chunk - if spyre_metrics := data.get("spyre_metrics"): + spyre_metrics = data.get("spyre_metrics") + print( + f"[SPYRE DEBUG client] usage chunk keys: {list(data.keys())}, spyre_metrics present: {spyre_metrics is not None}", + flush=True, + ) + if spyre_metrics: output.custom_metrics_dict = spyre_metrics most_recent_timestamp = timestamp diff --git a/sendnn_inference/v1/metrics/patch_serving.py b/sendnn_inference/v1/metrics/patch_serving.py index fcab9c8e5..87b733301 100644 --- a/sendnn_inference/v1/metrics/patch_serving.py +++ b/sendnn_inference/v1/metrics/patch_serving.py @@ -37,25 +37,34 @@ def patch_serving() -> None: async def _patched_generator(self, request, result_generator, request_id, *args, **kwargs): registry = get_registry() + print( + f"[SPYRE DEBUG server] _patched_generator called, request_id={request_id}, registry={registry is not None}", + flush=True, + ) async for chunk in _original(self, request, result_generator, request_id, *args, **kwargs): - if ( - registry is not None - and isinstance(chunk, str) - and '"usage"' in chunk - and '"choices":[]' in chunk - ): - try: - prefix = "data: " - data_str = chunk.removeprefix(prefix).rstrip("\n") - data = json.loads(data_str) - metrics = registry.get_and_clear(request_id) - if metrics: - data["spyre_metrics"] = dataclasses.asdict(metrics) - chunk = f"{prefix}{json.dumps(data)}\n\n" - except (json.JSONDecodeError, KeyError, TypeError): - pass # yield original chunk unchanged on any error + if isinstance(chunk, str) and '"usage"' in chunk and '"choices":[]' in chunk: + print( + f"[SPYRE DEBUG server] final usage chunk detected, registry={registry is not None}", + flush=True, + ) + if registry is not None: + try: + prefix = "data: " + data_str = chunk.removeprefix(prefix).rstrip("\n") + data = json.loads(data_str) + metrics = registry.get_and_clear(request_id) + print( + f"[SPYRE DEBUG server] registry.get_and_clear({request_id!r}) -> {metrics}", + flush=True, + ) + if metrics: + data["spyre_metrics"] = dataclasses.asdict(metrics) + chunk = f"{prefix}{json.dumps(data)}\n\n" + except (json.JSONDecodeError, KeyError, TypeError) as e: + print(f"[SPYRE DEBUG server] exception injecting metrics: {e}", flush=True) yield chunk OpenAIServingChat.chat_completion_stream_generator = _patched_generator # ty: ignore[invalid-assignment] _patched = True + print("[SPYRE DEBUG server] patch_serving() applied successfully", flush=True) logger.debug("Spyre serving patch applied: spyre_metrics will be injected in final SSE chunk") diff --git a/sendnn_inference/v1/metrics/stats_logger.py b/sendnn_inference/v1/metrics/stats_logger.py index 7c5967b81..4deed7db7 100644 --- a/sendnn_inference/v1/metrics/stats_logger.py +++ b/sendnn_inference/v1/metrics/stats_logger.py @@ -194,6 +194,12 @@ def record( records_to_write.append(record.to_json()) # Feed the bench metrics registry when enabled + print( + f"[SPYRE DEBUG server] FileStatLogger.record:" + f" BENCH_METRICS_ENABLED={envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED}," + f" _REGISTRY={_REGISTRY is not None}, request_id={r.request_id!r}", + flush=True, + ) if ( envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED and _REGISTRY is not None @@ -202,18 +208,16 @@ def record( chunk_stats = ( _SCHEDULER.get_and_clear_chunk_stats(r.request_id) if _SCHEDULER else None ) - _REGISTRY.put( - SpyreRequestMetrics( - request_id=r.request_id, - queued_time_s=r.queued_time, - num_chunked_prefills=chunk_stats["num_chunked_prefills"] - if chunk_stats - else 0, - chunk_prefill_latencies_s=chunk_stats["chunk_prefill_latencies_s"] - if chunk_stats - else [], - ) + m = SpyreRequestMetrics( + request_id=r.request_id, + queued_time_s=r.queued_time, + num_chunked_prefills=chunk_stats["num_chunked_prefills"] if chunk_stats else 0, + chunk_prefill_latencies_s=chunk_stats["chunk_prefill_latencies_s"] + if chunk_stats + else [], ) + print(f"[SPYRE DEBUG server] putting in registry: {m}", flush=True) + _REGISTRY.put(m) self.open_file_pointer.write("\n".join(records_to_write) + "\n") self.open_file_pointer.flush() From acfc886a3b8d3fe4f8e22438e611bb8972b5f0d8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 9 Jun 2026 22:13:58 +0200 Subject: [PATCH 037/106] bugfix processes isolations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 37 ++++++++++ sendnn_inference/v1/metrics/patch_serving.py | 73 +++++++++++--------- sendnn_inference/v1/metrics/stats_logger.py | 26 ------- 3 files changed, 77 insertions(+), 59 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 56d4f3a07..c24b52e6b 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 import math +import time from collections import deque from typing import TYPE_CHECKING, Iterable, Union @@ -206,6 +207,9 @@ def __init__(self, *args, **kwargs) -> None: # Per-request chunk prefill latency accumulator. # Only populated when SENDNN_INFERENCE_BENCH_METRICS_ENABLED is set. self._chunk_latencies: dict[str, list[float]] = {} + # Timestamps for queue-wait-time calculation. + self._arrival_ts: dict[str, float] = {} + self._first_scheduled_ts: dict[str, float] = {} if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED: from sendnn_inference.v1.metrics.stats_logger import register_scheduler @@ -244,6 +248,12 @@ def update_from_output(self, scheduler_output, model_runner_output): t = model_runner_output.chunk_prefill_time_s.get(req.request_id) if t is not None: self._chunk_latencies.setdefault(req.request_id, []).append(t) + # Track first-scheduled time and arrival time for queue-wait calculation + now = time.time() + for req in self.ongoing_prefills: + if req.request_id not in self._first_scheduled_ts: + self._first_scheduled_ts[req.request_id] = now + self._arrival_ts[req.request_id] = req.arrival_time # Remove completed prefills self.ongoing_prefills = [ @@ -263,6 +273,33 @@ def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: "chunk_prefill_latencies_s": lats, } + def _free_request(self, request, delay_free_blocks: bool = False): + """Override to inject Spyre bench metrics into kv_transfer_params so + they travel over ZMQ to the API server process in EngineCoreOutput.""" + kv_xfer_params = super()._free_request(request, delay_free_blocks) + + if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED: + req_id = request.request_id + chunk_stats = self.get_and_clear_chunk_stats(req_id) + first_ts = self._first_scheduled_ts.pop(req_id, None) + arrival_ts = self._arrival_ts.pop(req_id, None) + queued_time_s = ( + (first_ts - arrival_ts) if first_ts is not None and arrival_ts is not None else 0.0 + ) + spyre_data = { + "queued_time_s": queued_time_s, + "num_chunked_prefills": chunk_stats["num_chunked_prefills"] if chunk_stats else 0, + "chunk_prefill_latencies_s": chunk_stats["chunk_prefill_latencies_s"] + if chunk_stats + else [], + } + if kv_xfer_params is None: + kv_xfer_params = {"__spyre__": spyre_data} + else: + kv_xfer_params["__spyre__"] = spyre_data + + return kv_xfer_params + def adjust_computed_tokens( self, computed_tokens: int, left_padding: int, prefix_cache_len: int ) -> int: diff --git a/sendnn_inference/v1/metrics/patch_serving.py b/sendnn_inference/v1/metrics/patch_serving.py index 87b733301..6b79e0104 100644 --- a/sendnn_inference/v1/metrics/patch_serving.py +++ b/sendnn_inference/v1/metrics/patch_serving.py @@ -2,19 +2,17 @@ """Patch OpenAIServingChat to inject Spyre per-request metrics into the final SSE usage chunk when SENDNN_INFERENCE_BENCH_METRICS_ENABLED is set. -The final streaming chunk (empty choices, populated usage) is the natural -carrier for per-request metadata because the bench client already parses it. -We intercept only that one chunk per request (one json.loads + json.dumps), -so overhead is negligible. +Metrics are carried from the engine process to the API server process via +RequestOutput.kv_transfer_params["__spyre__"], which already travels over the +ZMQ IPC channel. The patched generator intercepts the result_generator to +capture the final RequestOutput, then injects the metrics into the final SSE +usage chunk before yielding it. """ -import dataclasses import json from vllm.logger import init_logger -from sendnn_inference.v1.metrics.stats_logger import get_registry - logger = init_logger(__name__) _patched = False @@ -36,32 +34,41 @@ def patch_serving() -> None: _original = OpenAIServingChat.chat_completion_stream_generator async def _patched_generator(self, request, result_generator, request_id, *args, **kwargs): - registry = get_registry() - print( - f"[SPYRE DEBUG server] _patched_generator called, request_id={request_id}, registry={registry is not None}", - flush=True, - ) - async for chunk in _original(self, request, result_generator, request_id, *args, **kwargs): - if isinstance(chunk, str) and '"usage"' in chunk and '"choices":[]' in chunk: - print( - f"[SPYRE DEBUG server] final usage chunk detected, registry={registry is not None}", - flush=True, - ) - if registry is not None: - try: - prefix = "data: " - data_str = chunk.removeprefix(prefix).rstrip("\n") - data = json.loads(data_str) - metrics = registry.get_and_clear(request_id) - print( - f"[SPYRE DEBUG server] registry.get_and_clear({request_id!r}) -> {metrics}", - flush=True, - ) - if metrics: - data["spyre_metrics"] = dataclasses.asdict(metrics) - chunk = f"{prefix}{json.dumps(data)}\n\n" - except (json.JSONDecodeError, KeyError, TypeError) as e: - print(f"[SPYRE DEBUG server] exception injecting metrics: {e}", flush=True) + # Wrap result_generator to capture the final RequestOutput's kv_transfer_params. + spyre_metrics: dict | None = None + + async def _capturing_generator(): + nonlocal spyre_metrics + async for res in result_generator: + if res.finished and res.kv_transfer_params: + spyre_metrics = res.kv_transfer_params.get("__spyre__") + print( + f"[SPYRE DEBUG server] captured spyre_metrics from res: {spyre_metrics}", + flush=True, + ) + yield res + + async for chunk in _original( + self, request, _capturing_generator(), request_id, *args, **kwargs + ): + if ( + spyre_metrics is not None + and isinstance(chunk, str) + and '"usage"' in chunk + and '"choices":[]' in chunk + ): + try: + prefix = "data: " + data_str = chunk.removeprefix(prefix).rstrip("\n") + data = json.loads(data_str) + data["spyre_metrics"] = spyre_metrics + chunk = f"{prefix}{json.dumps(data)}\n\n" + print( + f"[SPYRE DEBUG server] injected spyre_metrics into usage chunk for {request_id}", + flush=True, + ) + except (json.JSONDecodeError, KeyError, TypeError) as e: + print(f"[SPYRE DEBUG server] exception injecting metrics: {e}", flush=True) yield chunk OpenAIServingChat.chat_completion_stream_generator = _patched_generator # ty: ignore[invalid-assignment] diff --git a/sendnn_inference/v1/metrics/stats_logger.py b/sendnn_inference/v1/metrics/stats_logger.py index 4deed7db7..8b4ea97bb 100644 --- a/sendnn_inference/v1/metrics/stats_logger.py +++ b/sendnn_inference/v1/metrics/stats_logger.py @@ -193,32 +193,6 @@ def record( ) records_to_write.append(record.to_json()) - # Feed the bench metrics registry when enabled - print( - f"[SPYRE DEBUG server] FileStatLogger.record:" - f" BENCH_METRICS_ENABLED={envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED}," - f" _REGISTRY={_REGISTRY is not None}, request_id={r.request_id!r}", - flush=True, - ) - if ( - envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED - and _REGISTRY is not None - and r.request_id - ): - chunk_stats = ( - _SCHEDULER.get_and_clear_chunk_stats(r.request_id) if _SCHEDULER else None - ) - m = SpyreRequestMetrics( - request_id=r.request_id, - queued_time_s=r.queued_time, - num_chunked_prefills=chunk_stats["num_chunked_prefills"] if chunk_stats else 0, - chunk_prefill_latencies_s=chunk_stats["chunk_prefill_latencies_s"] - if chunk_stats - else [], - ) - print(f"[SPYRE DEBUG server] putting in registry: {m}", flush=True) - _REGISTRY.put(m) - self.open_file_pointer.write("\n".join(records_to_write) + "\n") self.open_file_pointer.flush() From a6a4590f433ed4be7baa6786e52c2894b3dbdbc0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 9 Jun 2026 22:46:58 +0200 Subject: [PATCH 038/106] bug fixed, remove prints MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/benchmarks/spyre_bench_serve.py | 13 +++++++++++-- sendnn_inference/benchmarks/spyre_request_func.py | 7 +------ sendnn_inference/v1/metrics/patch_serving.py | 13 +++---------- 3 files changed, 15 insertions(+), 18 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index fe2638da6..3bcda2001 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -11,6 +11,7 @@ import argparse import asyncio +import logging from typing import Any import numpy as np @@ -23,6 +24,8 @@ from sendnn_inference.benchmarks.spyre_request_func import async_request_spyre_chat +logger = logging.getLogger(__name__) + _BACKEND_NAME = "spyre-chat" # Shared accumulator — populated by the wrapper below during the benchmark run. @@ -39,8 +42,14 @@ async def _wrapper( pbar=None, ): output = await async_request_spyre_chat(request_func_input, session, pbar) - if output.success and output.custom_metrics_dict: - _spyre_metrics_collected.append(output.custom_metrics_dict) + if output.success: + if output.custom_metrics_dict: + _spyre_metrics_collected.append(output.custom_metrics_dict) + else: + logger.warning( + "Spyre metrics absent from response — is " + "SENDNN_INFERENCE_BENCH_METRICS_ENABLED set on the server?" + ) return output return _wrapper diff --git a/sendnn_inference/benchmarks/spyre_request_func.py b/sendnn_inference/benchmarks/spyre_request_func.py index fcd195735..519f9a9c8 100644 --- a/sendnn_inference/benchmarks/spyre_request_func.py +++ b/sendnn_inference/benchmarks/spyre_request_func.py @@ -106,12 +106,7 @@ async def async_request_spyre_chat( if (pt := usage.get("prompt_tokens")) is not None: output.prompt_len = pt # Parse Spyre-specific metrics from the same chunk - spyre_metrics = data.get("spyre_metrics") - print( - f"[SPYRE DEBUG client] usage chunk keys: {list(data.keys())}, spyre_metrics present: {spyre_metrics is not None}", - flush=True, - ) - if spyre_metrics: + if spyre_metrics := data.get("spyre_metrics"): output.custom_metrics_dict = spyre_metrics most_recent_timestamp = timestamp diff --git a/sendnn_inference/v1/metrics/patch_serving.py b/sendnn_inference/v1/metrics/patch_serving.py index 6b79e0104..080d37bde 100644 --- a/sendnn_inference/v1/metrics/patch_serving.py +++ b/sendnn_inference/v1/metrics/patch_serving.py @@ -42,10 +42,6 @@ async def _capturing_generator(): async for res in result_generator: if res.finished and res.kv_transfer_params: spyre_metrics = res.kv_transfer_params.get("__spyre__") - print( - f"[SPYRE DEBUG server] captured spyre_metrics from res: {spyre_metrics}", - flush=True, - ) yield res async for chunk in _original( @@ -63,15 +59,12 @@ async def _capturing_generator(): data = json.loads(data_str) data["spyre_metrics"] = spyre_metrics chunk = f"{prefix}{json.dumps(data)}\n\n" - print( - f"[SPYRE DEBUG server] injected spyre_metrics into usage chunk for {request_id}", - flush=True, - ) except (json.JSONDecodeError, KeyError, TypeError) as e: - print(f"[SPYRE DEBUG server] exception injecting metrics: {e}", flush=True) + logger.warning( + "Failed to inject spyre_metrics into SSE chunk for %s: %s", request_id, e + ) yield chunk OpenAIServingChat.chat_completion_stream_generator = _patched_generator # ty: ignore[invalid-assignment] _patched = True - print("[SPYRE DEBUG server] patch_serving() applied successfully", flush=True) logger.debug("Spyre serving patch applied: spyre_metrics will be injected in final SSE chunk") From 2250d9e663790431b3586e08265e734e73ecf8a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 9 Jun 2026 22:55:41 +0200 Subject: [PATCH 039/106] cleaner output print MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_bench_serve.py | 42 ++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 3bcda2001..8be0081dc 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -124,6 +124,7 @@ def _section(header: str, values: list[float], label: str) -> None: def main() -> None: + import io import sys # Allow `sendnn-bench serve ` as an alias (the word "serve" is ignored). @@ -143,10 +144,49 @@ def main() -> None: _spyre_metrics_collected.clear() - asyncio.run(main_async(args)) + # Capture any stdout noise (warnings, plot-saved lines) that vllm emits after + # its metrics table so we can print it after the SenDNN section. + _vllm_stdout_buf = io.StringIO() + _vllm_stderr_buf = io.StringIO() + _orig_stdout, _orig_stderr = sys.stdout, sys.stderr + + # Split stdout: lines that look like the vllm table go straight through; + # everything else is buffered. + class _SplitWriter: + def __init__(self, passthrough, buf): + self._passthrough = passthrough + self._buf = buf + self._in_table = True # vllm table comes first + + def write(self, s): + if self._in_table: + self._passthrough.write(s) + # Once we see the closing "===...===" the table is done. + if s.strip() == "=" * 50: + self._in_table = False + else: + self._buf.write(s) + + def flush(self): + self._passthrough.flush() + + sys.stdout = _SplitWriter(_orig_stdout, _vllm_stdout_buf) + sys.stderr = _vllm_stderr_buf + try: + asyncio.run(main_async(args)) + finally: + sys.stdout = _orig_stdout + sys.stderr = _orig_stderr + print("\n" + "=" * 50) + print("{s:{c}^{n}}".format(s=" SenDNN Metrics ", n=50, c="=")) + print("=" * 50) _print_spyre_section(_spyre_metrics_collected, selected_percentiles) + trailing = _vllm_stdout_buf.getvalue() + _vllm_stderr_buf.getvalue() + if trailing.strip(): + print(trailing, end="") + if __name__ == "__main__": main() From baf53882953dc5f254f0c43191633d290c687b5b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 10 Jun 2026 13:50:18 +0200 Subject: [PATCH 040/106] pretty print final output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_bench_serve.py | 86 +++++++++++-------- 1 file changed, 49 insertions(+), 37 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 8be0081dc..4c489e3b5 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -123,10 +123,56 @@ def _section(header: str, values: list[float], label: str) -> None: print("=" * 50) -def main() -> None: +def _run_vllm_and_capture_trailing(args: Any) -> tuple[str, str]: + """Run vllm's main_async, letting stdout/stderr pass through live until the + closing '=' * 50 line of the metrics table. Everything written after that + marker is captured and returned as (stdout_trailing, stderr_trailing).""" import io import sys + stdout_buf = io.StringIO() + stderr_buf = io.StringIO() + orig_stdout, orig_stderr = sys.stdout, sys.stderr + done = {"v": False} + + class _StdoutSplitter: + def write(self, s): + if not done["v"]: + if s.strip() == "=" * 50: + done["v"] = True + else: + orig_stdout.write(s) + else: + stdout_buf.write(s) + + def flush(self): + orig_stdout.flush() + + class _StderrSplitter: + def write(self, s): + if not done["v"]: + orig_stderr.write(s) + else: + stderr_buf.write(s) + + def flush(self): + if not done["v"]: + orig_stderr.flush() + + sys.stdout = _StdoutSplitter() + sys.stderr = _StderrSplitter() + try: + asyncio.run(main_async(args)) + finally: + sys.stdout = orig_stdout + sys.stderr = orig_stderr + + return stdout_buf.getvalue(), stderr_buf.getvalue() + + +def main() -> None: + import sys + # Allow `sendnn-bench serve ` as an alias (the word "serve" is ignored). argv = sys.argv[1:] if argv and argv[0] == "serve": @@ -144,46 +190,12 @@ def main() -> None: _spyre_metrics_collected.clear() - # Capture any stdout noise (warnings, plot-saved lines) that vllm emits after - # its metrics table so we can print it after the SenDNN section. - _vllm_stdout_buf = io.StringIO() - _vllm_stderr_buf = io.StringIO() - _orig_stdout, _orig_stderr = sys.stdout, sys.stderr + stdout_trailing, stderr_trailing = _run_vllm_and_capture_trailing(args) - # Split stdout: lines that look like the vllm table go straight through; - # everything else is buffered. - class _SplitWriter: - def __init__(self, passthrough, buf): - self._passthrough = passthrough - self._buf = buf - self._in_table = True # vllm table comes first - - def write(self, s): - if self._in_table: - self._passthrough.write(s) - # Once we see the closing "===...===" the table is done. - if s.strip() == "=" * 50: - self._in_table = False - else: - self._buf.write(s) - - def flush(self): - self._passthrough.flush() - - sys.stdout = _SplitWriter(_orig_stdout, _vllm_stdout_buf) - sys.stderr = _vllm_stderr_buf - try: - asyncio.run(main_async(args)) - finally: - sys.stdout = _orig_stdout - sys.stderr = _orig_stderr - - print("\n" + "=" * 50) print("{s:{c}^{n}}".format(s=" SenDNN Metrics ", n=50, c="=")) - print("=" * 50) _print_spyre_section(_spyre_metrics_collected, selected_percentiles) - trailing = _vllm_stdout_buf.getvalue() + _vllm_stderr_buf.getvalue() + trailing = stdout_trailing + stderr_trailing if trailing.strip(): print(trailing, end="") From f8460410ee9169edf0d8181861058e7181236ecb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 10 Jun 2026 14:51:17 +0200 Subject: [PATCH 041/106] inject per request results to vllm bench serve .json file MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_bench_serve.py | 77 +++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 4c489e3b5..cfd8657c5 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -11,7 +11,10 @@ import argparse import asyncio +import json import logging +import os +import time from typing import Any import numpy as np @@ -123,6 +126,77 @@ def _section(header: str, values: list[float], label: str) -> None: print("=" * 50) +def _inject_spyre_metrics_into_result_file( + args: Any, + metrics_list: list[dict[str, Any]], + run_started_at: float, +) -> None: + """If vllm wrote a result JSON (--save-result / --append-result / --result-filename), + find it and inject per-request Spyre metric lists alongside vllm's own per-request + fields (ttfts, itls, …).""" + if not metrics_list: + return + if not ( + getattr(args, "save_result", False) + or getattr(args, "append_result", False) + or getattr(args, "result_filename", None) + ): + return + + # Locate the file vllm just wrote by finding the newest .json in the result dir + # that was modified after we started the run. + result_dir = getattr(args, "result_dir", None) or "." + explicit_name = getattr(args, "result_filename", None) + + if explicit_name: + candidate = ( + explicit_name + if os.path.isabs(explicit_name) + else os.path.join(result_dir, explicit_name) + ) + candidates = [candidate] if os.path.isfile(candidate) else [] + else: + try: + candidates = [ + os.path.join(result_dir, f) + for f in os.listdir(result_dir) + if f.endswith(".json") + and os.path.getmtime(os.path.join(result_dir, f)) >= run_started_at + ] + except OSError: + candidates = [] + + if not candidates: + logger.warning("Could not locate vllm result JSON to inject Spyre metrics into.") + return + + file_path = max(candidates, key=os.path.getmtime) + + try: + with open(file_path, encoding="utf-8") as fh: + result = json.load(fh) + except Exception as exc: + logger.warning("Failed to read vllm result JSON %s: %s", file_path, exc) + return + + result["spyre_queue_times_s"] = [ + m["queued_time_s"] for m in metrics_list if "queued_time_s" in m + ] + result["spyre_num_chunked_prefills"] = [ + m["num_chunked_prefills"] for m in metrics_list if "num_chunked_prefills" in m + ] + result["spyre_chunk_prefill_latencies_s"] = [ + m.get("chunk_prefill_latencies_s", []) for m in metrics_list + ] + + try: + with open(file_path, "w", encoding="utf-8") as fh: + json.dump(result, fh) + logger.info("Spyre metrics injected into %s", file_path) + except Exception as exc: + logger.warning("Failed to write Spyre metrics into result JSON %s: %s", file_path, exc) + + def _run_vllm_and_capture_trailing(args: Any) -> tuple[str, str]: """Run vllm's main_async, letting stdout/stderr pass through live until the closing '=' * 50 line of the metrics table. Everything written after that @@ -190,11 +264,14 @@ def main() -> None: _spyre_metrics_collected.clear() + run_started_at = time.time() stdout_trailing, stderr_trailing = _run_vllm_and_capture_trailing(args) print("{s:{c}^{n}}".format(s=" SenDNN Metrics ", n=50, c="=")) _print_spyre_section(_spyre_metrics_collected, selected_percentiles) + _inject_spyre_metrics_into_result_file(args, _spyre_metrics_collected, run_started_at) + trailing = stdout_trailing + stderr_trailing if trailing.strip(): print(trailing, end="") From cae6cc7e01a7e4e17e573a28327b5b9c2edaa49c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 10 Jun 2026 15:33:43 +0200 Subject: [PATCH 042/106] standalone metric: num prefill chunks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/benchmarks/spyre_bench_serve.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index cfd8657c5..087c49497 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -103,6 +103,10 @@ def _print_spyre_section( chunk_lats_ms = [ lat * 1000 for m in metrics_list for lat in m.get("chunk_prefill_latencies_s", []) ] + total_prefill_chunks = sum(num_chunks_list) + + # Scalar summary line (mirrors vllm's plain-count header section) + print("{:<40} {:<10}".format("Total prefill chunks processed:", total_prefill_chunks)) def _section(header: str, values: list[float], label: str) -> None: if not values: @@ -188,6 +192,7 @@ def _inject_spyre_metrics_into_result_file( result["spyre_chunk_prefill_latencies_s"] = [ m.get("chunk_prefill_latencies_s", []) for m in metrics_list ] + result["spyre_total_prefill_chunks"] = sum(result["spyre_num_chunked_prefills"]) try: with open(file_path, "w", encoding="utf-8") as fh: From 92b78a57766c73656b8d7025d9d120c417365fd6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 10 Jun 2026 18:49:50 +0200 Subject: [PATCH 043/106] add benchmark tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- tests/benchmarks/__init__.py | 0 tests/benchmarks/test_bench_metrics.py | 461 +++++++++++++++++++++++++ 2 files changed, 461 insertions(+) create mode 100644 tests/benchmarks/__init__.py create mode 100644 tests/benchmarks/test_bench_metrics.py diff --git a/tests/benchmarks/__init__.py b/tests/benchmarks/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py new file mode 100644 index 000000000..e4c8169c4 --- /dev/null +++ b/tests/benchmarks/test_bench_metrics.py @@ -0,0 +1,461 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for sendnn-bench serve custom metric pipeline. + +Tests cover four layers: + 1. _inject_spyre_metrics_into_result_file — JSON result file injection + 2. _print_spyre_section — stdout output format + 3. ChunkedPrefillSpyreScheduler accumulation — _chunk_latencies / _arrival_ts / + _first_scheduled_ts, and get_and_clear_chunk_stats + 4. async_request_spyre_chat — client-side SSE parsing +""" + +import argparse +import asyncio +import json +import pathlib +import time +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from sendnn_inference.benchmarks.spyre_bench_serve import ( + _inject_spyre_metrics_into_result_file, + _print_spyre_section, +) + +# --------------------------------------------------------------------------- +# Shared test data +# --------------------------------------------------------------------------- + +# Deliberately synthetic values — order-of-magnitude differences make clear these +# are not real system measurements. +FAKE_METRICS: list[dict[str, Any]] = [ + { + "queued_time_s": 42.0, + "num_chunked_prefills": 7, + "chunk_prefill_latencies_s": [0.001, 999.9, 0.003, 500.0, 0.002, 750.0, 1.0], + }, + { + "queued_time_s": 0.00001, + "num_chunked_prefills": 3, + "chunk_prefill_latencies_s": [12345.6, 0.0001, 99999.9], + }, +] + +SELECTED_PERCENTILES = [99.0] + + +def _write_fake_result(tmp_path) -> pathlib.Path: + p = tmp_path / "result.json" + p.write_text(json.dumps({"backend": "spyre-chat", "num_prompts": 2})) + return p + + +def _make_args(tmp_path, *, save_result: bool = True, result_filename=None): + return argparse.Namespace( + save_result=save_result, + append_result=False, + result_filename=str(result_filename) if result_filename else None, + result_dir=str(tmp_path), + ) + + +# --------------------------------------------------------------------------- +# Test 1 — _inject_spyre_metrics_into_result_file +# --------------------------------------------------------------------------- + + +@pytest.mark.cpu +def test_inject_adds_spyre_keys(tmp_path): + result_file = _write_fake_result(tmp_path) + _inject_spyre_metrics_into_result_file(_make_args(tmp_path), FAKE_METRICS, time.time() - 1) + data = json.loads(result_file.read_text()) + assert "spyre_queue_times_s" in data + assert "spyre_num_chunked_prefills" in data + assert "spyre_chunk_prefill_latencies_s" in data + assert "spyre_total_prefill_chunks" in data + + +@pytest.mark.cpu +def test_inject_values_correct(tmp_path): + result_file = _write_fake_result(tmp_path) + _inject_spyre_metrics_into_result_file(_make_args(tmp_path), FAKE_METRICS, time.time() - 1) + data = json.loads(result_file.read_text()) + + assert data["spyre_queue_times_s"] == pytest.approx([42.0, 0.00001]) + assert data["spyre_num_chunked_prefills"] == [7, 3] + assert data["spyre_total_prefill_chunks"] == 10 + assert data["spyre_chunk_prefill_latencies_s"] == [ + [0.001, 999.9, 0.003, 500.0, 0.002, 750.0, 1.0], + [12345.6, 0.0001, 99999.9], + ] + # Original keys preserved + assert data["backend"] == "spyre-chat" + + +@pytest.mark.cpu +def test_inject_noop_when_save_result_false(tmp_path): + result_file = _write_fake_result(tmp_path) + original = result_file.read_text() + _inject_spyre_metrics_into_result_file( + _make_args(tmp_path, save_result=False), FAKE_METRICS, time.time() - 1 + ) + assert result_file.read_text() == original + + +@pytest.mark.cpu +def test_inject_noop_when_metrics_empty(tmp_path): + result_file = _write_fake_result(tmp_path) + original = result_file.read_text() + _inject_spyre_metrics_into_result_file(_make_args(tmp_path), [], time.time() - 1) + assert result_file.read_text() == original + + +@pytest.mark.cpu +def test_inject_explicit_result_filename(tmp_path): + result_file = _write_fake_result(tmp_path) + args = _make_args(tmp_path, result_filename=str(result_file)) + _inject_spyre_metrics_into_result_file(args, FAKE_METRICS, time.time() - 1) + data = json.loads(result_file.read_text()) + assert "spyre_queue_times_s" in data + + +# --------------------------------------------------------------------------- +# Test 2 — _print_spyre_section +# --------------------------------------------------------------------------- + + +@pytest.mark.cpu +def test_print_scalar_total_line(capsys): + _print_spyre_section(FAKE_METRICS, SELECTED_PERCENTILES) + out = capsys.readouterr().out + assert "Total prefill chunks processed:" in out + assert "10" in out + + +@pytest.mark.cpu +def test_print_sendnn_header(capsys): + # The SenDNN header is printed by main() just before _print_spyre_section. + # We verify the format string produces the expected centred header. + header = "{s:{c}^{n}}".format(s=" SenDNN Metrics ", n=50, c="=") + assert "SenDNN Metrics" in header + assert header.startswith("=") + assert header.endswith("=") + assert len(header) == 50 + + +@pytest.mark.cpu +def test_print_section_headers(capsys): + _print_spyre_section(FAKE_METRICS, SELECTED_PERCENTILES) + out = capsys.readouterr().out + assert "Queue Wait Time" in out + assert "Chunked Prefill Count" in out + assert "Chunked Prefill Latency" in out + + +@pytest.mark.cpu +def test_print_mean_median_p99(capsys): + _print_spyre_section(FAKE_METRICS, SELECTED_PERCENTILES) + out = capsys.readouterr().out + assert "Mean Queue Wait Time (ms):" in out + assert "Median Queue Wait Time (ms):" in out + assert "P99 Queue Wait Time (ms):" in out + + +@pytest.mark.cpu +def test_print_noop_when_empty(capsys): + _print_spyre_section([], SELECTED_PERCENTILES) + out = capsys.readouterr().out + assert out == "" + + +@pytest.mark.cpu +def test_print_missing_keys_tolerated(capsys): + # Metrics without chunk_prefill_latencies_s — that section should be absent, + # but queue time and chunk count sections should still print. + metrics = [ + {"queued_time_s": 77777.7, "num_chunked_prefills": 13}, + {"queued_time_s": 0.000003, "num_chunked_prefills": 99}, + ] + _print_spyre_section(metrics, SELECTED_PERCENTILES) + out = capsys.readouterr().out + assert "Queue Wait Time" in out + assert "Chunked Prefill Count" in out + assert "Chunked Prefill Latency" not in out + + +# --------------------------------------------------------------------------- +# Test 3B — get_and_clear_chunk_stats (pure unit, no engine) +# --------------------------------------------------------------------------- + + +def _make_bare_scheduler(): + """Instantiate ChunkedPrefillSpyreScheduler with __init__ bypassed so we can + call its methods without a full vllm engine setup.""" + from sendnn_inference.v1.core.scheduler import ChunkedPrefillSpyreScheduler + + with patch.object(ChunkedPrefillSpyreScheduler, "__init__", lambda *a, **kw: None): + s = ChunkedPrefillSpyreScheduler() + s._chunk_latencies = {} + return s + + +@pytest.mark.cpu +def test_get_and_clear_returns_correct_dict(): + s = _make_bare_scheduler() + s._chunk_latencies["r0"] = [88888.8, 0.000005] + result = s.get_and_clear_chunk_stats("r0") + assert result is not None + assert result["num_chunked_prefills"] == 2 + assert result["chunk_prefill_latencies_s"] == pytest.approx([88888.8, 0.000005]) + # Entry must be cleared after retrieval + assert "r0" not in s._chunk_latencies + + +@pytest.mark.cpu +def test_get_and_clear_unknown_req_returns_none(): + s = _make_bare_scheduler() + assert s.get_and_clear_chunk_stats("unknown") is None + + +# --------------------------------------------------------------------------- +# Test 3A — scheduler integration (real engine, requires model) +# --------------------------------------------------------------------------- + + +@pytest.mark.chunked_prefill +@pytest.mark.cpu +@pytest.mark.parametrize("max_num_seqs", [2]) +@pytest.mark.parametrize("max_model_len", [256]) +@pytest.mark.parametrize("max_num_batched_tokens", [64]) +@pytest.mark.parametrize("available_blocks", [None]) +def test_scheduler_bench_metrics_accumulated( + model, + backend, + monkeypatch, + set_random_seed, + max_num_seqs, + max_model_len, + max_num_batched_tokens, + available_blocks, +): + """Two requests with prompts longer than max_num_batched_tokens each trigger + multiple prefill chunks. Verify that _chunk_latencies, _arrival_ts, and + _first_scheduled_ts are populated correctly, and cleared once the request + finishes via _free_request.""" + from llm_cache import get_cached_engine + from scheduling_utils import create_request_for_scheduler_test, random_prompt + + monkeypatch.setenv("SENDNN_INFERENCE_BENCH_METRICS_ENABLED", "1") + # Re-read envs so the scheduler sees the updated value + import sendnn_inference.envs as envs_spyre + + monkeypatch.setattr(envs_spyre, "SENDNN_INFERENCE_BENCH_METRICS_ENABLED", True) + + engine = get_cached_engine( + model=model, + max_model_len=max_model_len, + max_num_seqs=max_num_seqs, + available_blocks=available_blocks, + max_num_batched_tokens=max_num_batched_tokens, + backend=backend, + monkeypatch=monkeypatch, + ) + scheduler = engine.scheduler + + # Patch _chunk_latencies accumulation guard for this test run since the + # engine was potentially cached before BENCH_METRICS_ENABLED was set. + scheduler._chunk_latencies = {} + scheduler._arrival_ts = {} + scheduler._first_scheduled_ts = {} + + # Prompts longer than max_num_batched_tokens (64) → ≥ 2 prefill chunks each + prompt_len = max_num_batched_tokens + 20 # 84 tokens → 2 chunks of 64 + req1 = create_request_for_scheduler_test( + model=model, + request_id=0, + add_step=0, + max_tokens=4, + prompt=random_prompt(model=model, seed=0, length=prompt_len), + use_golden_token_injection=False, + generate_hf_results=False, + ) + req2 = create_request_for_scheduler_test( + model=model, + request_id=1, + add_step=0, + max_tokens=4, + prompt=random_prompt(model=model, seed=1, length=prompt_len), + use_golden_token_injection=False, + generate_hf_results=False, + ) + + # Capture metrics just before _free_request clears them + captured: dict[str, dict] = {} + original_free = scheduler.__class__._free_request + + def _capturing_free(self, request, delay_free_blocks=False): + req_id = request.request_id + captured[req_id] = { + "chunk_latencies": list(self._chunk_latencies.get(req_id, [])), + "has_arrival_ts": req_id in self._arrival_ts, + "has_first_scheduled_ts": req_id in self._first_scheduled_ts, + } + return original_free(self, request, delay_free_blocks) + + scheduler._free_request = _capturing_free.__get__(scheduler) + + # Add both requests and step until both finish. + # engine.step() returns (engine_core_output_dict, ...) — outputs are keyed + # by worker rank; use rank 0. Each output has .outputs with per-request + # RequestOutput objects that expose .new_token_ids and .request_id. + engine.add_request(req1.request) + engine.add_request(req2.request) + + tokens_per_req: dict[str, int] = {"0": 0, "1": 0} + max_tokens = 4 + for _ in range(200): # generous upper bound + step_output = engine.step() + engine_core_output = step_output[0].get(0) + if engine_core_output is not None: + for out in engine_core_output.outputs: + tokens_per_req[out.request_id] = tokens_per_req.get(out.request_id, 0) + len( + out.new_token_ids + ) + if all(v >= max_tokens for v in tokens_per_req.values()): + break + + assert all(v >= max_tokens for v in tokens_per_req.values()), ( + f"Requests did not finish after 200 steps: {tokens_per_req}" + ) + + # Both requests must have been captured by _capturing_free + for req_id in ("0", "1"): + assert req_id in captured, f"_free_request was never called for req {req_id}" + info = captured[req_id] + + # At least 2 chunks recorded (prompt > chunk_size) + assert len(info["chunk_latencies"]) >= 2, ( + f"req {req_id}: expected ≥2 chunk latencies, got {info['chunk_latencies']}" + ) + # All latencies must be positive floats + for lat in info["chunk_latencies"]: + assert isinstance(lat, float) and lat > 0, f"req {req_id}: non-positive latency {lat}" + assert info["has_arrival_ts"], f"req {req_id}: _arrival_ts not set" + assert info["has_first_scheduled_ts"], f"req {req_id}: _first_scheduled_ts not set" + + # The two requests must have independent latency lists (no cross-contamination) + assert captured["0"]["chunk_latencies"] != captured["1"]["chunk_latencies"] or ( + # Allow equal only if prompts produced identical timings by coincidence — + # check lengths are both ≥ 2 as a minimum + len(captured["0"]["chunk_latencies"]) >= 2 and len(captured["1"]["chunk_latencies"]) >= 2 + ) + + # After all requests finished, the dicts must be empty + assert scheduler._chunk_latencies == {}, "Leftover entries in _chunk_latencies after run" + assert scheduler._arrival_ts == {}, "Leftover entries in _arrival_ts after run" + assert scheduler._first_scheduled_ts == {}, "Leftover entries in _first_scheduled_ts after run" + + +# --------------------------------------------------------------------------- +# Test 4 — async_request_spyre_chat SSE parsing +# --------------------------------------------------------------------------- + + +def _build_sse_stream(chunks: list[str]) -> list[bytes]: + """Encode a list of SSE message strings as byte chunks.""" + return [c.encode() for c in chunks] + + +def _make_session_mock(status: int, sse_chunks: list[str]): + """Return a mock aiohttp ClientSession whose post() returns a fake SSE stream.""" + + async def _iter_any(): + for chunk in _build_sse_stream(sse_chunks): + yield chunk + + response_mock = MagicMock() + response_mock.status = status + response_mock.reason = None + response_mock.content.iter_any = _iter_any + + cm = MagicMock() + cm.__aenter__ = AsyncMock(return_value=response_mock) + cm.__aexit__ = AsyncMock(return_value=False) + + session = MagicMock() + session.post.return_value = cm + return session + + +def _make_request_input(): + from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput + + return RequestFuncInput( + prompt="hello", + api_url="http://localhost:8000/v1/chat/completions", + prompt_len=1, + output_len=4, + model="test-model", + ) + + +_SPYRE_METRICS = { + "queued_time_s": 55555.5, + "num_chunked_prefills": 42, + "chunk_prefill_latencies_s": [0.000007, 66666.6], +} + +_SSE_WITH_METRICS = [ + 'data: {"id":"1","choices":[{"delta":{"content":"hi"},"index":0}]}\n\n', + f'data: {{"id":"1","choices":[],"usage":{{"prompt_tokens":3,"completion_tokens":2}},' + f'"spyre_metrics":{json.dumps(_SPYRE_METRICS)}}}\n\n', + "data: [DONE]\n\n", +] + +_SSE_WITHOUT_METRICS = [ + 'data: {"id":"1","choices":[{"delta":{"content":"hi"},"index":0}]}\n\n', + 'data: {"id":"1","choices":[],"usage":{"prompt_tokens":3,"completion_tokens":2}}\n\n', + "data: [DONE]\n\n", +] + + +@pytest.mark.cpu +def test_spyre_metrics_parsed(): + from sendnn_inference.benchmarks.spyre_request_func import async_request_spyre_chat + + session = _make_session_mock(200, _SSE_WITH_METRICS) + output = asyncio.run(async_request_spyre_chat(_make_request_input(), session)) + + assert output.success is True + assert output.custom_metrics_dict == _SPYRE_METRICS + + +@pytest.mark.cpu +def test_spyre_metrics_absent(): + from sendnn_inference.benchmarks.spyre_request_func import async_request_spyre_chat + + session = _make_session_mock(200, _SSE_WITHOUT_METRICS) + output = asyncio.run(async_request_spyre_chat(_make_request_input(), session)) + + assert output.success is True + assert output.custom_metrics_dict == {} + + +@pytest.mark.cpu +def test_success_flag_set(): + from sendnn_inference.benchmarks.spyre_request_func import async_request_spyre_chat + + session = _make_session_mock(200, _SSE_WITH_METRICS) + output = asyncio.run(async_request_spyre_chat(_make_request_input(), session)) + assert output.success is True + + +@pytest.mark.cpu +def test_output_tokens_parsed(): + from sendnn_inference.benchmarks.spyre_request_func import async_request_spyre_chat + + session = _make_session_mock(200, _SSE_WITH_METRICS) + output = asyncio.run(async_request_spyre_chat(_make_request_input(), session)) + assert output.output_tokens == 2 From 3fc70541cadd18964eb55a436f5203d6e1b1878a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 10 Jun 2026 19:46:32 +0200 Subject: [PATCH 044/106] scheduler: move bench attributes to dedicated dataclass MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 40 ++++++++++++++++---------- tests/benchmarks/test_bench_metrics.py | 35 +++++++++++----------- 2 files changed, 44 insertions(+), 31 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index c24b52e6b..f69ede913 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -3,6 +3,7 @@ import math import time from collections import deque +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Iterable, Union from vllm.logger import init_logger @@ -21,6 +22,17 @@ logger = init_logger(__name__) + +@dataclass +class SpyreBenchState: + """Bench-metrics-only per-request state. Only instantiated when + SENDNN_INFERENCE_BENCH_METRICS_ENABLED is set.""" + + chunk_latencies: dict[str, list[float]] = field(default_factory=dict) + arrival_ts: dict[str, float] = field(default_factory=dict) + first_scheduled_ts: dict[str, float] = field(default_factory=dict) + + # Ensure that block_size is 64 # This ensures the rounding function is correct assert SpyrePlatform.get_block_size() == 64 @@ -204,14 +216,10 @@ def __init__(self, *args, **kwargs) -> None: self.block_size = SpyrePlatform.get_block_size() self.max_batch_tkv_limit = SpyrePlatform.get_max_batch_tkv_limit() - # Per-request chunk prefill latency accumulator. - # Only populated when SENDNN_INFERENCE_BENCH_METRICS_ENABLED is set. - self._chunk_latencies: dict[str, list[float]] = {} - # Timestamps for queue-wait-time calculation. - self._arrival_ts: dict[str, float] = {} - self._first_scheduled_ts: dict[str, float] = {} + self._bench: SpyreBenchState | None = None if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED: + self._bench = SpyreBenchState() from sendnn_inference.v1.metrics.stats_logger import register_scheduler register_scheduler(self) @@ -243,17 +251,17 @@ def update_from_output(self, scheduler_output, model_runner_output): ) # Accumulate per-chunk timings before removing completed prefills - if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED: + if self._bench is not None: for req in self.ongoing_prefills: t = model_runner_output.chunk_prefill_time_s.get(req.request_id) if t is not None: - self._chunk_latencies.setdefault(req.request_id, []).append(t) + self._bench.chunk_latencies.setdefault(req.request_id, []).append(t) # Track first-scheduled time and arrival time for queue-wait calculation now = time.time() for req in self.ongoing_prefills: - if req.request_id not in self._first_scheduled_ts: - self._first_scheduled_ts[req.request_id] = now - self._arrival_ts[req.request_id] = req.arrival_time + if req.request_id not in self._bench.first_scheduled_ts: + self._bench.first_scheduled_ts[req.request_id] = now + self._bench.arrival_ts[req.request_id] = req.arrival_time # Remove completed prefills self.ongoing_prefills = [ @@ -265,7 +273,9 @@ def update_from_output(self, scheduler_output, model_runner_output): def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: """Return and clear accumulated chunk timing for a finished request.""" - lats = self._chunk_latencies.pop(req_id, None) + if self._bench is None: + return None + lats = self._bench.chunk_latencies.pop(req_id, None) if lats is None: return None return { @@ -278,11 +288,11 @@ def _free_request(self, request, delay_free_blocks: bool = False): they travel over ZMQ to the API server process in EngineCoreOutput.""" kv_xfer_params = super()._free_request(request, delay_free_blocks) - if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED: + if self._bench is not None: req_id = request.request_id chunk_stats = self.get_and_clear_chunk_stats(req_id) - first_ts = self._first_scheduled_ts.pop(req_id, None) - arrival_ts = self._arrival_ts.pop(req_id, None) + first_ts = self._bench.first_scheduled_ts.pop(req_id, None) + arrival_ts = self._bench.arrival_ts.pop(req_id, None) queued_time_s = ( (first_ts - arrival_ts) if first_ts is not None and arrival_ts is not None else 0.0 ) diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index e4c8169c4..2ee45f2e6 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -193,24 +193,24 @@ def test_print_missing_keys_tolerated(capsys): def _make_bare_scheduler(): """Instantiate ChunkedPrefillSpyreScheduler with __init__ bypassed so we can call its methods without a full vllm engine setup.""" - from sendnn_inference.v1.core.scheduler import ChunkedPrefillSpyreScheduler + from sendnn_inference.v1.core.scheduler import ChunkedPrefillSpyreScheduler, SpyreBenchState with patch.object(ChunkedPrefillSpyreScheduler, "__init__", lambda *a, **kw: None): s = ChunkedPrefillSpyreScheduler() - s._chunk_latencies = {} + s._bench = SpyreBenchState() return s @pytest.mark.cpu def test_get_and_clear_returns_correct_dict(): s = _make_bare_scheduler() - s._chunk_latencies["r0"] = [88888.8, 0.000005] + s._bench.chunk_latencies["r0"] = [88888.8, 0.000005] result = s.get_and_clear_chunk_stats("r0") assert result is not None assert result["num_chunked_prefills"] == 2 assert result["chunk_prefill_latencies_s"] == pytest.approx([88888.8, 0.000005]) # Entry must be cleared after retrieval - assert "r0" not in s._chunk_latencies + assert "r0" not in s._bench.chunk_latencies @pytest.mark.cpu @@ -264,11 +264,10 @@ def test_scheduler_bench_metrics_accumulated( ) scheduler = engine.scheduler - # Patch _chunk_latencies accumulation guard for this test run since the - # engine was potentially cached before BENCH_METRICS_ENABLED was set. - scheduler._chunk_latencies = {} - scheduler._arrival_ts = {} - scheduler._first_scheduled_ts = {} + # Reset bench state in case the engine was cached before BENCH_METRICS_ENABLED was set. + from sendnn_inference.v1.core.scheduler import SpyreBenchState + + scheduler._bench = SpyreBenchState() # Prompts longer than max_num_batched_tokens (64) → ≥ 2 prefill chunks each prompt_len = max_num_batched_tokens + 20 # 84 tokens → 2 chunks of 64 @@ -297,10 +296,11 @@ def test_scheduler_bench_metrics_accumulated( def _capturing_free(self, request, delay_free_blocks=False): req_id = request.request_id + bench = self._bench captured[req_id] = { - "chunk_latencies": list(self._chunk_latencies.get(req_id, [])), - "has_arrival_ts": req_id in self._arrival_ts, - "has_first_scheduled_ts": req_id in self._first_scheduled_ts, + "chunk_latencies": list(bench.chunk_latencies.get(req_id, [])) if bench else [], + "has_arrival_ts": (req_id in bench.arrival_ts) if bench else False, + "has_first_scheduled_ts": (req_id in bench.first_scheduled_ts) if bench else False, } return original_free(self, request, delay_free_blocks) @@ -352,10 +352,13 @@ def _capturing_free(self, request, delay_free_blocks=False): len(captured["0"]["chunk_latencies"]) >= 2 and len(captured["1"]["chunk_latencies"]) >= 2 ) - # After all requests finished, the dicts must be empty - assert scheduler._chunk_latencies == {}, "Leftover entries in _chunk_latencies after run" - assert scheduler._arrival_ts == {}, "Leftover entries in _arrival_ts after run" - assert scheduler._first_scheduled_ts == {}, "Leftover entries in _first_scheduled_ts after run" + # After all requests finished, the bench state dicts must be empty + assert scheduler._bench is not None + assert scheduler._bench.chunk_latencies == {}, "Leftover entries in chunk_latencies after run" + assert scheduler._bench.arrival_ts == {}, "Leftover entries in arrival_ts after run" + assert scheduler._bench.first_scheduled_ts == {}, ( + "Leftover entries in first_scheduled_ts after run" + ) # --------------------------------------------------------------------------- From 8a1756504f802e90cce0c8deb6c328bd6db6da53 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 10 Jun 2026 21:03:55 +0200 Subject: [PATCH 045/106] custom detailed timeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_bench_serve.py | 80 ++++- sendnn_inference/benchmarks/spyre_plot.py | 301 ++++++++++++++++++ sendnn_inference/v1/core/scheduler.py | 34 +- sendnn_inference/v1/metrics/stats_logger.py | 4 + .../v1/worker/spyre_model_runner.py | 19 +- 5 files changed, 431 insertions(+), 7 deletions(-) create mode 100644 sendnn_inference/benchmarks/spyre_plot.py diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 087c49497..44c2a2ed4 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -31,13 +31,15 @@ _BACKEND_NAME = "spyre-chat" -# Shared accumulator — populated by the wrapper below during the benchmark run. +# Shared accumulators — populated by the wrapper below during the benchmark run. _spyre_metrics_collected: list[dict[str, Any]] = [] +_request_outputs_collected: list[dict[str, Any]] = [] def _make_collecting_func(): """Return a wrapper around async_request_spyre_chat that accumulates - custom_metrics_dict into _spyre_metrics_collected.""" + custom_metrics_dict into _spyre_metrics_collected and per-request vLLM + timing into _request_outputs_collected.""" async def _wrapper( request_func_input: RequestFuncInput, @@ -48,6 +50,17 @@ async def _wrapper( if output.success: if output.custom_metrics_dict: _spyre_metrics_collected.append(output.custom_metrics_dict) + _request_outputs_collected.append( + { + "start_time": output.start_time, + "ttft": output.ttft, + "itl": output.itl, + "latency": output.latency, + "prompt_len": request_func_input.prompt_len, + "output_tokens": output.output_tokens, + **output.custom_metrics_dict, + } + ) else: logger.warning( "Spyre metrics absent from response — is " @@ -85,6 +98,29 @@ def _build_parser() -> argparse.ArgumentParser: action.default = _BACKEND_NAME break + parser.add_argument( + "--detailed-timeline", + action="store_true", + default=False, + help=( + "Write a detailed per-request Gantt-chart timeline HTML alongside the " + "JSON result file (same name with a _detailed.html suffix). " + "Requires --save-result and SENDNN_INFERENCE_BENCH_METRICS_ENABLED on the server." + ), + ) + parser.add_argument( + "--itl-thresholds", + type=float, + nargs=2, + metavar=("LOW", "HIGH"), + default=None, + help=( + "Two ITL thresholds in seconds for decode coloring in the detailed timeline. " + "Decode steps below LOW are green, between LOW and HIGH are orange, " + "above HIGH are red. When omitted, all decode steps are green." + ), + ) + return parser @@ -268,6 +304,7 @@ def main() -> None: selected_percentiles = [float(p) for p in args.metric_percentiles.split(",")] _spyre_metrics_collected.clear() + _request_outputs_collected.clear() run_started_at = time.time() stdout_trailing, stderr_trailing = _run_vllm_and_capture_trailing(args) @@ -277,6 +314,45 @@ def main() -> None: _inject_spyre_metrics_into_result_file(args, _spyre_metrics_collected, run_started_at) + if getattr(args, "detailed_timeline", False): + from pathlib import Path + + from sendnn_inference.benchmarks.spyre_plot import generate_detailed_timeline_plot + + # Derive the HTML path from the JSON result file: same name, _detailed.html suffix. + result_dir = getattr(args, "result_dir", None) or "." + explicit_name = getattr(args, "result_filename", None) + if explicit_name: + json_candidate = ( + explicit_name + if os.path.isabs(explicit_name) + else os.path.join(result_dir, explicit_name) + ) + candidates = [json_candidate] if os.path.isfile(json_candidate) else [] + else: + try: + candidates = [ + os.path.join(result_dir, f) + for f in os.listdir(result_dir) + if f.endswith(".json") + and os.path.getmtime(os.path.join(result_dir, f)) >= run_started_at + ] + except OSError: + candidates = [] + + if candidates: + json_path = Path(max(candidates, key=os.path.getmtime)) + html_path = json_path.with_name(json_path.stem + "_detailed.html") + itl_thresholds = getattr(args, "itl_thresholds", None) + generate_detailed_timeline_plot( + _request_outputs_collected, html_path, itl_thresholds=itl_thresholds + ) + else: + logger.warning( + "--detailed-timeline requires --save-result so the JSON path is known; " + "no result file found, skipping timeline." + ) + trailing = stdout_trailing + stderr_trailing if trailing.strip(): print(trailing, end="") diff --git a/sendnn_inference/benchmarks/spyre_plot.py b/sendnn_inference/benchmarks/spyre_plot.py new file mode 100644 index 000000000..8a5b0a5f9 --- /dev/null +++ b/sendnn_inference/benchmarks/spyre_plot.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Detailed per-request Gantt-chart timeline plot for sendnn-bench serve. + +Generates an HTML file showing, for each request: + - Queue wait time (before first prefill) + - Each individual prefill chunk as a separate segment + - Waiting gaps between segments (absorbed if < 10% of segment duration) + - Each individual decode step (with TKV in hover) +""" + +import logging +from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) + +# Color scheme +_COLOR_QUEUE_WAIT = "#636EFA" # blue-purple — queue wait before first prefill +_COLOR_WAITING = "#777777" # dark grey — inter-segment gaps +_COLOR_PREFILL = "#FF0092" # pink — all prefill chunks +_COLOR_DECODE_FAST = "#109618" # green — decode ITL below lower threshold (or all decodes) +_COLOR_DECODE_MID = "#FF7F0E" # orange — decode ITL between thresholds (vLLM colors) +_COLOR_DECODE_SLOW = "#D62728" # red — decode ITL above upper threshold +_GAP_ABSORPTION_THRESHOLD = 0.10 # absorb gap if < 10% of current segment duration + + +def _tostr(sec: float) -> str: + """Convert elapsed seconds to HH:MM:SS.mmm (same format as vLLM's plot.py).""" + h = int(sec // 3600) + m = int((sec % 3600) // 60) + s = sec % 60 + return f"{h:02d}:{m:02d}:{s:06.3f}" + + +def _decode_type(lat: float, thresholds: list[float] | None) -> str: + if not thresholds: + return "Decode" + if lat < thresholds[0]: + return "Decode" + if lat < thresholds[1]: + return "Decode (mid)" + return "Decode (slow)" + + +def _build_detailed_segments( + request: dict[str, Any], + t0_global: float, + itl_thresholds: list[float] | None = None, +) -> list[dict[str, Any]]: + """Convert one request's timing data into ordered Gantt segments. + + All timestamps are elapsed seconds relative to t0_global (min client start_time + across all requests), formatted as HH:MM:SS.mmm strings for px.timeline. + + Absolute server-side timestamps (chunk_prefill_start_times_s, decode_start_times_s) + are only used to derive *gaps between consecutive segments on the same request* — + never mixed with client-side start_time values — to avoid cross-clock skew. + """ + client_start = (request.get("start_time") or 0.0) - t0_global + latency = request.get("latency") + prompt_len = request.get("prompt_len") + output_tokens = request.get("output_tokens") + req_finish = client_start + latency if latency is not None else None + + queued_time_s = request.get("queued_time_s") or 0.0 + prefill_lats = request.get("chunk_prefill_latencies_s") or [] + prefill_starts_abs = request.get("chunk_prefill_start_times_s") or [] + decode_lats = request.get("decode_latencies_s") or [] + decode_starts_abs = request.get("decode_start_times_s") or [] + decode_tkvs = request.get("decode_tkvs") or [] + + if not prefill_lats: + return [] + + segments: list[dict[str, Any]] = [] + req_label = request.get("_label", "Req ?") + + common = { + "request_id": req_label, + "prompt_tokens": prompt_len, + "output_tokens": output_tokens, + "req_start_time": _tostr(client_start), + "req_finish_time": _tostr(req_finish) if req_finish is not None else "—", + } + + # --- Queue wait --- + # Anchored to client start_time; queued_time_s is server-measured but relative. + first_prefill_t = client_start + queued_time_s + segments.append( + { + **common, + "start": _tostr(client_start), + "end": _tostr(first_prefill_t), + "type": "Queue wait", + "duration": f"{queued_time_s * 1000:.1f}ms", + "tkv": "—", + } + ) + + # --- Prefill chunks --- + # Use server-side absolute timestamps only to derive inter-chunk gaps + # (t_start[i+1] - t_start[i] - lat[i] = gap between chunk i and i+1). + # Cursor advances from first_prefill_t using latencies + derived gaps. + cursor = first_prefill_t + prev_end_cursor = cursor # tracks where previous segment ended + + for i, lat in enumerate(prefill_lats): + seg_start = cursor + + # Derive gap from server-side start-time diff when available + if i > 0 and len(prefill_starts_abs) > i: + gap = prefill_starts_abs[i] - prefill_starts_abs[i - 1] - prefill_lats[i - 1] + gap = max(gap, 0.0) + if gap > lat * _GAP_ABSORPTION_THRESHOLD: + # Insert explicit waiting segment + segments.append( + { + **common, + "start": _tostr(prev_end_cursor), + "end": _tostr(prev_end_cursor + gap), + "type": "Waiting", + "duration": f"{gap * 1000:.1f}ms", + "tkv": "—", + } + ) + seg_start = prev_end_cursor + gap + # else: absorb gap — seg_start stays at prev_end_cursor (no waiting bar) + cursor = seg_start + + seg_end = seg_start + lat + segments.append( + { + **common, + "start": _tostr(seg_start), + "end": _tostr(seg_end), + "type": "Prefill", + "duration": f"{lat * 1000:.1f}ms", + "tkv": "—", + } + ) + prev_end_cursor = seg_end + cursor = seg_end + + # --- Transition gap: last prefill → first decode --- + if decode_lats and prefill_starts_abs and decode_starts_abs: + gap = decode_starts_abs[0] - prefill_starts_abs[-1] - prefill_lats[-1] + gap = max(gap, 0.0) + if gap > decode_lats[0] * _GAP_ABSORPTION_THRESHOLD: + segments.append( + { + **common, + "start": _tostr(prev_end_cursor), + "end": _tostr(prev_end_cursor + gap), + "type": "Waiting", + "duration": f"{gap * 1000:.1f}ms", + "tkv": "—", + } + ) + cursor = prev_end_cursor + gap + # else absorb + + # --- Decode steps --- + for i, lat in enumerate(decode_lats): + seg_start = cursor + + if i > 0 and len(decode_starts_abs) > i: + gap = decode_starts_abs[i] - decode_starts_abs[i - 1] - decode_lats[i - 1] + gap = max(gap, 0.0) + if gap > lat * _GAP_ABSORPTION_THRESHOLD: + segments.append( + { + **common, + "start": _tostr(cursor), + "end": _tostr(cursor + gap), + "type": "Waiting", + "duration": f"{gap * 1000:.1f}ms", + "tkv": "—", + } + ) + seg_start = cursor + gap + + tkv = decode_tkvs[i] if i < len(decode_tkvs) else None + seg_end = seg_start + lat + segments.append( + { + **common, + "start": _tostr(seg_start), + "end": _tostr(seg_end), + "type": _decode_type(lat, itl_thresholds), + "duration": f"{lat * 1000:.1f}ms", + "tkv": str(tkv) if tkv is not None else "—", + } + ) + cursor = seg_end + + return segments + + +def generate_detailed_timeline_plot( + requests: list[dict[str, Any]], + output_path: Path, + itl_thresholds: list[float] | None = None, +) -> None: + """Build a per-request Gantt-chart HTML and write it to output_path. + + Args: + itl_thresholds: Two thresholds in seconds [low, high]. Decode steps below + low are green, between low and high are orange, above high are red. + When None (default), all decode steps are green. + """ + try: + import pandas as pd + import plotly.express as px + import plotly.io as pio + except ImportError as exc: + logger.warning( + "Cannot generate detailed timeline plot — missing dependency: %s. " + "Install with: pip install plotly pandas", + exc, + ) + return + + if not requests: + logger.warning("No request data to plot — skipping detailed timeline.") + return + + valid_starts = [r["start_time"] for r in requests if r.get("start_time") is not None] + if not valid_starts: + logger.warning("No start_time in collected requests — skipping detailed timeline.") + return + t0_global = min(valid_starts) + + sorted_requests = sorted(requests, key=lambda r: r.get("start_time") or 0.0) + for idx, req in enumerate(sorted_requests): + req["_label"] = f"Req {idx}" + + all_segments: list[dict[str, Any]] = [] + for req in sorted_requests: + all_segments.extend(_build_detailed_segments(req, t0_global, itl_thresholds)) + + if not all_segments: + logger.warning("No plottable segments found — skipping detailed timeline.") + return + + df = pd.DataFrame(all_segments) + + color_map = { + "Queue wait": _COLOR_QUEUE_WAIT, + "Waiting": _COLOR_WAITING, + "Prefill": _COLOR_PREFILL, + "Decode": _COLOR_DECODE_FAST, + "Decode (mid)": _COLOR_DECODE_MID, + "Decode (slow)": _COLOR_DECODE_SLOW, + } + category_order = ["Queue wait", "Prefill", "Waiting", "Decode", "Decode (mid)", "Decode (slow)"] + + fig = px.timeline( + df, + x_start="start", + x_end="end", + y="request_id", + color="type", + color_discrete_map=color_map, + category_orders={"type": category_order}, + hover_data=[ + "prompt_tokens", + "output_tokens", + "req_start_time", + "req_finish_time", + "duration", + "tkv", + ], + ) + + fig.update_traces( + hovertemplate=( + "%{y}
" + "Type: %{fullData.name}
" + "Duration: %{customdata[4]}
" + "TKV: %{customdata[5]}
" + "Prompt tokens: %{customdata[0]}
" + "Output tokens: %{customdata[1]}
" + "Req start: %{customdata[2]}
" + "Req end: %{customdata[3]}
" + "" + ) + ) + + fig.update_yaxes(autorange="reversed") + fig.update_layout( + xaxis_title="Time (HH:MM:SS from run start)", + yaxis_title="Request", + legend_title_text="Segment type", + ) + + try: + pio.write_html(fig, str(output_path)) + logger.info("Detailed timeline written to %s", output_path) + print(f"Detailed timeline plot written to: {output_path}") + except Exception as exc: + logger.warning("Failed to write detailed timeline HTML: %s", exc) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index f69ede913..137f6823e 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -31,6 +31,10 @@ class SpyreBenchState: chunk_latencies: dict[str, list[float]] = field(default_factory=dict) arrival_ts: dict[str, float] = field(default_factory=dict) first_scheduled_ts: dict[str, float] = field(default_factory=dict) + chunk_start_times: dict[str, list[float]] = field(default_factory=dict) + decode_latencies: dict[str, list[float]] = field(default_factory=dict) + decode_start_times: dict[str, list[float]] = field(default_factory=dict) + decode_tkvs: dict[str, list[int]] = field(default_factory=dict) # Ensure that block_size is 64 @@ -256,12 +260,22 @@ def update_from_output(self, scheduler_output, model_runner_output): t = model_runner_output.chunk_prefill_time_s.get(req.request_id) if t is not None: self._bench.chunk_latencies.setdefault(req.request_id, []).append(t) + t_start = model_runner_output.chunk_prefill_start_s.get(req.request_id) + if t_start is not None: + self._bench.chunk_start_times.setdefault(req.request_id, []).append(t_start) # Track first-scheduled time and arrival time for queue-wait calculation now = time.time() for req in self.ongoing_prefills: if req.request_id not in self._bench.first_scheduled_ts: self._bench.first_scheduled_ts[req.request_id] = now self._bench.arrival_ts[req.request_id] = req.arrival_time + # Accumulate decode timings + for req_id, t1 in model_runner_output.decode_time_s.items(): + self._bench.decode_latencies.setdefault(req_id, []).append(t1) + for req_id, t0 in model_runner_output.decode_start_s.items(): + self._bench.decode_start_times.setdefault(req_id, []).append(t0) + for req_id, tkv in model_runner_output.decode_tkv.items(): + self._bench.decode_tkvs.setdefault(req_id, []).append(tkv) # Remove completed prefills self.ongoing_prefills = [ @@ -276,11 +290,19 @@ def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: if self._bench is None: return None lats = self._bench.chunk_latencies.pop(req_id, None) - if lats is None: + starts = self._bench.chunk_start_times.pop(req_id, None) + dec_lats = self._bench.decode_latencies.pop(req_id, None) + dec_starts = self._bench.decode_start_times.pop(req_id, None) + dec_tkvs = self._bench.decode_tkvs.pop(req_id, None) + if lats is None and dec_lats is None: return None return { - "num_chunked_prefills": len(lats), - "chunk_prefill_latencies_s": lats, + "num_chunked_prefills": len(lats) if lats else 0, + "chunk_prefill_latencies_s": lats or [], + "chunk_prefill_start_times_s": starts or [], + "decode_latencies_s": dec_lats or [], + "decode_start_times_s": dec_starts or [], + "decode_tkvs": dec_tkvs or [], } def _free_request(self, request, delay_free_blocks: bool = False): @@ -302,6 +324,12 @@ def _free_request(self, request, delay_free_blocks: bool = False): "chunk_prefill_latencies_s": chunk_stats["chunk_prefill_latencies_s"] if chunk_stats else [], + "chunk_prefill_start_times_s": chunk_stats["chunk_prefill_start_times_s"] + if chunk_stats + else [], + "decode_latencies_s": chunk_stats["decode_latencies_s"] if chunk_stats else [], + "decode_start_times_s": chunk_stats["decode_start_times_s"] if chunk_stats else [], + "decode_tkvs": chunk_stats["decode_tkvs"] if chunk_stats else [], } if kv_xfer_params is None: kv_xfer_params = {"__spyre__": spyre_data} diff --git a/sendnn_inference/v1/metrics/stats_logger.py b/sendnn_inference/v1/metrics/stats_logger.py index 8b4ea97bb..9000a80b4 100644 --- a/sendnn_inference/v1/metrics/stats_logger.py +++ b/sendnn_inference/v1/metrics/stats_logger.py @@ -36,6 +36,10 @@ class SpyreRequestMetrics: queued_time_s: float num_chunked_prefills: int chunk_prefill_latencies_s: list[float] + chunk_prefill_start_times_s: list[float] = dataclasses.field(default_factory=list) + decode_latencies_s: list[float] = dataclasses.field(default_factory=list) + decode_start_times_s: list[float] = dataclasses.field(default_factory=list) + decode_tkvs: list[int] = dataclasses.field(default_factory=list) class SpyreMetricsRegistry: diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index 7e7721f92..a90519216 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -101,6 +101,12 @@ class SpyreModelRunnerOutput(ModelRunnerOutput): # Per-chunk prefill wall-clock time in seconds, keyed by request_id. # Only populated when SENDNN_INFERENCE_BENCH_METRICS_ENABLED is set. chunk_prefill_time_s: dict[str, float] = field(default_factory=dict) + # Absolute wall-clock start time (time.time()) of each prefill chunk. + chunk_prefill_start_s: dict[str, float] = field(default_factory=dict) + # Decode step wall-clock duration, start time, and tkv per request. + decode_time_s: dict[str, float] = field(default_factory=dict) + decode_start_s: dict[str, float] = field(default_factory=dict) + decode_tkv: dict[str, int] = field(default_factory=dict) InputBatchT = TypeVar("InputBatchT", bound=BaseInputBatch) @@ -1571,10 +1577,12 @@ def execute_model( t1 = time.time() - t0 logger.debug("t_forward_pass: %.2fms [prefill single chunk][batch size 1]", (t1 * 1000)) output = self.prefill_output() + assert isinstance(output, SpyreModelRunnerOutput) if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED: req_ids = list(scheduler_output.num_scheduled_tokens) if req_ids: output.chunk_prefill_time_s[req_ids[0]] = t1 + output.chunk_prefill_start_s[req_ids[0]] = t0 return output # Apply grammar bitmask for structured output requests. @@ -1619,10 +1627,17 @@ def execute_model( return self.get_empty_output() model_output = self.sampled_output(output, is_prefill) - if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED and is_prefill: + if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED: req_ids = list(scheduler_output.num_scheduled_tokens) if req_ids: - model_output.chunk_prefill_time_s[req_ids[0]] = t1 + if is_prefill: + model_output.chunk_prefill_time_s[req_ids[0]] = t1 + model_output.chunk_prefill_start_s[req_ids[0]] = t0 + else: + for req_id in req_ids: + model_output.decode_time_s[req_id] = t1 + model_output.decode_start_s[req_id] = t0 + model_output.decode_tkv[req_id] = model_output.tkv return model_output def prefill_output(self) -> SpyreModelRunnerOutput: From 71dbc31b6563765fdf7912079a53e5a57d2c7fcb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 10 Jun 2026 22:21:44 +0200 Subject: [PATCH 046/106] surface decode metrics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/benchmarks/spyre_bench_serve.py | 3 +++ tests/benchmarks/test_bench_metrics.py | 13 +++++++++++-- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 44c2a2ed4..7c25bf2fa 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -139,6 +139,7 @@ def _print_spyre_section( chunk_lats_ms = [ lat * 1000 for m in metrics_list for lat in m.get("chunk_prefill_latencies_s", []) ] + decode_lats_ms = [lat * 1000 for m in metrics_list for lat in m.get("decode_latencies_s", [])] total_prefill_chunks = sum(num_chunks_list) # Scalar summary line (mirrors vllm's plain-count header section) @@ -162,6 +163,7 @@ def _section(header: str, values: list[float], label: str) -> None: _section("Queue Wait Time", queue_times_ms, "Queue Wait Time (ms)") _section("Chunked Prefill Count", num_chunks_list, "Num Chunked Prefills") _section("Chunked Prefill Latency", chunk_lats_ms, "Chunk Prefill Latency (ms)") + _section("Decode Step Latency", decode_lats_ms, "Decode Step Latency (ms)") print("=" * 50) @@ -229,6 +231,7 @@ def _inject_spyre_metrics_into_result_file( m.get("chunk_prefill_latencies_s", []) for m in metrics_list ] result["spyre_total_prefill_chunks"] = sum(result["spyre_num_chunked_prefills"]) + result["spyre_decode_latencies_s"] = [m.get("decode_latencies_s", []) for m in metrics_list] try: with open(file_path, "w", encoding="utf-8") as fh: diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index 2ee45f2e6..379cfaa13 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -35,11 +35,13 @@ "queued_time_s": 42.0, "num_chunked_prefills": 7, "chunk_prefill_latencies_s": [0.001, 999.9, 0.003, 500.0, 0.002, 750.0, 1.0], + "decode_latencies_s": [88888.8, 0.000005, 44444.4], }, { "queued_time_s": 0.00001, "num_chunked_prefills": 3, "chunk_prefill_latencies_s": [12345.6, 0.0001, 99999.9], + "decode_latencies_s": [0.000002, 77777.7], }, ] @@ -75,6 +77,7 @@ def test_inject_adds_spyre_keys(tmp_path): assert "spyre_num_chunked_prefills" in data assert "spyre_chunk_prefill_latencies_s" in data assert "spyre_total_prefill_chunks" in data + assert "spyre_decode_latencies_s" in data @pytest.mark.cpu @@ -90,6 +93,10 @@ def test_inject_values_correct(tmp_path): [0.001, 999.9, 0.003, 500.0, 0.002, 750.0, 1.0], [12345.6, 0.0001, 99999.9], ] + assert data["spyre_decode_latencies_s"] == [ + [88888.8, 0.000005, 44444.4], + [0.000002, 77777.7], + ] # Original keys preserved assert data["backend"] == "spyre-chat" @@ -152,6 +159,7 @@ def test_print_section_headers(capsys): assert "Queue Wait Time" in out assert "Chunked Prefill Count" in out assert "Chunked Prefill Latency" in out + assert "Decode Step Latency" in out @pytest.mark.cpu @@ -172,8 +180,8 @@ def test_print_noop_when_empty(capsys): @pytest.mark.cpu def test_print_missing_keys_tolerated(capsys): - # Metrics without chunk_prefill_latencies_s — that section should be absent, - # but queue time and chunk count sections should still print. + # Metrics without chunk_prefill_latencies_s or decode_latencies_s — those + # sections should be absent, but queue time and chunk count should still print. metrics = [ {"queued_time_s": 77777.7, "num_chunked_prefills": 13}, {"queued_time_s": 0.000003, "num_chunked_prefills": 99}, @@ -183,6 +191,7 @@ def test_print_missing_keys_tolerated(capsys): assert "Queue Wait Time" in out assert "Chunked Prefill Count" in out assert "Chunked Prefill Latency" not in out + assert "Decode Step Latency" not in out # --------------------------------------------------------------------------- From 6e4394e37b4387be087ddebb6d8067ad02f0167b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 10 Jun 2026 22:48:33 +0200 Subject: [PATCH 047/106] timeline legend show itl range MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/benchmarks/spyre_plot.py | 44 ++++++++++++++++------- 1 file changed, 32 insertions(+), 12 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_plot.py b/sendnn_inference/benchmarks/spyre_plot.py index 8a5b0a5f9..da4a681ed 100644 --- a/sendnn_inference/benchmarks/spyre_plot.py +++ b/sendnn_inference/benchmarks/spyre_plot.py @@ -32,20 +32,35 @@ def _tostr(sec: float) -> str: return f"{h:02d}:{m:02d}:{s:06.3f}" -def _decode_type(lat: float, thresholds: list[float] | None) -> str: +def _decode_labels(thresholds: list[float] | None) -> tuple[str, str, str]: + """Return (fast_label, mid_label, slow_label) matching vLLM's timeline format.""" if not thresholds: - return "Decode" + return ("Decode", "", "") + lo_ms = int(thresholds[0] * 1000) + hi_ms = int(thresholds[1] * 1000) + return ( + f"ITL < {lo_ms}ms", + f"{lo_ms}ms ≤ ITL < {hi_ms}ms", + f"ITL ≥ {hi_ms}ms", + ) + + +def _decode_type(lat: float, thresholds: list[float] | None, labels: tuple[str, str, str]) -> str: + fast, mid, slow = labels + if not thresholds: + return fast if lat < thresholds[0]: - return "Decode" + return fast if lat < thresholds[1]: - return "Decode (mid)" - return "Decode (slow)" + return mid + return slow def _build_detailed_segments( request: dict[str, Any], t0_global: float, itl_thresholds: list[float] | None = None, + decode_labels: tuple[str, str, str] | None = None, ) -> list[dict[str, Any]]: """Convert one request's timing data into ordered Gantt segments. @@ -186,7 +201,7 @@ def _build_detailed_segments( **common, "start": _tostr(seg_start), "end": _tostr(seg_end), - "type": _decode_type(lat, itl_thresholds), + "type": _decode_type(lat, itl_thresholds, decode_labels or ("Decode", "", "")), "duration": f"{lat * 1000:.1f}ms", "tkv": str(tkv) if tkv is not None else "—", } @@ -234,9 +249,12 @@ def generate_detailed_timeline_plot( for idx, req in enumerate(sorted_requests): req["_label"] = f"Req {idx}" + labels = _decode_labels(itl_thresholds) + fast_lbl, mid_lbl, slow_lbl = labels + all_segments: list[dict[str, Any]] = [] for req in sorted_requests: - all_segments.extend(_build_detailed_segments(req, t0_global, itl_thresholds)) + all_segments.extend(_build_detailed_segments(req, t0_global, itl_thresholds, labels)) if not all_segments: logger.warning("No plottable segments found — skipping detailed timeline.") @@ -244,15 +262,17 @@ def generate_detailed_timeline_plot( df = pd.DataFrame(all_segments) - color_map = { + color_map: dict[str, str] = { "Queue wait": _COLOR_QUEUE_WAIT, "Waiting": _COLOR_WAITING, "Prefill": _COLOR_PREFILL, - "Decode": _COLOR_DECODE_FAST, - "Decode (mid)": _COLOR_DECODE_MID, - "Decode (slow)": _COLOR_DECODE_SLOW, + fast_lbl: _COLOR_DECODE_FAST, } - category_order = ["Queue wait", "Prefill", "Waiting", "Decode", "Decode (mid)", "Decode (slow)"] + category_order = ["Queue wait", "Prefill", "Waiting", fast_lbl] + if itl_thresholds: + color_map[mid_lbl] = _COLOR_DECODE_MID + color_map[slow_lbl] = _COLOR_DECODE_SLOW + category_order += [mid_lbl, slow_lbl] fig = px.timeline( df, From bcf882e6219e92518bc125642503c32ca766a8f3 Mon Sep 17 00:00:00 2001 From: Max de Bayser Date: Thu, 11 Jun 2026 10:41:13 -0400 Subject: [PATCH 048/106] Fix worker out-of-order bug and improve tests Signed-off-by: Max de Bayser --- .../v1/worker/spyre_model_runner.py | 4 +-- tests/e2e/test_logits_processors.py | 2 +- ...test_spyre_decode_pause_scheduler_steps.py | 27 ++++++++++++++++++- tests/multimodal/test_llava_next.py | 2 +- tests/multimodal/test_mistral3.py | 2 +- 5 files changed, 31 insertions(+), 6 deletions(-) diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index c3da8936a..07edb1a10 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -1450,7 +1450,7 @@ def _update_batch(self, scheduler_output: SchedulerOutput): # Find requests that are in input_batch but not in scheduler output (paused) paused_req_ids = current_batch_req_ids - scheduled_req_ids - for req_id in paused_req_ids: + for req_id in sorted(paused_req_ids): # Only pause if it's not a finished request (finished requests are handled separately) if req_id not in (scheduler_output.finished_req_ids or []): self.input_batch.pause_request(req_id) @@ -1461,7 +1461,7 @@ def _update_batch(self, scheduler_output: SchedulerOutput): # Find requests that are in scheduler output but not in input_batch # (restore from pausing) restored_req_ids = scheduled_req_ids - current_batch_req_ids - for req_id in restored_req_ids: + for req_id in sorted(restored_req_ids): # Only restore requests that were previously paused if req_id in self.paused_req_ids and req_id in self.requests: req_state = self.requests[req_id] diff --git a/tests/e2e/test_logits_processors.py b/tests/e2e/test_logits_processors.py index c5bfc8dfd..d16c3c2cb 100644 --- a/tests/e2e/test_logits_processors.py +++ b/tests/e2e/test_logits_processors.py @@ -127,7 +127,7 @@ def test_logits_processor_advanced( Uses SpyLogitsProcessor as the inner processor to verify token generation. """ - from tests.v1.worker.mock_model import InstrumentedModelRunner + from v1.worker.mock_model import InstrumentedModelRunner from vllm.v1.sample.logits_processor.state import LogitsProcessors monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") diff --git a/tests/e2e/test_spyre_decode_pause_scheduler_steps.py b/tests/e2e/test_spyre_decode_pause_scheduler_steps.py index 267ee72e5..01c64375d 100644 --- a/tests/e2e/test_spyre_decode_pause_scheduler_steps.py +++ b/tests/e2e/test_spyre_decode_pause_scheduler_steps.py @@ -17,7 +17,7 @@ create_request_for_scheduler_test, random_prompt, ) -from spyre_util import ModelInfo +from spyre_util import ModelInfo, verify_block_tables @pytest.mark.chunked_prefill @@ -100,6 +100,7 @@ def test_max_batch_tkv_decode_pausing( "running": ["0"], "request_outputs": ["0"], "n_used_blocks": 1, + "block_tables": {"0": [1]}, }, { # Decode sequence 0 @@ -109,6 +110,7 @@ def test_max_batch_tkv_decode_pausing( "running": ["0"], "request_outputs": ["0"], "n_used_blocks": 1, + "block_tables": {"0": [1]}, }, { # Prefill sequence 1 @@ -118,6 +120,7 @@ def test_max_batch_tkv_decode_pausing( "running": ["1", "0"], "request_outputs": ["1"], "n_used_blocks": 2, + "block_tables": {"0": [1], "1": [2]}, }, { # Decode sequences 0 and 1 @@ -127,6 +130,7 @@ def test_max_batch_tkv_decode_pausing( "running": ["1", "0"], "request_outputs": ["1", "0"], "n_used_blocks": 2, + "block_tables": {"0": [1], "1": [2]}, }, { # Prefill sequence 2 @@ -139,6 +143,7 @@ def test_max_batch_tkv_decode_pausing( "running": ["2", "1", "0"], "request_outputs": ["2"], "n_used_blocks": 4, + "block_tables": {"0": [1], "1": [2], "2": [3, 4]}, }, { # Decode sequences 0, 1, and 2 @@ -148,6 +153,7 @@ def test_max_batch_tkv_decode_pausing( "running": ["2", "1", "0"], "request_outputs": ["2", "1", "0"], "n_used_blocks": 4, + "block_tables": {"0": [1], "1": [2], "2": [3, 4]}, }, { # Decode sequences 0, 1, and 2 @@ -157,6 +163,7 @@ def test_max_batch_tkv_decode_pausing( "running": ["2", "1", "0"], "request_outputs": ["2", "1", "0"], "n_used_blocks": 4, + "block_tables": {"0": [1], "1": [2], "2": [3, 4]}, }, { # Decode sequences 0 and 1 @@ -168,6 +175,7 @@ def test_max_batch_tkv_decode_pausing( "running": ["1", "0"], "request_outputs": ["1", "0"], "n_used_blocks": 4, + "block_tables": {"0": [1], "1": [2], "2": [3, 4]}, }, { # Decode sequences 0 and 1 @@ -179,6 +187,7 @@ def test_max_batch_tkv_decode_pausing( "request_outputs": ["1", "0"], "finished_requests": ["0"], "n_used_blocks": 3, + "block_tables": {"1": [2], "2": [3, 4]}, }, { # Decode sequences 1 and 2 @@ -189,6 +198,7 @@ def test_max_batch_tkv_decode_pausing( "running": ["1", "2"], "request_outputs": ["1", "2"], "n_used_blocks": 3, + "block_tables": {"1": [2], "2": [3, 4]}, }, { # Decode sequences 1 and 2 @@ -200,6 +210,7 @@ def test_max_batch_tkv_decode_pausing( "request_outputs": ["1", "2"], "finished_requests": ["1"], "n_used_blocks": 2, + "block_tables": {"2": [3, 4]}, }, { # Decode sequence 2 @@ -209,6 +220,7 @@ def test_max_batch_tkv_decode_pausing( "running": ["2"], "request_outputs": ["2"], "n_used_blocks": 2, + "block_tables": {"2": [3, 4]}, }, { # Sequence 2 finishes @@ -219,6 +231,7 @@ def test_max_batch_tkv_decode_pausing( "request_outputs": ["2"], "finished_requests": ["2"], "n_used_blocks": 0, + "block_tables": {}, }, { # tkv should be cleared one step later @@ -243,6 +256,7 @@ def test_max_batch_tkv_decode_pausing( available_blocks=available_blocks, max_batch_tkv_limit=max_batch_tkv_limit, max_num_batched_tokens=max_num_batched_tokens, + extra_assert_funcs=[verify_block_tables], ) @@ -316,6 +330,7 @@ def test_prefill_exceeds_max_batch_tkv( "running": [], "request_outputs": [], "n_used_blocks": 0, + "block_tables": {}, }, { # Prefill sequence 0 @@ -325,6 +340,7 @@ def test_prefill_exceeds_max_batch_tkv( "running": ["0"], "request_outputs": ["0"], "n_used_blocks": 1, + "block_tables": {"0": [1]}, }, { # Decode sequence 0 @@ -334,6 +350,7 @@ def test_prefill_exceeds_max_batch_tkv( "running": ["0"], "request_outputs": ["0"], "n_used_blocks": 1, + "block_tables": {"0": [1]}, }, { # Prefill sequence 1 @@ -343,6 +360,7 @@ def test_prefill_exceeds_max_batch_tkv( "running": ["1", "0"], "request_outputs": ["1"], "n_used_blocks": 2, + "block_tables": {"0": [1], "1": [2]}, }, { # Decode sequences 0 and 1 @@ -352,6 +370,7 @@ def test_prefill_exceeds_max_batch_tkv( "running": ["1", "0"], "request_outputs": ["1", "0"], "n_used_blocks": 2, + "block_tables": {"0": [1], "1": [2]}, }, { # Decode sequences 0 and 1 @@ -363,6 +382,7 @@ def test_prefill_exceeds_max_batch_tkv( "running": ["1", "0"], "request_outputs": ["1", "0"], "n_used_blocks": 2, + "block_tables": {"0": [1], "1": [2]}, }, { # Sequences 0 and 1 both finish @@ -373,6 +393,7 @@ def test_prefill_exceeds_max_batch_tkv( "request_outputs": ["1", "0"], "finished_requests": ["1", "0"], "n_used_blocks": 0, + "block_tables": {}, }, { # Prefill sequence 2 @@ -382,6 +403,7 @@ def test_prefill_exceeds_max_batch_tkv( "running": ["2"], "request_outputs": ["2"], "n_used_blocks": 2, + "block_tables": {"2": [3, 4]}, }, { # Decode sequence 2 @@ -391,6 +413,7 @@ def test_prefill_exceeds_max_batch_tkv( "running": ["2"], "request_outputs": ["2"], "n_used_blocks": 2, + "block_tables": {"2": [3, 4]}, }, { # Sequence 2 finishes @@ -401,6 +424,7 @@ def test_prefill_exceeds_max_batch_tkv( "request_outputs": ["2"], "finished_requests": ["2"], "n_used_blocks": 0, + "block_tables": {}, }, { # tkv should be cleared one step later @@ -425,4 +449,5 @@ def test_prefill_exceeds_max_batch_tkv( available_blocks=available_blocks, max_batch_tkv_limit=max_batch_tkv_limit, max_num_batched_tokens=max_num_batched_tokens, + extra_assert_funcs=[verify_block_tables], ) diff --git a/tests/multimodal/test_llava_next.py b/tests/multimodal/test_llava_next.py index fb71e60f7..e93e7badb 100644 --- a/tests/multimodal/test_llava_next.py +++ b/tests/multimodal/test_llava_next.py @@ -15,7 +15,7 @@ from vllm.multimodal.inputs import MultiModalFeatureSpec import sendnn_inference.multimodal as spyre_mm -from tests.spyre_util import REFERENCE_MODELS +from spyre_util import REFERENCE_MODELS GVISION_MODEL = REFERENCE_MODELS["ibm-granite/granite-vision-3.2-2b"] # Marks all tests in this file as multimodal and CPU to match diff --git a/tests/multimodal/test_mistral3.py b/tests/multimodal/test_mistral3.py index 887498d08..2c0503970 100644 --- a/tests/multimodal/test_mistral3.py +++ b/tests/multimodal/test_mistral3.py @@ -15,7 +15,7 @@ from vllm.multimodal.inputs import MultiModalFeatureSpec import sendnn_inference.multimodal as spyre_mm -from tests.spyre_util import REFERENCE_MODELS +from spyre_util import REFERENCE_MODELS # MISTRAL3_MODEL = REFERENCE_MODELS["mistralai/Mistral-Small-3.2-24B-Instruct-2506"] MISTRAL3_MODEL = REFERENCE_MODELS["mistralai/Mistral-Small-3.1-24B-Instruct-2503"] From 0845da88e0b352090b1ca76ecbca52aec92a0588 Mon Sep 17 00:00:00 2001 From: Max de Bayser Date: Thu, 11 Jun 2026 11:54:05 -0400 Subject: [PATCH 049/106] update hf_cache.json Signed-off-by: Max de Bayser --- tests/hf_cache.json | 94 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 93 insertions(+), 1 deletion(-) diff --git a/tests/hf_cache.json b/tests/hf_cache.json index 71959fa49..b678dfd73 100644 --- a/tests/hf_cache.json +++ b/tests/hf_cache.json @@ -272,6 +272,52 @@ "tokens": [ "n", " " ], "logprobs": [ -2.626779556274414, -4.909548759460449 ] } + }, + "__tokens__41511_37260_20675_12728_25134_19906_38530_14911_23429_28678_44642_24810_13855_37154_30398_12315_44722_48312_39829_44349_15247_35878_44186_33624_23210": { + "7": { + "text": ": 10000", + "token_ids": [ 44, 225, 35, 34, 34, 34, 34 ], + "tokens": [ ":", " ", "1", "0", "0", "0", "0" ], + "logprobs": [ -1.9627763032913208, -2.7621099948883057, -1.441766381263733, -2.154932975769043, -1.3412963151931763, -0.8329601287841797, -0.7286175489425659 ] + } + }, + "__tokens__6606_41659_37546_12539_24355_22097_32032_38773_4614_1394_41085_21274_37473_104_21895_35470_11246_46468_44313_1504_1251_26615_46167_18740_10648": { + "6": { + "text": "niques, 201", + "token_ids": [ 20455, 30, 225, 36, 34, 35 ], + "tokens": [ "niques", ",", " ", "2", "0", "1" ], + "logprobs": [ -1.7665514945983887, -3.048448085784912, -3.3999183177948, -1.7820366621017456, -0.6218030452728271, -0.6172934174537659 ] + } + }, + "__tokens__46997_46594_2780_4173_41072_36179_32923_15148_29788_29830_28571_7786_21171_19346_35542_48904_46671_26751_21869_13187_1766_1350_22854_15656_18681_43839_25845_27554_11608_1173_15984_6720_25082_49094_33157_8940_43927_39168_36102_44567_37502_38823_17392_48223_47286_7924_37066_35156_22682_26072_24089_45463_24621_40877_17399_43400_44228_22663_27908_45242_35580_23921_10904_15960_34390_8164": { + "3": { + "text": "etterThan", + "token_ids": [ 364, 391, 11546 ], + "tokens": [ "et", "ter", "Than" ], + "logprobs": [ -3.224111795425415, -1.6213057041168213, -3.286860704421997 ] + }, + "10": { + "text": "etterThanThanThanThanThanThanThanThan", + "token_ids": [ 364, 391, 11546, 11546, 11546, 11546, 11546, 11546, 11546, 11546 ], + "tokens": [ "et", "ter", "Than", "Than", "Than", "Than", "Than", "Than", "Than", "Than" ], + "logprobs": [ -3.224111795425415, -1.6213057041168213, -3.286860704421997, -3.065337896347046, -0.04596678167581558, -0.02546372078359127, -0.016097459942102432, -0.012399725615978241, -0.010239332914352417, -0.009012495167553425 ] + } + }, + "__tokens__41511_37260_20675_12728_25134_19906_38530_14911_23429_28678_44642_24810_13855_37154_30398": { + "11": { + "text": "enance SpanTransactionalvelvelvelvelvelvelvelvel", + "token_ids": [ 12988, 14911, 23429, 1203, 1203, 1203, 1203, 1203, 1203, 1203, 1203 ], + "tokens": [ "enance", " Span", "Transactional", "vel", "vel", "vel", "vel", "vel", "vel", "vel", "vel" ], + "logprobs": [ -0.7761213779449463, -2.555943489074707, -1.8137869834899902, -5.274265289306641, -2.291093111038208, -0.07470891624689102, -0.026308227330446243, -0.015697719529271126, -0.007885280065238476, -0.005254506133496761, -0.0035511308815330267 ] + } + }, + "__tokens__6606_41659_37546_12539_24355_22097_32032_38773_4614_1394_41085_21274_37473_104_21895": { + "13": { + "text": ",100000000000", + "token_ids": [ 30, 35, 34, 34, 34, 34, 34, 34, 34, 34, 34, 34, 34 ], + "tokens": [ ",", "1", "0", "0", "0", "0", "0", "0", "0", "0", "0", "0", "0" ], + "logprobs": [ -3.362755537033081, -3.4085299968719482, -1.7558674812316895, -0.9610953330993652, -1.2473288774490356, -0.8631182312965393, -0.88742995262146, -0.7514090538024902, -0.8336255550384521, -0.70525062084198, -0.7933809161186218, -0.8284340500831604, -0.6348282098770142 ] + } } } }, @@ -548,7 +594,53 @@ "tokens": [ " t", "oler" ], "logprobs": [ -4.828008651733398, -3.2714998722076416 ] } + }, + "__tokens__41511_37260_20675_12728_25134_19906_38530_14911_23429_28678_44642_24810_13855_37154_30398_12315_44722_48312_39829_44349_15247_35878_44186_33624_23210": { + "7": { + "text": " Latin-1 Latin-", + "token_ids": [ 19190, 266, 31, 35, 19190, 266, 31 ], + "tokens": [ " Lat", "in", "-", "1", " Lat", "in", "-" ], + "logprobs": [ -3.864952802658081, -1.5731655359268188, -2.660896062850952, -1.8971174955368042, -2.8469624519348145, -0.007352791260927916, -0.3517385423183441 ] + } + }, + "__tokens__6606_41659_37546_12539_24355_22097_32032_38773_4614_1394_41085_21274_37473_104_21895_35470_11246_46468_44313_1504_1251_26615_46167_18740_10648": { + "6": { + "text": "nique(s)\n\n", + "token_ids": [ 37881, 26, 101, 27, 203, 203 ], + "tokens": [ "nique", "(", "s", ")", "\n", "\n" ], + "logprobs": [ -0.4146469235420227, -3.3161118030548096, -3.53086519241333, -0.7896995544433594, -2.6630074977874756, -0.6078359484672546 ] + } + }, + "__tokens__46997_46594_2780_4173_41072_36179_32923_15148_29788_29830_28571_7786_21171_19346_35542_48904_46671_26751_21869_13187_1766_1350_22854_15656_18681_43839_25845_27554_11608_1173_15984_6720_25082_49094_33157_8940_43927_39168_36102_44567_37502_38823_17392_48223_47286_7924_37066_35156_22682_26072_24089_45463_24621_40877_17399_43400_44228_22663_27908_45242_35580_23921_10904_15960_34390_8164": { + "3": { + "text": "ol - ", + "token_ids": [ 362, 429, 225 ], + "tokens": [ "ol", " -", " " ], + "logprobs": [ -4.403185844421387, -4.8154520988464355, -3.007040023803711 ] + }, + "10": { + "text": "ol - 1.0.0.0", + "token_ids": [ 362, 429, 225, 35, 32, 34, 32, 34, 32, 34 ], + "tokens": [ "ol", " -", " ", "1", ".", "0", ".", "0", ".", "0" ], + "logprobs": [ -4.403185844421387, -4.8154520988464355, -3.007040023803711, -0.6841981410980225, -1.199366569519043, -1.2438615560531616, -1.0571844577789307, -0.7776751518249512, -1.0070819854736328, -0.5469136238098145 ] + } + }, + "__tokens__41511_37260_20675_12728_25134_19906_38530_14911_23429_28678_44642_24810_13855_37154_30398": { + "11": { + "text": "enance.\n\n# 1. 2.", + "token_ids": [ 12988, 32, 203, 203, 21, 225, 35, 32, 225, 36, 32 ], + "tokens": [ "enance", ".", "\n", "\n", "#", " ", "1", ".", " ", "2", "." ], + "logprobs": [ -0.7462671399116516, -2.5164425373077393, -3.604614734649658, -0.8915535807609558, -2.2450551986694336, -3.185319662094116, -1.5186692476272583, -0.649973452091217, -2.2771198749542236, -0.923798680305481, -0.21222370862960815 ] + } + }, + "__tokens__6606_41659_37546_12539_24355_22097_32032_38773_4614_1394_41085_21274_37473_104_21895": { + "13": { + "text": "'';\n\n//\n// //\n// //", + "token_ids": [ 25, 920, 203, 203, 306, 203, 306, 225, 434, 203, 306, 225, 434 ], + "tokens": [ "'", "';", "\n", "\n", "//", "\n", "//", " ", " //", "\n", "//", " ", " //" ], + "logprobs": [ -3.038684129714966, -2.2691264152526855, -1.329238772392273, -0.6360853910446167, -2.914003372192383, -2.864417552947998, -0.6330344080924988, -1.2776681184768677, -2.7238693237304688, -2.050839424133301, -0.7158605456352234, -0.7345200181007385, -0.7729841470718384 ] + } } } } -} +} \ No newline at end of file From b54fd0c2754cd191b39f6e918892cce8e4a0561c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 12 Jun 2026 12:49:03 +0200 Subject: [PATCH 050/106] modelrunner: use dedicated dataclass for bench data MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 12 ++--- .../v1/worker/spyre_model_runner.py | 50 ++++++++++++------- 2 files changed, 38 insertions(+), 24 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 137f6823e..2d98110d9 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -255,12 +255,12 @@ def update_from_output(self, scheduler_output, model_runner_output): ) # Accumulate per-chunk timings before removing completed prefills - if self._bench is not None: + if self._bench is not None and model_runner_output._bench is not None: for req in self.ongoing_prefills: - t = model_runner_output.chunk_prefill_time_s.get(req.request_id) + t = model_runner_output._bench.chunk_prefill_time_s.get(req.request_id) if t is not None: self._bench.chunk_latencies.setdefault(req.request_id, []).append(t) - t_start = model_runner_output.chunk_prefill_start_s.get(req.request_id) + t_start = model_runner_output._bench.chunk_prefill_start_s.get(req.request_id) if t_start is not None: self._bench.chunk_start_times.setdefault(req.request_id, []).append(t_start) # Track first-scheduled time and arrival time for queue-wait calculation @@ -270,11 +270,11 @@ def update_from_output(self, scheduler_output, model_runner_output): self._bench.first_scheduled_ts[req.request_id] = now self._bench.arrival_ts[req.request_id] = req.arrival_time # Accumulate decode timings - for req_id, t1 in model_runner_output.decode_time_s.items(): + for req_id, t1 in model_runner_output._bench.decode_time_s.items(): self._bench.decode_latencies.setdefault(req_id, []).append(t1) - for req_id, t0 in model_runner_output.decode_start_s.items(): + for req_id, t0 in model_runner_output._bench.decode_start_s.items(): self._bench.decode_start_times.setdefault(req_id, []).append(t0) - for req_id, tkv in model_runner_output.decode_tkv.items(): + for req_id, tkv in model_runner_output._bench.decode_tkv.items(): self._bench.decode_tkvs.setdefault(req_id, []).append(tkv) # Remove completed prefills diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index a90519216..abc7f9344 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -88,6 +88,18 @@ class SamplingForwardInputs(ModelForwardInputs): slot_mapping: torch.Tensor +@dataclass +class SpyreModelRunnerBenchState: + """Bench-metrics-only state in SpyreModelRunnerOutput. + Only instantiated when SENDNN_INFERENCE_BENCH_METRICS_ENABLED is set.""" + + chunk_prefill_time_s: dict[str, float] = field(default_factory=dict) + chunk_prefill_start_s: dict[str, float] = field(default_factory=dict) + decode_time_s: dict[str, float] = field(default_factory=dict) + decode_start_s: dict[str, float] = field(default_factory=dict) + decode_tkv: dict[str, int] = field(default_factory=dict) + + @dataclass class SpyreModelRunnerOutput(ModelRunnerOutput): # Current tkv: this is the maximum padded request length in the batch) @@ -98,15 +110,8 @@ class SpyreModelRunnerOutput(ModelRunnerOutput): # available than the number of scheduled tokens. In that case, the scheduler # needs to update its state to reflect the correct number of computed tokens prefix_cache_hit_len: dict[str, int] = field(default_factory=dict) - # Per-chunk prefill wall-clock time in seconds, keyed by request_id. - # Only populated when SENDNN_INFERENCE_BENCH_METRICS_ENABLED is set. - chunk_prefill_time_s: dict[str, float] = field(default_factory=dict) - # Absolute wall-clock start time (time.time()) of each prefill chunk. - chunk_prefill_start_s: dict[str, float] = field(default_factory=dict) - # Decode step wall-clock duration, start time, and tkv per request. - decode_time_s: dict[str, float] = field(default_factory=dict) - decode_start_s: dict[str, float] = field(default_factory=dict) - decode_tkv: dict[str, int] = field(default_factory=dict) + # Bench-metrics-only state. None when SENDNN_INFERENCE_BENCH_METRICS_ENABLED is off. + _bench: SpyreModelRunnerBenchState | None = None InputBatchT = TypeVar("InputBatchT", bound=BaseInputBatch) @@ -1530,6 +1535,13 @@ def execute_model( ) -> ModelRunnerOutput: t0 = time.time() + # Create bench state dict only if bench metrics are enabled + bench_state = ( + SpyreModelRunnerBenchState() + if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED + else None + ) + self.update_states(scheduler_output) if not scheduler_output.total_num_scheduled_tokens: @@ -1578,11 +1590,12 @@ def execute_model( logger.debug("t_forward_pass: %.2fms [prefill single chunk][batch size 1]", (t1 * 1000)) output = self.prefill_output() assert isinstance(output, SpyreModelRunnerOutput) - if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED: + output._bench = bench_state + if bench_state is not None: req_ids = list(scheduler_output.num_scheduled_tokens) if req_ids: - output.chunk_prefill_time_s[req_ids[0]] = t1 - output.chunk_prefill_start_s[req_ids[0]] = t0 + bench_state.chunk_prefill_time_s[req_ids[0]] = t1 + bench_state.chunk_prefill_start_s[req_ids[0]] = t0 return output # Apply grammar bitmask for structured output requests. @@ -1627,17 +1640,18 @@ def execute_model( return self.get_empty_output() model_output = self.sampled_output(output, is_prefill) - if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED: + model_output._bench = bench_state + if bench_state is not None: req_ids = list(scheduler_output.num_scheduled_tokens) if req_ids: if is_prefill: - model_output.chunk_prefill_time_s[req_ids[0]] = t1 - model_output.chunk_prefill_start_s[req_ids[0]] = t0 + bench_state.chunk_prefill_time_s[req_ids[0]] = t1 + bench_state.chunk_prefill_start_s[req_ids[0]] = t0 else: for req_id in req_ids: - model_output.decode_time_s[req_id] = t1 - model_output.decode_start_s[req_id] = t0 - model_output.decode_tkv[req_id] = model_output.tkv + bench_state.decode_time_s[req_id] = t1 + bench_state.decode_start_s[req_id] = t0 + bench_state.decode_tkv[req_id] = model_output.tkv return model_output def prefill_output(self) -> SpyreModelRunnerOutput: From ffeac1f1588efba169f8de803b4e7787c333a966 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 12 Jun 2026 12:52:43 +0200 Subject: [PATCH 051/106] remove decode_tkv MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/benchmarks/spyre_plot.py | 10 +++++++--- sendnn_inference/v1/core/scheduler.py | 13 +++++++------ sendnn_inference/v1/metrics/stats_logger.py | 2 +- sendnn_inference/v1/worker/spyre_model_runner.py | 5 +++-- 4 files changed, 18 insertions(+), 12 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_plot.py b/sendnn_inference/benchmarks/spyre_plot.py index da4a681ed..e7e743d9d 100644 --- a/sendnn_inference/benchmarks/spyre_plot.py +++ b/sendnn_inference/benchmarks/spyre_plot.py @@ -82,7 +82,7 @@ def _build_detailed_segments( prefill_starts_abs = request.get("chunk_prefill_start_times_s") or [] decode_lats = request.get("decode_latencies_s") or [] decode_starts_abs = request.get("decode_start_times_s") or [] - decode_tkvs = request.get("decode_tkvs") or [] + tkvs = request.get("tkvs") or [] if not prefill_lats: return [] @@ -118,6 +118,7 @@ def _build_detailed_segments( # Cursor advances from first_prefill_t using latencies + derived gaps. cursor = first_prefill_t prev_end_cursor = cursor # tracks where previous segment ended + tkv_idx = 0 # index into all_tkvs (one per prefill, then one per decode) for i, lat in enumerate(prefill_lats): seg_start = cursor @@ -142,6 +143,8 @@ def _build_detailed_segments( # else: absorb gap — seg_start stays at prev_end_cursor (no waiting bar) cursor = seg_start + tkv = tkvs[tkv_idx] if tkv_idx < len(tkvs) else None + tkv_idx += 1 seg_end = seg_start + lat segments.append( { @@ -150,7 +153,7 @@ def _build_detailed_segments( "end": _tostr(seg_end), "type": "Prefill", "duration": f"{lat * 1000:.1f}ms", - "tkv": "—", + "tkv": str(tkv) if tkv is not None else "—", } ) prev_end_cursor = seg_end @@ -194,7 +197,8 @@ def _build_detailed_segments( ) seg_start = cursor + gap - tkv = decode_tkvs[i] if i < len(decode_tkvs) else None + tkv = tkvs[tkv_idx] if tkv_idx < len(tkvs) else None + tkv_idx += 1 seg_end = seg_start + lat segments.append( { diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 2d98110d9..91c53478e 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -34,7 +34,7 @@ class SpyreBenchState: chunk_start_times: dict[str, list[float]] = field(default_factory=dict) decode_latencies: dict[str, list[float]] = field(default_factory=dict) decode_start_times: dict[str, list[float]] = field(default_factory=dict) - decode_tkvs: dict[str, list[int]] = field(default_factory=dict) + tkvs: dict[str, list[int]] = field(default_factory=dict) # Ensure that block_size is 64 @@ -274,8 +274,9 @@ def update_from_output(self, scheduler_output, model_runner_output): self._bench.decode_latencies.setdefault(req_id, []).append(t1) for req_id, t0 in model_runner_output._bench.decode_start_s.items(): self._bench.decode_start_times.setdefault(req_id, []).append(t0) - for req_id, tkv in model_runner_output._bench.decode_tkv.items(): - self._bench.decode_tkvs.setdefault(req_id, []).append(tkv) + # Accumulate tkvs (one per prefill chunk and decode step) + for req_id, tkv_list in model_runner_output._bench.tkvs.items(): + self._bench.tkvs.setdefault(req_id, []).extend(tkv_list) # Remove completed prefills self.ongoing_prefills = [ @@ -293,7 +294,7 @@ def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: starts = self._bench.chunk_start_times.pop(req_id, None) dec_lats = self._bench.decode_latencies.pop(req_id, None) dec_starts = self._bench.decode_start_times.pop(req_id, None) - dec_tkvs = self._bench.decode_tkvs.pop(req_id, None) + tkvs = self._bench.tkvs.pop(req_id, None) if lats is None and dec_lats is None: return None return { @@ -302,7 +303,7 @@ def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: "chunk_prefill_start_times_s": starts or [], "decode_latencies_s": dec_lats or [], "decode_start_times_s": dec_starts or [], - "decode_tkvs": dec_tkvs or [], + "tkvs": tkvs or [], } def _free_request(self, request, delay_free_blocks: bool = False): @@ -329,7 +330,7 @@ def _free_request(self, request, delay_free_blocks: bool = False): else [], "decode_latencies_s": chunk_stats["decode_latencies_s"] if chunk_stats else [], "decode_start_times_s": chunk_stats["decode_start_times_s"] if chunk_stats else [], - "decode_tkvs": chunk_stats["decode_tkvs"] if chunk_stats else [], + "tkvs": chunk_stats["tkvs"] if chunk_stats else [], } if kv_xfer_params is None: kv_xfer_params = {"__spyre__": spyre_data} diff --git a/sendnn_inference/v1/metrics/stats_logger.py b/sendnn_inference/v1/metrics/stats_logger.py index 9000a80b4..e1c204970 100644 --- a/sendnn_inference/v1/metrics/stats_logger.py +++ b/sendnn_inference/v1/metrics/stats_logger.py @@ -39,7 +39,7 @@ class SpyreRequestMetrics: chunk_prefill_start_times_s: list[float] = dataclasses.field(default_factory=list) decode_latencies_s: list[float] = dataclasses.field(default_factory=list) decode_start_times_s: list[float] = dataclasses.field(default_factory=list) - decode_tkvs: list[int] = dataclasses.field(default_factory=list) + tkvs: list[int] = dataclasses.field(default_factory=list) class SpyreMetricsRegistry: diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index abc7f9344..9bf58b4ef 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -97,7 +97,7 @@ class SpyreModelRunnerBenchState: chunk_prefill_start_s: dict[str, float] = field(default_factory=dict) decode_time_s: dict[str, float] = field(default_factory=dict) decode_start_s: dict[str, float] = field(default_factory=dict) - decode_tkv: dict[str, int] = field(default_factory=dict) + tkvs: dict[str, list[int]] = field(default_factory=dict) @dataclass @@ -1647,11 +1647,12 @@ def execute_model( if is_prefill: bench_state.chunk_prefill_time_s[req_ids[0]] = t1 bench_state.chunk_prefill_start_s[req_ids[0]] = t0 + bench_state.tkvs.setdefault(req_ids[0], []).append(model_output.tkv) else: for req_id in req_ids: bench_state.decode_time_s[req_id] = t1 bench_state.decode_start_s[req_id] = t0 - bench_state.decode_tkv[req_id] = model_output.tkv + bench_state.tkvs.setdefault(req_id, []).append(model_output.tkv) return model_output def prefill_output(self) -> SpyreModelRunnerOutput: From dcca355a84b6bdbe1d3f17f98334a49e95e1851b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 12 Jun 2026 13:48:48 +0200 Subject: [PATCH 052/106] update tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_bench_serve.py | 5 +++++ tests/benchmarks/test_bench_metrics.py | 21 +++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 7c25bf2fa..cf27d990c 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -230,8 +230,13 @@ def _inject_spyre_metrics_into_result_file( result["spyre_chunk_prefill_latencies_s"] = [ m.get("chunk_prefill_latencies_s", []) for m in metrics_list ] + result["spyre_chunk_prefill_start_times_s"] = [ + m.get("chunk_prefill_start_times_s", []) for m in metrics_list + ] result["spyre_total_prefill_chunks"] = sum(result["spyre_num_chunked_prefills"]) result["spyre_decode_latencies_s"] = [m.get("decode_latencies_s", []) for m in metrics_list] + result["spyre_decode_start_times_s"] = [m.get("decode_start_times_s", []) for m in metrics_list] + result["spyre_tkvs"] = [m.get("tkvs", []) for m in metrics_list] try: with open(file_path, "w", encoding="utf-8") as fh: diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index 379cfaa13..0e2c3d4af 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -35,13 +35,19 @@ "queued_time_s": 42.0, "num_chunked_prefills": 7, "chunk_prefill_latencies_s": [0.001, 999.9, 0.003, 500.0, 0.002, 750.0, 1.0], + "chunk_prefill_start_times_s": [1000.0, 1000.01, 2000.0, 2000.5, 3000.0, 3000.1, 4000.0], "decode_latencies_s": [88888.8, 0.000005, 44444.4], + "decode_start_times_s": [5000.0, 5088888.8, 5088888.8], + "tkvs": [512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072], }, { "queued_time_s": 0.00001, "num_chunked_prefills": 3, "chunk_prefill_latencies_s": [12345.6, 0.0001, 99999.9], + "chunk_prefill_start_times_s": [1000.0, 1012345.6, 1012345.6], "decode_latencies_s": [0.000002, 77777.7], + "decode_start_times_s": [1112345.5, 1112345.5], + "tkvs": [256, 512, 1024, 2048, 4096], }, ] @@ -76,8 +82,11 @@ def test_inject_adds_spyre_keys(tmp_path): assert "spyre_queue_times_s" in data assert "spyre_num_chunked_prefills" in data assert "spyre_chunk_prefill_latencies_s" in data + assert "spyre_chunk_prefill_start_times_s" in data assert "spyre_total_prefill_chunks" in data assert "spyre_decode_latencies_s" in data + assert "spyre_decode_start_times_s" in data + assert "spyre_tkvs" in data @pytest.mark.cpu @@ -93,10 +102,22 @@ def test_inject_values_correct(tmp_path): [0.001, 999.9, 0.003, 500.0, 0.002, 750.0, 1.0], [12345.6, 0.0001, 99999.9], ] + assert data["spyre_chunk_prefill_start_times_s"] == [ + [1000.0, 1000.01, 2000.0, 2000.5, 3000.0, 3000.1, 4000.0], + [1000.0, 1012345.6, 1012345.6], + ] assert data["spyre_decode_latencies_s"] == [ [88888.8, 0.000005, 44444.4], [0.000002, 77777.7], ] + assert data["spyre_decode_start_times_s"] == [ + [5000.0, 5088888.8, 5088888.8], + [1112345.5, 1112345.5], + ] + assert data["spyre_tkvs"] == [ + [512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072], + [256, 512, 1024, 2048, 4096], + ] # Original keys preserved assert data["backend"] == "spyre-chat" From d23408e3bfbe744044fd07138384e676de887e48 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 12 Jun 2026 15:07:24 +0200 Subject: [PATCH 053/106] track all metrics in scheduler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 64 +++++++++++++------ .../v1/worker/spyre_model_runner.py | 44 +------------ 2 files changed, 45 insertions(+), 63 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 91c53478e..b4f199fc4 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -25,8 +25,8 @@ @dataclass class SpyreBenchState: - """Bench-metrics-only per-request state. Only instantiated when - SENDNN_INFERENCE_BENCH_METRICS_ENABLED is set.""" + """Bench-metrics-only state for tracking per-request and per-step timing. + Only instantiated when SENDNN_INFERENCE_BENCH_METRICS_ENABLED is set.""" chunk_latencies: dict[str, list[float]] = field(default_factory=dict) arrival_ts: dict[str, float] = field(default_factory=dict) @@ -35,6 +35,8 @@ class SpyreBenchState: decode_latencies: dict[str, list[float]] = field(default_factory=dict) decode_start_times: dict[str, list[float]] = field(default_factory=dict) tkvs: dict[str, list[int]] = field(default_factory=dict) + prefill_step_start: float | None = None + decode_step_start: float | None = None # Ensure that block_size is 64 @@ -254,29 +256,40 @@ def update_from_output(self, scheduler_output, model_runner_output): prefix_cache_len=prefix_cache_len, ) - # Accumulate per-chunk timings before removing completed prefills - if self._bench is not None and model_runner_output._bench is not None: - for req in self.ongoing_prefills: - t = model_runner_output._bench.chunk_prefill_time_s.get(req.request_id) - if t is not None: - self._bench.chunk_latencies.setdefault(req.request_id, []).append(t) - t_start = model_runner_output._bench.chunk_prefill_start_s.get(req.request_id) - if t_start is not None: - self._bench.chunk_start_times.setdefault(req.request_id, []).append(t_start) - # Track first-scheduled time and arrival time for queue-wait calculation + # Measure timing durations and accumulate metrics (scheduler-side timing injection) + if self._bench is not None: now = time.time() + + # Prefill duration measurement (if prefill step was scheduled) + if self._bench.prefill_step_start is not None: + assert self.previous_step_was_prefill and self._bench.decode_step_start is None + t0 = self._bench.prefill_step_start + duration = now - t0 + all_prefill_reqs = [ + r.req_id for r in scheduler_output.scheduled_new_reqs + ] + scheduler_output.scheduled_cached_reqs.req_ids + for req_id in all_prefill_reqs: + self._bench.chunk_latencies.setdefault(req_id, []).append(duration) + self._bench.tkvs.setdefault(req_id, []).append(model_runner_output.tkv) + self._bench.prefill_step_start = None + self._bench.decode_step_start = None + + # Decode duration measurement (if decode step was scheduled) + elif self._bench.decode_step_start is not None: + assert not self.previous_step_was_prefill and self._bench.prefill_step_start is None + t0 = self._bench.decode_step_start + duration = now - t0 + for req_id in scheduler_output.scheduled_cached_reqs.req_ids: + self._bench.decode_latencies.setdefault(req_id, []).append(duration) + self._bench.tkvs.setdefault(req_id, []).append(model_runner_output.tkv) + self._bench.prefill_step_start = None + self._bench.decode_step_start = None + + # Track first-scheduled time and arrival time for queue-wait calculation for req in self.ongoing_prefills: if req.request_id not in self._bench.first_scheduled_ts: self._bench.first_scheduled_ts[req.request_id] = now self._bench.arrival_ts[req.request_id] = req.arrival_time - # Accumulate decode timings - for req_id, t1 in model_runner_output._bench.decode_time_s.items(): - self._bench.decode_latencies.setdefault(req_id, []).append(t1) - for req_id, t0 in model_runner_output._bench.decode_start_s.items(): - self._bench.decode_start_times.setdefault(req_id, []).append(t0) - # Accumulate tkvs (one per prefill chunk and decode step) - for req_id, tkv_list in model_runner_output._bench.tkvs.items(): - self._bench.tkvs.setdefault(req_id, []).extend(tkv_list) # Remove completed prefills self.ongoing_prefills = [ @@ -487,6 +500,17 @@ def schedule(self) -> "SchedulerOutput": # TODO: Implement sample_tokens() in SpyreModelRunner to enable async grammar # collection for better performance. outputs._spyre_grammar_output = self.get_grammar_bitmask(outputs) # type: ignore[attr-defined] + + # Inject scheduler-side step-start timestamps for accurate timing measurement + if self._bench is not None: + now = time.time() + if self.previous_step_was_prefill: + self._bench.prefill_step_start = now + self._bench.decode_step_start = None + else: + self._bench.decode_step_start = now + self._bench.prefill_step_start = None + return outputs def can_schedule_prefill(self, request: Request) -> bool: diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index 9bf58b4ef..d7accd3b2 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -88,18 +88,6 @@ class SamplingForwardInputs(ModelForwardInputs): slot_mapping: torch.Tensor -@dataclass -class SpyreModelRunnerBenchState: - """Bench-metrics-only state in SpyreModelRunnerOutput. - Only instantiated when SENDNN_INFERENCE_BENCH_METRICS_ENABLED is set.""" - - chunk_prefill_time_s: dict[str, float] = field(default_factory=dict) - chunk_prefill_start_s: dict[str, float] = field(default_factory=dict) - decode_time_s: dict[str, float] = field(default_factory=dict) - decode_start_s: dict[str, float] = field(default_factory=dict) - tkvs: dict[str, list[int]] = field(default_factory=dict) - - @dataclass class SpyreModelRunnerOutput(ModelRunnerOutput): # Current tkv: this is the maximum padded request length in the batch) @@ -110,8 +98,6 @@ class SpyreModelRunnerOutput(ModelRunnerOutput): # available than the number of scheduled tokens. In that case, the scheduler # needs to update its state to reflect the correct number of computed tokens prefix_cache_hit_len: dict[str, int] = field(default_factory=dict) - # Bench-metrics-only state. None when SENDNN_INFERENCE_BENCH_METRICS_ENABLED is off. - _bench: SpyreModelRunnerBenchState | None = None InputBatchT = TypeVar("InputBatchT", bound=BaseInputBatch) @@ -1535,13 +1521,6 @@ def execute_model( ) -> ModelRunnerOutput: t0 = time.time() - # Create bench state dict only if bench metrics are enabled - bench_state = ( - SpyreModelRunnerBenchState() - if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED - else None - ) - self.update_states(scheduler_output) if not scheduler_output.total_num_scheduled_tokens: @@ -1588,15 +1567,7 @@ def execute_model( t1 = time.time() - t0 logger.debug("t_forward_pass: %.2fms [prefill single chunk][batch size 1]", (t1 * 1000)) - output = self.prefill_output() - assert isinstance(output, SpyreModelRunnerOutput) - output._bench = bench_state - if bench_state is not None: - req_ids = list(scheduler_output.num_scheduled_tokens) - if req_ids: - bench_state.chunk_prefill_time_s[req_ids[0]] = t1 - bench_state.chunk_prefill_start_s[req_ids[0]] = t0 - return output + return self.prefill_output() # Apply grammar bitmask for structured output requests. self.apply_grammar_bitmask( @@ -1640,19 +1611,6 @@ def execute_model( return self.get_empty_output() model_output = self.sampled_output(output, is_prefill) - model_output._bench = bench_state - if bench_state is not None: - req_ids = list(scheduler_output.num_scheduled_tokens) - if req_ids: - if is_prefill: - bench_state.chunk_prefill_time_s[req_ids[0]] = t1 - bench_state.chunk_prefill_start_s[req_ids[0]] = t0 - bench_state.tkvs.setdefault(req_ids[0], []).append(model_output.tkv) - else: - for req_id in req_ids: - bench_state.decode_time_s[req_id] = t1 - bench_state.decode_start_s[req_id] = t0 - bench_state.tkvs.setdefault(req_id, []).append(model_output.tkv) return model_output def prefill_output(self) -> SpyreModelRunnerOutput: From 6014899d52665cbdc5d1a0089130c23d8687a6f5 Mon Sep 17 00:00:00 2001 From: Max de Bayser Date: Fri, 12 Jun 2026 12:47:10 -0400 Subject: [PATCH 054/106] Always allow last prefill chunks onto the decode batch if it doesn't cause pausing Signed-off-by: Max de Bayser --- sendnn_inference/v1/core/scheduler.py | 51 +++------------------------ 1 file changed, 5 insertions(+), 46 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 9a5bdeb62..5bb8ca29c 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -172,7 +172,7 @@ class ChunkedPrefillSpyreScheduler(SpyreScheduler): - Volumetric constraint: the product of batch_size and current TKV must not exceed `VLLM_DT_MAX_BATCH_TKV_LIMIT` when adding a new - request. See `check_batch_tkv_limit()` method for details. + request. See `_can_decode_all_requests()` method for details. - Decode pausing: requests may be temporarily paused from decoding when the batch TKV limit would be exceeded in the next decode step. @@ -464,22 +464,10 @@ def _satisfies_last_chunk_constraints(self, request: Request) -> bool: num_running = len(decoding_requests) cond1 = num_running + len(self.waiting) < self.max_num_running_reqs - # calculate new max tkv of the batch given the new sequence joins - # considers all possible cases: - # - prompt_len > self.tkv and fall into different blocks - # - prompt_len and self.tkv fall within the same block - # - prompt_len < self.tkv and fall into different blocks - prompt_len = request.num_prompt_tokens - n_blocks = math.floor(max(self.tkv, prompt_len) / self.block_size) - new_req_tkv = n_blocks * self.block_size + prompt_len % self.block_size - - # check that adding the new request to the decode batch still have - # batch size x tkv value being smaller than the accepted limit - cond2 = lambda: self.check_batch_tkv_limit( - request=request, - new_req_tkv=new_req_tkv, - running=decoding_requests, - ) + # Check that the current decode batch is not about to have requests paused. + # This avoids adding more request to be paused and seems to slightly improve + # metrics. + cond2 = lambda: self._can_decode_all_requests(self.running) return cond1 and cond2() @@ -501,35 +489,6 @@ def _has_scheduling_priority(self, request): num_prefills = len(self.waiting) + len(self.ongoing_prefills) return num_prefills < max_concurrent_prefills - def check_batch_tkv_limit(self, request: Request, new_req_tkv: int, running) -> bool: - """ - Check whether adding a new sequence to the decode batch would immediately - violate Spyre's maximum batch volume constraint for chunked prefill. - - In Spyre, the product of `batch_size` and the current `tkv` - (tokens-per-sequence) must not exceed the limit defined by - `VLLM_DT_MAX_BATCH_TKV_LIMIT`. This checks the immediate constraint - only, not future states. - """ - # Calculate the current max tkv across all sequences - # new_req_tkv is already the current tkv for the new request - current_max_tkv = new_req_tkv - - # Check current tkv for all running requests - n_blocks = math.floor(max(self.tkv, request.num_prompt_tokens) / self.block_size) - for req in running: - dec_req_tkv = n_blocks * self.block_size + req.num_computed_tokens % self.block_size - current_max_tkv = max(current_max_tkv, dec_req_tkv) - - # Calculate batch size (including the new request if not already running) - batch_size = len(running) - if request not in running: - batch_size += 1 - - # Check immediate volume constraint - current_batch_tkv = batch_size * current_max_tkv - return current_batch_tkv <= self.max_batch_tkv_limit - def _can_decode_all_requests(self, decoding_requests: list[Request]) -> bool: """ Check if all decoding requests can be decoded in the next step without From 5f28438a844b72663255615a7df19d66f0d32e36 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 12 Jun 2026 18:01:22 +0200 Subject: [PATCH 055/106] debug tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_bench_serve.py | 26 +- sendnn_inference/benchmarks/spyre_plot.py | 31 +- sendnn_inference/v1/core/scheduler.py | 2 + tests/benchmarks/test_bench_metrics.py | 299 +++++++++++++----- 4 files changed, 263 insertions(+), 95 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index cf27d990c..4435a7f36 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -110,12 +110,12 @@ def _build_parser() -> argparse.ArgumentParser: ) parser.add_argument( "--itl-thresholds", - type=float, - nargs=2, - metavar=("LOW", "HIGH"), + type=str, + metavar="LOW,HIGH", default=None, help=( - "Two ITL thresholds in seconds for decode coloring in the detailed timeline. " + "Two ITL thresholds in milliseconds (comma-separated) for decode " + " coloring in the detailed timeline. " "Decode steps below LOW are green, between LOW and HIGH are orange, " "above HIGH are red. When omitted, all decode steps are green." ), @@ -221,7 +221,7 @@ def _inject_spyre_metrics_into_result_file( logger.warning("Failed to read vllm result JSON %s: %s", file_path, exc) return - result["spyre_queue_times_s"] = [ + result["spyre_queued_time_s"] = [ m["queued_time_s"] for m in metrics_list if "queued_time_s" in m ] result["spyre_num_chunked_prefills"] = [ @@ -351,7 +351,21 @@ def main() -> None: if candidates: json_path = Path(max(candidates, key=os.path.getmtime)) html_path = json_path.with_name(json_path.stem + "_detailed.html") - itl_thresholds = getattr(args, "itl_thresholds", None) + itl_thresholds_str = getattr(args, "itl_thresholds", None) + # Parse comma-separated milliseconds and convert to seconds + itl_thresholds = None + if itl_thresholds_str: + try: + thresholds_ms = [float(x.strip()) for x in itl_thresholds_str.split(",")] + if len(thresholds_ms) != 2: + raise ValueError("Expected exactly 2 comma-separated values") + itl_thresholds = [ms / 1000.0 for ms in thresholds_ms] + except (ValueError, AttributeError) as e: + logger.warning( + "Invalid --itl-thresholds format: %s (expected LOW,HIGH in ms)", + e, + ) + itl_thresholds = None generate_detailed_timeline_plot( _request_outputs_collected, html_path, itl_thresholds=itl_thresholds ) diff --git a/sendnn_inference/benchmarks/spyre_plot.py b/sendnn_inference/benchmarks/spyre_plot.py index e7e743d9d..901e84c5c 100644 --- a/sendnn_inference/benchmarks/spyre_plot.py +++ b/sendnn_inference/benchmarks/spyre_plot.py @@ -2,7 +2,7 @@ """Detailed per-request Gantt-chart timeline plot for sendnn-bench serve. Generates an HTML file showing, for each request: - - Queue wait time (before first prefill) + - Waiting time before first prefill (ttft - sum of prefill latencies) - Each individual prefill chunk as a separate segment - Waiting gaps between segments (absorbed if < 10% of segment duration) - Each individual decode step (with TKV in hover) @@ -15,8 +15,7 @@ logger = logging.getLogger(__name__) # Color scheme -_COLOR_QUEUE_WAIT = "#636EFA" # blue-purple — queue wait before first prefill -_COLOR_WAITING = "#777777" # dark grey — inter-segment gaps +_COLOR_WAITING = "#636EFA" # blue-purple — waiting (time before first prefill) _COLOR_PREFILL = "#FF0092" # pink — all prefill chunks _COLOR_DECODE_FAST = "#109618" # green — decode ITL below lower threshold (or all decodes) _COLOR_DECODE_MID = "#FF7F0E" # orange — decode ITL between thresholds (vLLM colors) @@ -67,9 +66,9 @@ def _build_detailed_segments( All timestamps are elapsed seconds relative to t0_global (min client start_time across all requests), formatted as HH:MM:SS.mmm strings for px.timeline. - Absolute server-side timestamps (chunk_prefill_start_times_s, decode_start_times_s) - are only used to derive *gaps between consecutive segments on the same request* — - never mixed with client-side start_time values — to avoid cross-clock skew. + ttft - sum(chunk_prefill_latencies_s) gives the waiting time before the first + prefill (plus any inter-prefill gaps, which are shown as separate "Waiting" segments). + All values are client-clock, so there is no cross-clock skew. """ client_start = (request.get("start_time") or 0.0) - t0_global latency = request.get("latency") @@ -77,7 +76,6 @@ def _build_detailed_segments( output_tokens = request.get("output_tokens") req_finish = client_start + latency if latency is not None else None - queued_time_s = request.get("queued_time_s") or 0.0 prefill_lats = request.get("chunk_prefill_latencies_s") or [] prefill_starts_abs = request.get("chunk_prefill_start_times_s") or [] decode_lats = request.get("decode_latencies_s") or [] @@ -98,16 +96,22 @@ def _build_detailed_segments( "req_finish_time": _tostr(req_finish) if req_finish is not None else "—", } - # --- Queue wait --- - # Anchored to client start_time; queued_time_s is server-measured but relative. - first_prefill_t = client_start + queued_time_s + # --- Waiting (client send → first prefill) --- + # ttft covers: waiting + all prefill chunks + inter-prefill gaps. + # Subtracting the sum of prefill latencies leaves waiting + gaps; gaps are then + # shown separately as "Waiting" segments driven by server timestamps, so this + # initial bar ends up covering only the pre-first-prefill wait. + # All values are client-clock, so there is no cross-clock skew. + ttft = request.get("ttft") or 0.0 + waiting_time_s = max(ttft - sum(prefill_lats), 0.0) + first_prefill_t = client_start + waiting_time_s segments.append( { **common, "start": _tostr(client_start), "end": _tostr(first_prefill_t), - "type": "Queue wait", - "duration": f"{queued_time_s * 1000:.1f}ms", + "type": "Waiting", + "duration": f"{waiting_time_s * 1000:.1f}ms", "tkv": "—", } ) @@ -267,12 +271,11 @@ def generate_detailed_timeline_plot( df = pd.DataFrame(all_segments) color_map: dict[str, str] = { - "Queue wait": _COLOR_QUEUE_WAIT, "Waiting": _COLOR_WAITING, "Prefill": _COLOR_PREFILL, fast_lbl: _COLOR_DECODE_FAST, } - category_order = ["Queue wait", "Prefill", "Waiting", fast_lbl] + category_order = ["Waiting", "Prefill", fast_lbl] if itl_thresholds: color_map[mid_lbl] = _COLOR_DECODE_MID color_map[slow_lbl] = _COLOR_DECODE_SLOW diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index b4f199fc4..26d953a10 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -270,6 +270,7 @@ def update_from_output(self, scheduler_output, model_runner_output): ] + scheduler_output.scheduled_cached_reqs.req_ids for req_id in all_prefill_reqs: self._bench.chunk_latencies.setdefault(req_id, []).append(duration) + self._bench.chunk_start_times.setdefault(req_id, []).append(t0) self._bench.tkvs.setdefault(req_id, []).append(model_runner_output.tkv) self._bench.prefill_step_start = None self._bench.decode_step_start = None @@ -281,6 +282,7 @@ def update_from_output(self, scheduler_output, model_runner_output): duration = now - t0 for req_id in scheduler_output.scheduled_cached_reqs.req_ids: self._bench.decode_latencies.setdefault(req_id, []).append(duration) + self._bench.decode_start_times.setdefault(req_id, []).append(t0) self._bench.tkvs.setdefault(req_id, []).append(model_runner_output.tkv) self._bench.prefill_step_start = None self._bench.decode_step_start = None diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index 0e2c3d4af..230ec71ec 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -28,8 +28,7 @@ # Shared test data # --------------------------------------------------------------------------- -# Deliberately synthetic values — order-of-magnitude differences make clear these -# are not real system measurements. +# Synthetic values FAKE_METRICS: list[dict[str, Any]] = [ { "queued_time_s": 42.0, @@ -51,7 +50,7 @@ }, ] -SELECTED_PERCENTILES = [99.0] +SELECTED_PERCENTILES = [90.0, 99.0, 100.0] def _write_fake_result(tmp_path) -> pathlib.Path: @@ -79,14 +78,9 @@ def test_inject_adds_spyre_keys(tmp_path): result_file = _write_fake_result(tmp_path) _inject_spyre_metrics_into_result_file(_make_args(tmp_path), FAKE_METRICS, time.time() - 1) data = json.loads(result_file.read_text()) - assert "spyre_queue_times_s" in data - assert "spyre_num_chunked_prefills" in data - assert "spyre_chunk_prefill_latencies_s" in data - assert "spyre_chunk_prefill_start_times_s" in data - assert "spyre_total_prefill_chunks" in data - assert "spyre_decode_latencies_s" in data - assert "spyre_decode_start_times_s" in data - assert "spyre_tkvs" in data + expected_keys = {"spyre_" + k for k in FAKE_METRICS[0]} | {"spyre_total_prefill_chunks"} + for key in expected_keys: + assert key in data, f"expected key {key!r} missing from result JSON" @pytest.mark.cpu @@ -95,29 +89,16 @@ def test_inject_values_correct(tmp_path): _inject_spyre_metrics_into_result_file(_make_args(tmp_path), FAKE_METRICS, time.time() - 1) data = json.loads(result_file.read_text()) - assert data["spyre_queue_times_s"] == pytest.approx([42.0, 0.00001]) - assert data["spyre_num_chunked_prefills"] == [7, 3] - assert data["spyre_total_prefill_chunks"] == 10 - assert data["spyre_chunk_prefill_latencies_s"] == [ - [0.001, 999.9, 0.003, 500.0, 0.002, 750.0, 1.0], - [12345.6, 0.0001, 99999.9], - ] - assert data["spyre_chunk_prefill_start_times_s"] == [ - [1000.0, 1000.01, 2000.0, 2000.5, 3000.0, 3000.1, 4000.0], - [1000.0, 1012345.6, 1012345.6], - ] - assert data["spyre_decode_latencies_s"] == [ - [88888.8, 0.000005, 44444.4], - [0.000002, 77777.7], - ] - assert data["spyre_decode_start_times_s"] == [ - [5000.0, 5088888.8, 5088888.8], - [1112345.5, 1112345.5], - ] - assert data["spyre_tkvs"] == [ - [512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072], - [256, 512, 1024, 2048, 4096], - ] + # Per-request lists map directly: "spyre_" + key → [m[key] for m in FAKE_METRICS] + for key in FAKE_METRICS[0]: + if key == "num_chunked_prefills": + continue + expected = [m[key] for m in FAKE_METRICS] + assert data["spyre_" + key] == expected + # Derived run-level scalar + assert data["spyre_total_prefill_chunks"] == sum( + m["num_chunked_prefills"] for m in FAKE_METRICS + ) # Original keys preserved assert data["backend"] == "spyre-chat" @@ -146,7 +127,7 @@ def test_inject_explicit_result_filename(tmp_path): args = _make_args(tmp_path, result_filename=str(result_file)) _inject_spyre_metrics_into_result_file(args, FAKE_METRICS, time.time() - 1) data = json.loads(result_file.read_text()) - assert "spyre_queue_times_s" in data + assert "spyre_queued_time_s" in data # --------------------------------------------------------------------------- @@ -154,16 +135,42 @@ def test_inject_explicit_result_filename(tmp_path): # --------------------------------------------------------------------------- -@pytest.mark.cpu -def test_print_scalar_total_line(capsys): - _print_spyre_section(FAKE_METRICS, SELECTED_PERCENTILES) - out = capsys.readouterr().out - assert "Total prefill chunks processed:" in out - assert "10" in out +class _TrackedOutput: + """Wraps the captured stdout of _print_spyre_section and tracks which + non-blank, non-separator lines have been covered by assert_contains(). + Call assert_all_lines_covered() at the end of the test to fail if any + content line was never asserted against.""" + + def __init__(self, text: str) -> None: + self._text = text + # Content lines: all non-empty lines except pure "=" footer/header lines + self._content_lines = [ + line for line in text.splitlines() if line.strip() and set(line.strip()) != {"="} + ] + self._covered: set[int] = set() + + def assert_contains(self, substring: str) -> None: + assert substring in self._text, f"{substring!r} not found in output" + for i, line in enumerate(self._content_lines): + if substring in line: + self._covered.add(i) + + def assert_not_contains(self, substring: str) -> None: + assert substring not in self._text, f"{substring!r} unexpectedly found in output" + + def assert_all_lines_covered(self) -> None: + uncovered = [ + self._content_lines[i] + for i in range(len(self._content_lines)) + if i not in self._covered + ] + assert not uncovered, "The following output lines were never asserted:\n" + "\n".join( + f" {line!r}" for line in uncovered + ) @pytest.mark.cpu -def test_print_sendnn_header(capsys): +def test_print_sendnn_header(): # The SenDNN header is printed by main() just before _print_spyre_section. # We verify the format string produces the expected centred header. header = "{s:{c}^{n}}".format(s=" SenDNN Metrics ", n=50, c="=") @@ -174,22 +181,35 @@ def test_print_sendnn_header(capsys): @pytest.mark.cpu -def test_print_section_headers(capsys): +def test_print_spyre_section_output(capsys): + """Single test covering every line emitted by _print_spyre_section. + Uses _TrackedOutput to enforce that no output line goes unasserted — + add assertions here when a new metric section is added.""" _print_spyre_section(FAKE_METRICS, SELECTED_PERCENTILES) - out = capsys.readouterr().out - assert "Queue Wait Time" in out - assert "Chunked Prefill Count" in out - assert "Chunked Prefill Latency" in out - assert "Decode Step Latency" in out + out = _TrackedOutput(capsys.readouterr().out) + # Run-level scalar + out.assert_contains("Total prefill chunks processed:") + out.assert_contains("10") -@pytest.mark.cpu -def test_print_mean_median_p99(capsys): - _print_spyre_section(FAKE_METRICS, SELECTED_PERCENTILES) - out = capsys.readouterr().out - assert "Mean Queue Wait Time (ms):" in out - assert "Median Queue Wait Time (ms):" in out - assert "P99 Queue Wait Time (ms):" in out + # Section separators and mean/median/percentile lines for all four sections + out.assert_contains("Queue Wait Time") + out.assert_contains("Chunked Prefill Count") + out.assert_contains("Chunked Prefill Latency") + out.assert_contains("Decode Step Latency") + + for label in ( + "Queue Wait Time (ms)", + "Num Chunked Prefills", + "Chunk Prefill Latency (ms)", + "Decode Step Latency (ms)", + ): + out.assert_contains(f"Mean {label}:") + out.assert_contains(f"Median {label}:") + for percentile in SELECTED_PERCENTILES: + out.assert_contains(f"P{int(percentile)} {label}:") + + out.assert_all_lines_covered() @pytest.mark.cpu @@ -231,16 +251,94 @@ def _make_bare_scheduler(): return s +# One entry per field of SpyreBenchState — test_bench_fixture_covers_all_per_req_fields +# will fail if this is out of sync. For dict fields the value is used as the test +# payload (populated as bench.["r0"] = value). Scalar fields are set directly. +_BENCH_FIXTURE: dict[str, Any] = { + "chunk_latencies": [88888.8, 0.000005], + "arrival_ts": 1000.0, + "first_scheduled_ts": 1001.0, + "chunk_start_times": [1000.0, 1088888.8], + "decode_latencies": [0.1, 0.2], + "decode_start_times": [2000.0, 2000.1], + "tkvs": [64, 128, 192, 256], + "prefill_step_start": 999.0, + "decode_step_start": 1999.0, +} + +# Expected return value of get_and_clear_chunk_stats. Add an entry here when adding +# a new key to its return dict — the structural test will fail until you do. +_EXPECTED_RESULT: dict[str, Any] = { + "num_chunked_prefills": 2, + "chunk_prefill_latencies_s": pytest.approx([88888.8, 0.000005]), + "chunk_prefill_start_times_s": pytest.approx([1000.0, 1088888.8]), + "decode_latencies_s": pytest.approx([0.1, 0.2]), + "decode_start_times_s": pytest.approx([2000.0, 2000.1]), + "tkvs": [64, 128, 192, 256], +} + + +@pytest.mark.cpu +def test_bench_fixture_covers_all_per_req_fields(): + """_BENCH_FIXTURE must contain one entry per field of SpyreBenchState. + Fails when a field is added to SpyreBenchState without updating _BENCH_FIXTURE.""" + from dataclasses import fields as dc_fields + + s = _make_bare_scheduler() + assert s._bench is not None + all_fields = {f.name for f in dc_fields(s._bench)} + assert _BENCH_FIXTURE.keys() == all_fields, ( + f"_BENCH_FIXTURE is out of sync with SpyreBenchState fields.\n" + f" extra : {_BENCH_FIXTURE.keys() - all_fields}\n" + f" missing: {all_fields - _BENCH_FIXTURE.keys()}" + ) + + +@pytest.mark.cpu +def test_get_and_clear_result_keys(): + """get_and_clear_chunk_stats must return exactly the keys in _EXPECTED_RESULT. + Fails when a key is added or removed without updating _EXPECTED_RESULT.""" + s = _make_bare_scheduler() + assert s._bench is not None + bench = s._bench + for field, value in _BENCH_FIXTURE.items(): + attr = getattr(bench, field) + if isinstance(attr, dict): + attr["r0"] = value + else: + setattr(bench, field, value) + result = s.get_and_clear_chunk_stats("r0") + assert result is not None + assert result.keys() == _EXPECTED_RESULT.keys(), ( + f"get_and_clear_chunk_stats returned unexpected keys.\n" + f" extra : {result.keys() - _EXPECTED_RESULT.keys()}\n" + f" missing: {_EXPECTED_RESULT.keys() - result.keys()}" + ) + + @pytest.mark.cpu def test_get_and_clear_returns_correct_dict(): s = _make_bare_scheduler() - s._bench.chunk_latencies["r0"] = [88888.8, 0.000005] + assert s._bench is not None + bench = s._bench + + for field, value in _BENCH_FIXTURE.items(): + attr = getattr(bench, field) + if isinstance(attr, dict): + attr["r0"] = value + else: + setattr(bench, field, value) + result = s.get_and_clear_chunk_stats("r0") assert result is not None - assert result["num_chunked_prefills"] == 2 - assert result["chunk_prefill_latencies_s"] == pytest.approx([88888.8, 0.000005]) - # Entry must be cleared after retrieval - assert "r0" not in s._bench.chunk_latencies + + for key, expected in _EXPECTED_RESULT.items(): + assert result[key] == expected, f"result[{key!r}] mismatch" + + # All per-request bench fields must be cleared after retrieval + for field, value in _BENCH_FIXTURE.items(): + if isinstance(value, list): + assert "r0" not in getattr(bench, field), f"bench.{field} was not cleared for r0" @pytest.mark.cpu @@ -325,13 +423,24 @@ def test_scheduler_bench_metrics_accumulated( original_free = scheduler.__class__._free_request def _capturing_free(self, request, delay_free_blocks=False): + from dataclasses import fields as dc_fields + req_id = request.request_id bench = self._bench - captured[req_id] = { - "chunk_latencies": list(bench.chunk_latencies.get(req_id, [])) if bench else [], - "has_arrival_ts": (req_id in bench.arrival_ts) if bench else False, - "has_first_scheduled_ts": (req_id in bench.first_scheduled_ts) if bench else False, - } + # Snapshot all dict fields dynamically so new metrics are captured automatically. + # List-valued dicts are copied; scalar-valued dicts (arrival_ts, first_scheduled_ts) + # are stored as-is so truthiness checks work correctly. + if bench: + snap = {} + for f in dc_fields(bench): + val = getattr(bench, f.name) + if not isinstance(val, dict): + continue + entry = val.get(req_id) + snap[f.name] = list(entry) if isinstance(entry, list) else entry + captured[req_id] = snap + else: + captured[req_id] = {} return original_free(self, request, delay_free_blocks) scheduler._free_request = _capturing_free.__get__(scheduler) @@ -365,15 +474,53 @@ def _capturing_free(self, request, delay_free_blocks=False): assert req_id in captured, f"_free_request was never called for req {req_id}" info = captured[req_id] - # At least 2 chunks recorded (prompt > chunk_size) + # At least 2 prefill chunks recorded (prompt > chunk_size) assert len(info["chunk_latencies"]) >= 2, ( f"req {req_id}: expected ≥2 chunk latencies, got {info['chunk_latencies']}" ) - # All latencies must be positive floats + # All prefill latencies must be positive floats for lat in info["chunk_latencies"]: assert isinstance(lat, float) and lat > 0, f"req {req_id}: non-positive latency {lat}" - assert info["has_arrival_ts"], f"req {req_id}: _arrival_ts not set" - assert info["has_first_scheduled_ts"], f"req {req_id}: _first_scheduled_ts not set" + + # chunk_start_times must match chunk_latencies in length + assert len(info["chunk_start_times"]) == len(info["chunk_latencies"]), ( + f"req {req_id}: chunk_start_times length {len(info['chunk_start_times'])} " + f"!= chunk_latencies length {len(info['chunk_latencies'])}" + ) + for ts in info["chunk_start_times"]: + assert isinstance(ts, float) and ts > 0, ( + f"req {req_id}: non-positive chunk_start_time {ts}" + ) + + # Decode latencies: at least 1 decode step after the prefill phase + assert len(info["decode_latencies"]) >= 1, ( + f"req {req_id}: expected ≥1 decode latency, got {info['decode_latencies']}" + ) + for lat in info["decode_latencies"]: + assert isinstance(lat, float) and lat > 0, ( + f"req {req_id}: non-positive decode latency {lat}" + ) + + # decode_start_times must match decode_latencies in length + assert len(info["decode_start_times"]) == len(info["decode_latencies"]), ( + f"req {req_id}: decode_start_times length {len(info['decode_start_times'])} " + f"!= decode_latencies length {len(info['decode_latencies'])}" + ) + for ts in info["decode_start_times"]: + assert isinstance(ts, float) and ts > 0, ( + f"req {req_id}: non-positive decode_start_time {ts}" + ) + + # tkvs: one entry per prefill chunk + per decode step + expected_tkvs = len(info["chunk_latencies"]) + len(info["decode_latencies"]) + assert len(info["tkvs"]) == expected_tkvs, ( + f"req {req_id}: expected {expected_tkvs} tkvs, got {info['tkvs']}" + ) + for tkv in info["tkvs"]: + assert isinstance(tkv, int) and tkv > 0, f"req {req_id}: non-positive tkv {tkv}" + + assert info["arrival_ts"] is not None, f"req {req_id}: arrival_ts not set" + assert info["first_scheduled_ts"] is not None, f"req {req_id}: first_scheduled_ts not set" # The two requests must have independent latency lists (no cross-contamination) assert captured["0"]["chunk_latencies"] != captured["1"]["chunk_latencies"] or ( @@ -382,13 +529,15 @@ def _capturing_free(self, request, delay_free_blocks=False): len(captured["0"]["chunk_latencies"]) >= 2 and len(captured["1"]["chunk_latencies"]) >= 2 ) - # After all requests finished, the bench state dicts must be empty + # After all requests finished, all bench state dicts must be empty. + # Checked dynamically so new fields are covered automatically. + from dataclasses import fields as dc_fields + assert scheduler._bench is not None - assert scheduler._bench.chunk_latencies == {}, "Leftover entries in chunk_latencies after run" - assert scheduler._bench.arrival_ts == {}, "Leftover entries in arrival_ts after run" - assert scheduler._bench.first_scheduled_ts == {}, ( - "Leftover entries in first_scheduled_ts after run" - ) + for f in dc_fields(scheduler._bench): + val = getattr(scheduler._bench, f.name) + if isinstance(val, dict): + assert val == {}, f"Leftover entries in {f.name} after run" # --------------------------------------------------------------------------- From 636d4c1f8f1bdf8d530b889508f6e343a2ac47d4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 12 Jun 2026 20:44:31 +0200 Subject: [PATCH 056/106] update timeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_bench_serve.py | 24 ++++---- sendnn_inference/benchmarks/spyre_plot.py | 61 +++++++++---------- 2 files changed, 41 insertions(+), 44 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 4435a7f36..ee30e9e1a 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -109,13 +109,13 @@ def _build_parser() -> argparse.ArgumentParser: ), ) parser.add_argument( - "--itl-thresholds", + "--decode-thresholds", type=str, metavar="LOW,HIGH", default=None, help=( - "Two ITL thresholds in milliseconds (comma-separated) for decode " - " coloring in the detailed timeline. " + "Two decode latency thresholds in milliseconds (comma-separated) for " + "coloring in the detailed timeline. " "Decode steps below LOW are green, between LOW and HIGH are orange, " "above HIGH are red. When omitted, all decode steps are green." ), @@ -350,24 +350,24 @@ def main() -> None: if candidates: json_path = Path(max(candidates, key=os.path.getmtime)) - html_path = json_path.with_name(json_path.stem + "_detailed.html") - itl_thresholds_str = getattr(args, "itl_thresholds", None) + html_path = json_path.with_name(json_path.stem + "_detailed_timeline.html") + decode_thresholds_str = getattr(args, "decode_thresholds", None) # Parse comma-separated milliseconds and convert to seconds - itl_thresholds = None - if itl_thresholds_str: + decode_thresholds = None + if decode_thresholds_str: try: - thresholds_ms = [float(x.strip()) for x in itl_thresholds_str.split(",")] + thresholds_ms = [float(x.strip()) for x in decode_thresholds_str.split(",")] if len(thresholds_ms) != 2: raise ValueError("Expected exactly 2 comma-separated values") - itl_thresholds = [ms / 1000.0 for ms in thresholds_ms] + decode_thresholds = [ms / 1000.0 for ms in thresholds_ms] except (ValueError, AttributeError) as e: logger.warning( - "Invalid --itl-thresholds format: %s (expected LOW,HIGH in ms)", + "Invalid --decode-thresholds format: %s (expected LOW,HIGH in ms)", e, ) - itl_thresholds = None + decode_thresholds = None generate_detailed_timeline_plot( - _request_outputs_collected, html_path, itl_thresholds=itl_thresholds + _request_outputs_collected, html_path, decode_thresholds=decode_thresholds ) else: logger.warning( diff --git a/sendnn_inference/benchmarks/spyre_plot.py b/sendnn_inference/benchmarks/spyre_plot.py index 901e84c5c..40a07217e 100644 --- a/sendnn_inference/benchmarks/spyre_plot.py +++ b/sendnn_inference/benchmarks/spyre_plot.py @@ -2,7 +2,7 @@ """Detailed per-request Gantt-chart timeline plot for sendnn-bench serve. Generates an HTML file showing, for each request: - - Waiting time before first prefill (ttft - sum of prefill latencies) + - Queue wait time (before first prefill) - Each individual prefill chunk as a separate segment - Waiting gaps between segments (absorbed if < 10% of segment duration) - Each individual decode step (with TKV in hover) @@ -15,11 +15,12 @@ logger = logging.getLogger(__name__) # Color scheme -_COLOR_WAITING = "#636EFA" # blue-purple — waiting (time before first prefill) +_COLOR_QUEUE_WAIT = "#636EFA" # blue-purple — queue wait before first prefill +_COLOR_WAITING = "#777777" # dark grey — inter-segment gaps _COLOR_PREFILL = "#FF0092" # pink — all prefill chunks -_COLOR_DECODE_FAST = "#109618" # green — decode ITL below lower threshold (or all decodes) -_COLOR_DECODE_MID = "#FF7F0E" # orange — decode ITL between thresholds (vLLM colors) -_COLOR_DECODE_SLOW = "#D62728" # red — decode ITL above upper threshold +_COLOR_DECODE_FAST = "#109618" # green — decode latency below lower threshold (or all decodes) +_COLOR_DECODE_MID = "#FF7F0E" # orange — decode latency between thresholds (vLLM colors) +_COLOR_DECODE_SLOW = "#D62728" # red — decode latency above upper threshold _GAP_ABSORPTION_THRESHOLD = 0.10 # absorb gap if < 10% of current segment duration @@ -38,9 +39,9 @@ def _decode_labels(thresholds: list[float] | None) -> tuple[str, str, str]: lo_ms = int(thresholds[0] * 1000) hi_ms = int(thresholds[1] * 1000) return ( - f"ITL < {lo_ms}ms", - f"{lo_ms}ms ≤ ITL < {hi_ms}ms", - f"ITL ≥ {hi_ms}ms", + f"Decode latency < {lo_ms}ms", + f"{lo_ms}ms ≤ Decode latency < {hi_ms}ms", + f"Decode latency ≥ {hi_ms}ms", ) @@ -58,7 +59,7 @@ def _decode_type(lat: float, thresholds: list[float] | None, labels: tuple[str, def _build_detailed_segments( request: dict[str, Any], t0_global: float, - itl_thresholds: list[float] | None = None, + decode_thresholds: list[float] | None = None, decode_labels: tuple[str, str, str] | None = None, ) -> list[dict[str, Any]]: """Convert one request's timing data into ordered Gantt segments. @@ -66,9 +67,9 @@ def _build_detailed_segments( All timestamps are elapsed seconds relative to t0_global (min client start_time across all requests), formatted as HH:MM:SS.mmm strings for px.timeline. - ttft - sum(chunk_prefill_latencies_s) gives the waiting time before the first - prefill (plus any inter-prefill gaps, which are shown as separate "Waiting" segments). - All values are client-clock, so there is no cross-clock skew. + Absolute server-side timestamps (chunk_prefill_start_times_s, decode_start_times_s) + are only used to derive *gaps between consecutive segments on the same request* — + never mixed with client-side start_time values — to avoid cross-clock skew. """ client_start = (request.get("start_time") or 0.0) - t0_global latency = request.get("latency") @@ -76,6 +77,7 @@ def _build_detailed_segments( output_tokens = request.get("output_tokens") req_finish = client_start + latency if latency is not None else None + queued_time_s = request.get("queued_time_s") or 0.0 prefill_lats = request.get("chunk_prefill_latencies_s") or [] prefill_starts_abs = request.get("chunk_prefill_start_times_s") or [] decode_lats = request.get("decode_latencies_s") or [] @@ -96,22 +98,16 @@ def _build_detailed_segments( "req_finish_time": _tostr(req_finish) if req_finish is not None else "—", } - # --- Waiting (client send → first prefill) --- - # ttft covers: waiting + all prefill chunks + inter-prefill gaps. - # Subtracting the sum of prefill latencies leaves waiting + gaps; gaps are then - # shown separately as "Waiting" segments driven by server timestamps, so this - # initial bar ends up covering only the pre-first-prefill wait. - # All values are client-clock, so there is no cross-clock skew. - ttft = request.get("ttft") or 0.0 - waiting_time_s = max(ttft - sum(prefill_lats), 0.0) - first_prefill_t = client_start + waiting_time_s + # --- Queue wait --- + # Anchored to client start_time; queued_time_s is server-measured but relative. + first_prefill_t = client_start + queued_time_s segments.append( { **common, "start": _tostr(client_start), "end": _tostr(first_prefill_t), - "type": "Waiting", - "duration": f"{waiting_time_s * 1000:.1f}ms", + "type": "Queue wait", + "duration": f"{queued_time_s * 1000:.1f}ms", "tkv": "—", } ) @@ -155,7 +151,7 @@ def _build_detailed_segments( **common, "start": _tostr(seg_start), "end": _tostr(seg_end), - "type": "Prefill", + "type": "Chunked Prefill", "duration": f"{lat * 1000:.1f}ms", "tkv": str(tkv) if tkv is not None else "—", } @@ -209,7 +205,7 @@ def _build_detailed_segments( **common, "start": _tostr(seg_start), "end": _tostr(seg_end), - "type": _decode_type(lat, itl_thresholds, decode_labels or ("Decode", "", "")), + "type": _decode_type(lat, decode_thresholds, decode_labels or ("Decode", "", "")), "duration": f"{lat * 1000:.1f}ms", "tkv": str(tkv) if tkv is not None else "—", } @@ -222,12 +218,12 @@ def _build_detailed_segments( def generate_detailed_timeline_plot( requests: list[dict[str, Any]], output_path: Path, - itl_thresholds: list[float] | None = None, + decode_thresholds: list[float] | None = None, ) -> None: """Build a per-request Gantt-chart HTML and write it to output_path. Args: - itl_thresholds: Two thresholds in seconds [low, high]. Decode steps below + decode_thresholds: Two thresholds in seconds [low, high]. Decode steps below low are green, between low and high are orange, above high are red. When None (default), all decode steps are green. """ @@ -257,12 +253,12 @@ def generate_detailed_timeline_plot( for idx, req in enumerate(sorted_requests): req["_label"] = f"Req {idx}" - labels = _decode_labels(itl_thresholds) + labels = _decode_labels(decode_thresholds) fast_lbl, mid_lbl, slow_lbl = labels all_segments: list[dict[str, Any]] = [] for req in sorted_requests: - all_segments.extend(_build_detailed_segments(req, t0_global, itl_thresholds, labels)) + all_segments.extend(_build_detailed_segments(req, t0_global, decode_thresholds, labels)) if not all_segments: logger.warning("No plottable segments found — skipping detailed timeline.") @@ -271,12 +267,13 @@ def generate_detailed_timeline_plot( df = pd.DataFrame(all_segments) color_map: dict[str, str] = { + "Queue wait": _COLOR_QUEUE_WAIT, "Waiting": _COLOR_WAITING, - "Prefill": _COLOR_PREFILL, + "Chunked Prefill": _COLOR_PREFILL, fast_lbl: _COLOR_DECODE_FAST, } - category_order = ["Waiting", "Prefill", fast_lbl] - if itl_thresholds: + category_order = ["Queue wait", "Chunked Prefill", "Waiting", fast_lbl] + if decode_thresholds: color_map[mid_lbl] = _COLOR_DECODE_MID color_map[slow_lbl] = _COLOR_DECODE_SLOW category_order += [mid_lbl, slow_lbl] From 885f0374a618f219f1e39b2112dc32c56c2ab23b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 12 Jun 2026 20:45:34 +0200 Subject: [PATCH 057/106] add claude skill MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .claude/skills/add-bench-metric/SKILL.md | 270 +++++++++++++++++++++++ 1 file changed, 270 insertions(+) create mode 100644 .claude/skills/add-bench-metric/SKILL.md diff --git a/.claude/skills/add-bench-metric/SKILL.md b/.claude/skills/add-bench-metric/SKILL.md new file mode 100644 index 000000000..fad1fe184 --- /dev/null +++ b/.claude/skills/add-bench-metric/SKILL.md @@ -0,0 +1,270 @@ +--- +name: add-bench-metric +description: Add a new custom per-request benchmark metric to sendnn-bench serve. Covers every layer of the pipeline: scheduler-side timing → SpyreBenchState → _free_request → kv_transfer_params → SSE injection → client parsing → aggregation and printing → result JSON injection. Use when the user says "add a new bench metric", "track X per request", "surface Y in bench serve output", or similar. +argument-hint: +--- + +Add a new custom per-request benchmark metric: **$ARGUMENTS**. + +If `$ARGUMENTS` is empty, ask the user what they want to measure before doing anything else. + +--- + +## Background — the pipeline + +Every custom bench metric travels through five layers (all scheduler-process-side): + +``` +schedule() → SpyreBenchState (accumulate raw values) + → _free_request() → kv_transfer_params["__spyre__"] (ZMQ to API server) + → patch_serving.py SSE injection (into final usage SSE chunk) + → async_request_spyre_chat() → output.custom_metrics_dict + → _print_spyre_section() + _inject_spyre_metrics_into_result_file() +``` + +**Key architecture facts:** +- All timing is measured in the scheduler. +- `SpyreBenchState` is the single source of truth for per-request bench state. It is `None` when `SENDNN_INFERENCE_BENCH_METRICS_ENABLED` is off; every access must be guarded by `if self._bench is not None:`. +- Timing works by bracket: at the **end** of `schedule()`, a step-start timestamp (`prefill_step_start` or `decode_step_start`) is written into `_bench`. At the **start** of the next `update_from_output()`, the duration is computed from that timestamp and `time.time()`. +- `stats_logger.py` contains `SpyreRequestMetrics` / `SpyreMetricsRegistry` / `_SCHEDULER` which are **not part of the active metric pipeline** — do not add new metric fields there. +- `dataclasses.asdict()` is **not** used to serialize `SpyreBenchState` — the fields are serialised manually in `_free_request()` into a plain `dict` assigned to `kv_transfer_params["__spyre__"]`. + +--- + +## Step 0 — Understand the new metric + +Before writing any code, answer these questions (ask the user if unclear): + +1. **Where is the raw value available?** Scheduler state (e.g., a timestamp set during `schedule()`)? `FinishedRequestStats` from vllm (e.g., `r.queued_time`)? +2. **What is the cardinality?** One scalar per request? One value per prefill chunk? One value per decode step? A list per request? +3. **What unit?** Seconds (convert to ms in the output)? Count? Bytes? +4. **Aggregation?** Mean/Median/PXX like the existing metrics (use `_section()`), or a run-level scalar (use a plain print line)? + +--- + +## Step 1 — Add a field to `SpyreBenchState` + +**File**: `sendnn_inference/v1/core/scheduler.py` + +Add your field to the dataclass at the top of the file: + +```python +@dataclass +class SpyreBenchState: + chunk_latencies: dict[str, list[float]] = field(default_factory=dict) + chunk_start_times: dict[str, list[float]] = field(default_factory=dict) + arrival_ts: dict[str, float] = field(default_factory=dict) + first_scheduled_ts: dict[str, float] = field(default_factory=dict) + decode_latencies: dict[str, list[float]] = field(default_factory=dict) + decode_start_times: dict[str, list[float]] = field(default_factory=dict) + tkvs: dict[str, list[int]] = field(default_factory=dict) + prefill_step_start: float | None = None + decode_step_start: float | None = None + # NEW — keyed by request_id, choose the container type for your cardinality: + my_new_metric: dict[str, ] = field(default_factory=dict) +``` + +--- + +## Step 2 — Populate the field in `update_from_output()` or `schedule()` + +**File**: `sendnn_inference/v1/core/scheduler.py` + +**If the value is a timing measurement** (most common): set a start timestamp at the end of `schedule()` (alongside the existing `prefill_step_start`/`decode_step_start` pattern), then compute and accumulate the duration in `update_from_output()` (alongside the existing chunk/decode duration blocks). + +Example — adding a "time from first prefill to decode start" metric: + +```python +# In update_from_output(), under `if self._bench is not None:` +if self._bench.prefill_step_start is not None: + duration = now - self._bench.prefill_step_start + for req_id in all_prefill_reqs: + self._bench.my_new_metric.setdefault(req_id, []).append(duration) +``` + +**If the value comes from scheduler state** (e.g. request attributes): accumulate directly in `update_from_output()` under `if self._bench is not None:`: + +```python +for req in self.ongoing_prefills: + v = + if v is not None: + self._bench.my_new_metric[req.request_id] = v +``` + +Always guard with `if self._bench is not None:`. Never read `self._bench` without this guard. + +--- + +## Step 3 — Expose in `get_and_clear_chunk_stats()` + +**File**: `sendnn_inference/v1/core/scheduler.py` + +`get_and_clear_chunk_stats()` retrieves and removes per-request bench data in one atomic operation. Add your new field: + +```python +def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: + if self._bench is None: + return None + lats = self._bench.chunk_latencies.pop(req_id, None) + # ... existing pops ... + my_val = self._bench.my_new_metric.pop(req_id, None) # NEW + if lats is None and dec_lats is None: + return None + return { + "num_chunked_prefills": len(lats) if lats else 0, + "chunk_prefill_latencies_s": lats or [], + # ... existing fields ... + "my_new_metric": my_val or , # NEW + } +``` + +--- + +## Step 4 — Pack into `kv_transfer_params["__spyre__"]` in `_free_request()` + +**File**: `sendnn_inference/v1/core/scheduler.py` + +`_free_request()` builds the `spyre_data` dict that travels over ZMQ to the API server. Add your new field here: + +```python +spyre_data = { + "queued_time_s": queued_time_s, + "num_chunked_prefills": chunk_stats["num_chunked_prefills"] if chunk_stats else 0, + "chunk_prefill_latencies_s": chunk_stats["chunk_prefill_latencies_s"] if chunk_stats else [], + # ... existing fields ... + "my_new_metric": chunk_stats["my_new_metric"] if chunk_stats else , # NEW +} +``` + +This dict is what `patch_serving.py` picks up and injects into the SSE stream. No change to `patch_serving.py` itself is needed — it passes the whole dict through as `spyre_metrics`. + +--- + +## Step 5 — No changes needed in `patch_serving.py` or `spyre_request_func.py` + +`patch_serving.py` injects the entire `__spyre__` dict into the SSE chunk as-is. `async_request_spyre_chat` in `spyre_request_func.py` reads the `spyre_metrics` key from the final SSE usage chunk into `output.custom_metrics_dict` — your new key is automatically included. Verify by printing `output.custom_metrics_dict` in a test run if needed. + +--- + +## Step 6 — Aggregate and print + +**File**: `sendnn_inference/benchmarks/spyre_bench_serve.py` + +In `_print_spyre_section()`, choose the output style based on cardinality: + +**Option A — per-request distribution (Mean/Median/PXX)**: use `_section()`. + +```python +# Scalar per request: +my_values = [m["my_new_metric"] for m in metrics_list if "my_new_metric" in m] + +# List per request (flatten across all requests): +my_values = [v for m in metrics_list for v in m.get("my_new_metric", [])] + +_section("My New Metric", my_values, "My New Metric (unit)") +``` + +**Option B — run-level scalar summary**: print a plain line *above* the `_section()` calls, mirroring the `total_prefill_chunks` pattern. + +```python +total = sum(m.get("my_new_metric", 0) for m in metrics_list) +print("{:<40} {:<10}".format("My run-level total:", total)) +``` + +--- + +## Step 7 — Inject into the result JSON + +**File**: `sendnn_inference/benchmarks/spyre_bench_serve.py` + +In `_inject_spyre_metrics_into_result_file()`, add your new key alongside the existing ones: + +```python +# Per-request list (parallel to vllm's ttfts, itls, …) +result["spyre_my_new_metric"] = [m.get("my_new_metric", ) for m in metrics_list] + +# Or a run-level scalar derived from an already-written list: +result["spyre_my_total"] = sum(result["spyre_my_new_metric"]) +``` + +Use a `spyre_` prefix so the key is clearly SenDNN-owned in the vllm result JSON. + +--- + +## Step 8 — Add tests + +**File**: `tests/benchmarks/test_bench_metrics.py` + +1. **Update `FAKE_METRICS`**: Add your new key with obviously synthetic values to both dict entries in the list. + +2. **Update `test_inject_adds_spyre_keys`**: Assert the `spyre_`-prefixed key is present in the result JSON. + +3. **Update `test_inject_values_correct`**: Assert the exact values are injected correctly for both requests. + +4. **Update `test_scheduler_bench_metrics_accumulated`**: This is an integration test that runs a real engine and captures the bench state just before `_free_request` clears it. The capture (`_capturing_free`) and the post-run empty-dict check are both dynamic — they iterate over `dataclasses.fields(bench)` and need no changes. What you **do** need to add is a value assertion for your new field in the `for req_id in ("0", "1"):` block, following the pattern of the existing ones. For example, for a new list-per-step field: + +```python +# In the for req_id in ("0", "1"): block: +assert len(info["my_new_metric"]) >= 1, ( + f"req {req_id}: expected ≥1 my_new_metric entry, got {info['my_new_metric']}" +) +for val in info["my_new_metric"]: + assert isinstance(val, float) and val > 0, f"req {req_id}: non-positive my_new_metric {val}" +``` + +For a scalar field (`arrival_ts`-style), check `is not None`: +```python +assert info["my_scalar_field"] is not None, f"req {req_id}: my_scalar_field not set" +``` + +5. **Update `test_get_and_clear_returns_correct_dict`**: The test is driven by two dicts defined just above it — update both: + + - **`_BENCH_FIXTURE`** — add an entry for your new `SpyreBenchState` field. This dict must cover **every** field of `SpyreBenchState` (dict and scalar alike); `test_bench_fixture_covers_all_per_req_fields` compares `_BENCH_FIXTURE.keys()` against `dataclasses.fields(bench)` and fails if they diverge. For dict fields the value is used as the per-request payload (`bench.["r0"] = value`); for scalar fields it is set directly (`setattr(bench, field, value)`). + - **`_EXPECTED_RESULT`** — add an entry `"": `. `test_get_and_clear_result_keys` asserts `result.keys() == _EXPECTED_RESULT.keys()`, so it will fail if the returned dict has any extra or missing keys. + +Example: +```python +# In FAKE_METRICS entry 1: +"my_new_metric": [0.001, 0.002], + +# In test_inject_adds_spyre_keys: +assert "spyre_my_new_metric" in data + +# In test_inject_values_correct: +assert data["spyre_my_new_metric"] == [[0.001, 0.002], [0.003]] + +# In the sentinel dicts (dict field example): +_BENCH_FIXTURE: dict[str, Any] = { + ..., + "my_new_metric": [0.001, 0.002], # NEW — dict field, keyed by req_id at test time +} +_EXPECTED_RESULT: dict[str, Any] = { + ..., + "my_new_metric_s": pytest.approx([0.001, 0.002]), # NEW — key in returned dict +} + +# Scalar field example (e.g. a single float per request stored directly): +_BENCH_FIXTURE: dict[str, Any] = { + ..., + "my_scalar_field": 42.0, # NEW — set directly via setattr +} +``` + +6. **Update `test_print_spyre_section_output`** (if you added a new `_section()` call): Add `out.assert_contains("")` for the section separator, and add the label string (third argument to `_section()`) to the `for label in (...)` loop so mean/median/percentile lines are asserted. `assert_all_lines_covered()` is called at the end of the test and will fail if any output line was not covered by an assertion — the error message lists the exact uncovered lines. + +--- + +## Checklist + +Before declaring done: + +- [ ] Field added to `SpyreBenchState` with the correct container type (scalar, list, dict-of-list…) +- [ ] Populated in `update_from_output()` (or `schedule()`) under `if self._bench is not None:` guard +- [ ] Retrieved via `.pop()` in `get_and_clear_chunk_stats()` with a safe default if absent +- [ ] Key added to the `spyre_data` dict in `_free_request()` with a safe fallback when `chunk_stats is None` +- [ ] `_print_spyre_section()` updated with a new `_section()` call or run-level print line +- [ ] `_inject_spyre_metrics_into_result_file()` updated with a new `result["spyre_..."]` key +- [ ] Tests updated: `FAKE_METRICS`, `test_inject_adds_spyre_keys`, `test_inject_values_correct`, `_BENCH_FIXTURE`, `_EXPECTED_RESULT` (drives `test_get_and_clear_returns_correct_dict`), `test_scheduler_bench_metrics_accumulated` adapted accordingly, `test_print_spyre_section_output` updated if a new `_section()` was added +- [ ] No changes to `patch_serving.py`, `spyre_request_func.py`, or any model runner file +- [ ] No changes to `SpyreRequestMetrics` in `stats_logger.py` (not part of the active pipeline) +- [ ] All new tracking is gated by `self._bench is not None` (enforced by `SpyreBenchState` being `None` when env var is off) From b2556a9ae9727340806b6a22441f1e620c45eb79 Mon Sep 17 00:00:00 2001 From: Max de Bayser Date: Fri, 12 Jun 2026 15:44:20 -0400 Subject: [PATCH 058/106] cleanup block pool before tests Signed-off-by: Max de Bayser --- tests/e2e/test_spyre_decode_pause_scheduler_steps.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/e2e/test_spyre_decode_pause_scheduler_steps.py b/tests/e2e/test_spyre_decode_pause_scheduler_steps.py index 01c64375d..237df9ef3 100644 --- a/tests/e2e/test_spyre_decode_pause_scheduler_steps.py +++ b/tests/e2e/test_spyre_decode_pause_scheduler_steps.py @@ -257,6 +257,7 @@ def test_max_batch_tkv_decode_pausing( max_batch_tkv_limit=max_batch_tkv_limit, max_num_batched_tokens=max_num_batched_tokens, extra_assert_funcs=[verify_block_tables], + prefix_caching=True, ) @@ -450,4 +451,5 @@ def test_prefill_exceeds_max_batch_tkv( max_batch_tkv_limit=max_batch_tkv_limit, max_num_batched_tokens=max_num_batched_tokens, extra_assert_funcs=[verify_block_tables], + prefix_caching=True, ) From 3a51f7df4ff4ac530a7d79a439a171e403277935 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 12 Jun 2026 22:03:51 +0200 Subject: [PATCH 059/106] display prefix cache hit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/benchmarks/spyre_bench_serve.py | 8 ++++++++ sendnn_inference/v1/core/scheduler.py | 7 ++++++- tests/benchmarks/test_bench_metrics.py | 6 +++++- 3 files changed, 19 insertions(+), 2 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index ee30e9e1a..71663370f 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -160,10 +160,15 @@ def _section(header: str, values: list[float], label: str) -> None: ) ) + cache_hit_pcts = [ + m["prefix_cache_hit_pct"] * 100 for m in metrics_list if "prefix_cache_hit_pct" in m + ] + _section("Queue Wait Time", queue_times_ms, "Queue Wait Time (ms)") _section("Chunked Prefill Count", num_chunks_list, "Num Chunked Prefills") _section("Chunked Prefill Latency", chunk_lats_ms, "Chunk Prefill Latency (ms)") _section("Decode Step Latency", decode_lats_ms, "Decode Step Latency (ms)") + _section("Prefix Cache Hit", cache_hit_pcts, "Prefix Cache Hit (%)") print("=" * 50) @@ -237,6 +242,9 @@ def _inject_spyre_metrics_into_result_file( result["spyre_decode_latencies_s"] = [m.get("decode_latencies_s", []) for m in metrics_list] result["spyre_decode_start_times_s"] = [m.get("decode_start_times_s", []) for m in metrics_list] result["spyre_tkvs"] = [m.get("tkvs", []) for m in metrics_list] + result["spyre_prefix_cache_hit_pct"] = [ + m.get("prefix_cache_hit_pct", 0.0) for m in metrics_list + ] try: with open(file_path, "w", encoding="utf-8") as fh: diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 26d953a10..bf0b6fe66 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -334,9 +334,13 @@ def _free_request(self, request, delay_free_blocks: bool = False): queued_time_s = ( (first_ts - arrival_ts) if first_ts is not None and arrival_ts is not None else 0.0 ) + num_executed = chunk_stats["num_chunked_prefills"] if chunk_stats else 0 + num_expected = math.ceil(request.num_prompt_tokens / self.chunk_size) + num_skipped = max(0, num_expected - num_executed) + cache_hit_pct = num_skipped / num_expected if num_expected > 0 else 0.0 spyre_data = { "queued_time_s": queued_time_s, - "num_chunked_prefills": chunk_stats["num_chunked_prefills"] if chunk_stats else 0, + "num_chunked_prefills": num_executed, "chunk_prefill_latencies_s": chunk_stats["chunk_prefill_latencies_s"] if chunk_stats else [], @@ -346,6 +350,7 @@ def _free_request(self, request, delay_free_blocks: bool = False): "decode_latencies_s": chunk_stats["decode_latencies_s"] if chunk_stats else [], "decode_start_times_s": chunk_stats["decode_start_times_s"] if chunk_stats else [], "tkvs": chunk_stats["tkvs"] if chunk_stats else [], + "prefix_cache_hit_pct": cache_hit_pct, } if kv_xfer_params is None: kv_xfer_params = {"__spyre__": spyre_data} diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index 230ec71ec..2cad89234 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -38,6 +38,7 @@ "decode_latencies_s": [88888.8, 0.000005, 44444.4], "decode_start_times_s": [5000.0, 5088888.8, 5088888.8], "tkvs": [512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072], + "prefix_cache_hit_pct": 0.25, }, { "queued_time_s": 0.00001, @@ -47,6 +48,7 @@ "decode_latencies_s": [0.000002, 77777.7], "decode_start_times_s": [1112345.5, 1112345.5], "tkvs": [256, 512, 1024, 2048, 4096], + "prefix_cache_hit_pct": 0.0, }, ] @@ -192,17 +194,19 @@ def test_print_spyre_section_output(capsys): out.assert_contains("Total prefill chunks processed:") out.assert_contains("10") - # Section separators and mean/median/percentile lines for all four sections + # Section separators and mean/median/percentile lines for all five sections out.assert_contains("Queue Wait Time") out.assert_contains("Chunked Prefill Count") out.assert_contains("Chunked Prefill Latency") out.assert_contains("Decode Step Latency") + out.assert_contains("Prefix Cache Hit") for label in ( "Queue Wait Time (ms)", "Num Chunked Prefills", "Chunk Prefill Latency (ms)", "Decode Step Latency (ms)", + "Prefix Cache Hit (%)", ): out.assert_contains(f"Mean {label}:") out.assert_contains(f"Median {label}:") From c4233b7d644298a4b59cda1b715056c457d04201 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 12 Jun 2026 22:35:43 +0200 Subject: [PATCH 060/106] display left-padding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_bench_serve.py | 4 ++++ sendnn_inference/v1/core/scheduler.py | 15 ++++++++++++++- tests/benchmarks/test_bench_metrics.py | 18 +++++++++++++++++- 3 files changed, 35 insertions(+), 2 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 71663370f..366c1b97e 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -164,11 +164,14 @@ def _section(header: str, values: list[float], label: str) -> None: m["prefix_cache_hit_pct"] * 100 for m in metrics_list if "prefix_cache_hit_pct" in m ] + left_padding_blocks = [v for m in metrics_list for v in m.get("left_padding_blocks", [])] + _section("Queue Wait Time", queue_times_ms, "Queue Wait Time (ms)") _section("Chunked Prefill Count", num_chunks_list, "Num Chunked Prefills") _section("Chunked Prefill Latency", chunk_lats_ms, "Chunk Prefill Latency (ms)") _section("Decode Step Latency", decode_lats_ms, "Decode Step Latency (ms)") _section("Prefix Cache Hit", cache_hit_pcts, "Prefix Cache Hit (%)") + _section("Left Padding Blocks", left_padding_blocks, "Left Padding Blocks") print("=" * 50) @@ -245,6 +248,7 @@ def _inject_spyre_metrics_into_result_file( result["spyre_prefix_cache_hit_pct"] = [ m.get("prefix_cache_hit_pct", 0.0) for m in metrics_list ] + result["spyre_left_padding_blocks"] = [m.get("left_padding_blocks", []) for m in metrics_list] try: with open(file_path, "w", encoding="utf-8") as fh: diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index bf0b6fe66..bd05076bb 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -35,6 +35,7 @@ class SpyreBenchState: decode_latencies: dict[str, list[float]] = field(default_factory=dict) decode_start_times: dict[str, list[float]] = field(default_factory=dict) tkvs: dict[str, list[int]] = field(default_factory=dict) + left_padding_blocks: dict[str, list[int]] = field(default_factory=dict) prefill_step_start: float | None = None decode_step_start: float | None = None @@ -280,10 +281,19 @@ def update_from_output(self, scheduler_output, model_runner_output): assert not self.previous_step_was_prefill and self._bench.prefill_step_start is None t0 = self._bench.decode_step_start duration = now - t0 + tkv = model_runner_output.tkv + max_num_blocks = math.ceil(tkv / self.block_size) + req_by_id = {r.request_id: r for r in self.running} for req_id in scheduler_output.scheduled_cached_reqs.req_ids: self._bench.decode_latencies.setdefault(req_id, []).append(duration) self._bench.decode_start_times.setdefault(req_id, []).append(t0) - self._bench.tkvs.setdefault(req_id, []).append(model_runner_output.tkv) + self._bench.tkvs.setdefault(req_id, []).append(tkv) + req = req_by_id.get(req_id) + if req is not None: + req_num_blocks = math.ceil(req.num_computed_tokens / self.block_size) + self._bench.left_padding_blocks.setdefault(req_id, []).append( + max_num_blocks - req_num_blocks + ) self._bench.prefill_step_start = None self._bench.decode_step_start = None @@ -310,6 +320,7 @@ def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: dec_lats = self._bench.decode_latencies.pop(req_id, None) dec_starts = self._bench.decode_start_times.pop(req_id, None) tkvs = self._bench.tkvs.pop(req_id, None) + left_padding_blocks = self._bench.left_padding_blocks.pop(req_id, None) if lats is None and dec_lats is None: return None return { @@ -319,6 +330,7 @@ def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: "decode_latencies_s": dec_lats or [], "decode_start_times_s": dec_starts or [], "tkvs": tkvs or [], + "left_padding_blocks": left_padding_blocks or [], } def _free_request(self, request, delay_free_blocks: bool = False): @@ -351,6 +363,7 @@ def _free_request(self, request, delay_free_blocks: bool = False): "decode_start_times_s": chunk_stats["decode_start_times_s"] if chunk_stats else [], "tkvs": chunk_stats["tkvs"] if chunk_stats else [], "prefix_cache_hit_pct": cache_hit_pct, + "left_padding_blocks": chunk_stats["left_padding_blocks"] if chunk_stats else [], } if kv_xfer_params is None: kv_xfer_params = {"__spyre__": spyre_data} diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index 2cad89234..9db743814 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -39,6 +39,7 @@ "decode_start_times_s": [5000.0, 5088888.8, 5088888.8], "tkvs": [512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072], "prefix_cache_hit_pct": 0.25, + "left_padding_blocks": [2, 0, 1], }, { "queued_time_s": 0.00001, @@ -49,6 +50,7 @@ "decode_start_times_s": [1112345.5, 1112345.5], "tkvs": [256, 512, 1024, 2048, 4096], "prefix_cache_hit_pct": 0.0, + "left_padding_blocks": [3, 1], }, ] @@ -194,12 +196,13 @@ def test_print_spyre_section_output(capsys): out.assert_contains("Total prefill chunks processed:") out.assert_contains("10") - # Section separators and mean/median/percentile lines for all five sections + # Section separators and mean/median/percentile lines for all six sections out.assert_contains("Queue Wait Time") out.assert_contains("Chunked Prefill Count") out.assert_contains("Chunked Prefill Latency") out.assert_contains("Decode Step Latency") out.assert_contains("Prefix Cache Hit") + out.assert_contains("Left Padding Blocks") for label in ( "Queue Wait Time (ms)", @@ -207,6 +210,7 @@ def test_print_spyre_section_output(capsys): "Chunk Prefill Latency (ms)", "Decode Step Latency (ms)", "Prefix Cache Hit (%)", + "Left Padding Blocks", ): out.assert_contains(f"Mean {label}:") out.assert_contains(f"Median {label}:") @@ -266,6 +270,7 @@ def _make_bare_scheduler(): "decode_latencies": [0.1, 0.2], "decode_start_times": [2000.0, 2000.1], "tkvs": [64, 128, 192, 256], + "left_padding_blocks": [2, 0], "prefill_step_start": 999.0, "decode_step_start": 1999.0, } @@ -279,6 +284,7 @@ def _make_bare_scheduler(): "decode_latencies_s": pytest.approx([0.1, 0.2]), "decode_start_times_s": pytest.approx([2000.0, 2000.1]), "tkvs": [64, 128, 192, 256], + "left_padding_blocks": [2, 0], } @@ -526,6 +532,16 @@ def _capturing_free(self, request, delay_free_blocks=False): assert info["arrival_ts"] is not None, f"req {req_id}: arrival_ts not set" assert info["first_scheduled_ts"] is not None, f"req {req_id}: first_scheduled_ts not set" + # left_padding_blocks: one entry per decode step, matching decode_latencies length + assert len(info["left_padding_blocks"]) == len(info["decode_latencies"]), ( + f"req {req_id}: left_padding_blocks length {len(info['left_padding_blocks'])} " + f"!= decode_latencies length {len(info['decode_latencies'])}" + ) + for blocks in info["left_padding_blocks"]: + assert isinstance(blocks, int) and blocks >= 0, ( + f"req {req_id}: negative left_padding_blocks value {blocks}" + ) + # The two requests must have independent latency lists (no cross-contamination) assert captured["0"]["chunk_latencies"] != captured["1"]["chunk_latencies"] or ( # Allow equal only if prompts produced identical timings by coincidence — From 182306f9ac5e83d8802cf0a6801acc66f5ee0b87 Mon Sep 17 00:00:00 2001 From: Max de Bayser Date: Fri, 12 Jun 2026 16:36:59 -0400 Subject: [PATCH 061/106] refactor pause/resume Signed-off-by: Max de Bayser --- .../v1/worker/spyre_input_batch.py | 65 ++++++++----------- .../v1/worker/spyre_model_runner.py | 2 +- 2 files changed, 27 insertions(+), 40 deletions(-) diff --git a/sendnn_inference/v1/worker/spyre_input_batch.py b/sendnn_inference/v1/worker/spyre_input_batch.py index 28c2f0b25..88615a241 100644 --- a/sendnn_inference/v1/worker/spyre_input_batch.py +++ b/sendnn_inference/v1/worker/spyre_input_batch.py @@ -372,13 +372,12 @@ def get_available_index(self) -> int | None: available_indices_list = available_indices.squeeze(dim=-1).tolist() return available_indices_list[0] if available_indices_list else None - def add_request( + def _setup_request_data( self, request: "SamplingRequestState", req_index: int | None = None, - ) -> int: + ) -> tuple[int, int]: req_index = super().add_request(request, req_index) - req_id = request.req_id # NOTE: differently from gpu input batch, self.req_output_token_ids # is not synced with self._req_ids, it should use @@ -389,6 +388,23 @@ def add_request( dense_index = self.req_idx_to_dense_index(req_index) self.req_output_token_ids.insert(dense_index, request.output_token_ids) + # Copy the output token ids. + start_idx = len(request.prompt_token_ids) + end_idx = start_idx + len(request.output_token_ids) + self.token_ids_cpu[req_index, start_idx:end_idx] = request.output_token_ids + + return req_index, dense_index + + def add_request( + self, + request: "SamplingRequestState", + req_index: int | None = None, + ) -> int: + req_index, dense_index = self._setup_request_data(request, req_index) + req_id = request.req_id + + self._register_sampling_params(req_id, req_index, request) + params = request.sampling_params # TODO add pooling params tmp_dense = self.num_reqs - 1 self.batch_update_builder.added.append( @@ -401,12 +417,6 @@ def add_request( ) tmp_dense = tmp_dense - 1 - # Copy the output token ids. - start_idx = len(request.prompt_token_ids) - end_idx = start_idx + len(request.output_token_ids) - self.token_ids_cpu[req_index, start_idx:end_idx] = request.output_token_ids - - self._register_sampling_params(req_id, req_index, request) return req_index def remove_request(self, req_id: str): @@ -460,17 +470,13 @@ def pause_request(self, req_id: str) -> None: - The slot is freed (cleared from _req_ids) so new requests can use it. """ # Pop from id map and clear the slot - req_index = self.req_id_to_index.pop(req_id, None) + req_index = super().remove_request(req_id) if req_index is None: return # Must compute dense_index before masking dense_index = self.req_idx_to_dense_index(req_index) - - # Free the slot for new requests - self._req_ids[req_index] = None self.req_indices_mask[req_index] = False - self._num_requests -= 1 # Tell LogitProcessorWrapper to save state at this dense position. self.batch_update_builder.pause_append(dense_index, req_id) @@ -483,10 +489,10 @@ def pause_request(self, req_id: str) -> None: (tmp_dense, tmp_dense + 1, MoveDirectionality.UNIDIRECTIONAL) ) - self.req_output_token_ids.pop(dense_index) - self._unregister_sampling_params(req_id, req_index) + self.req_output_token_ids.pop(dense_index) + self._unregister_sampling_params(req_id, req_index) - def resume_request(self, req_id: str, request: "SamplingRequestState") -> None: + def resume_request(self, request: "SamplingRequestState") -> None: """Restore a previously paused request to the active batch. Emits an 'added' event so builtin processors (MinP, LogitBias, …) @@ -494,27 +500,10 @@ def resume_request(self, req_id: str, request: "SamplingRequestState") -> None: event then tells LogitProcessorWrapper to overwrite that freshly initialised slot with the exact saved state, preserving history. """ - # Get an available slot (same as add_request) - req_index = self.get_available_index() - assert req_index is not None - assert req_index < self.max_num_reqs - - # Set up the slot - self._req_ids[req_index] = req_id - self.req_indices_mask[req_index] = True - self.req_id_to_index[req_id] = req_index - self._num_requests += 1 - - # Copy prompt and output token ids - num_prompt_tokens = len(request.prompt_token_ids) - self.num_prompt_tokens[req_index] = num_prompt_tokens - self.token_ids_cpu[req_index, :num_prompt_tokens] = request.prompt_token_ids - start_idx = num_prompt_tokens - end_idx = start_idx + len(request.output_token_ids) - self.token_ids_cpu[req_index, start_idx:end_idx] = request.output_token_ids + req_index, dense_index = self._setup_request_data(request) + req_id = request.req_id - dense_index = self.req_idx_to_dense_index(req_index) - self.req_output_token_ids.insert(dense_index, request.output_token_ids) + self._register_sampling_params(req_id, req_index, request) # Tell LogitProcessorWrapper to restore the saved state at dense_index. # No 'added' event needed - the saved LogitsProcessor instance already @@ -530,8 +519,6 @@ def resume_request(self, req_id: str, request: "SamplingRequestState") -> None: ) tmp_dense -= 1 - self._register_sampling_params(req_id, req_index, request) - def _register_sampling_params( self, req_id: str, req_index: int, request: "SamplingRequestState" ) -> None: diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index 07edb1a10..fce9a29b3 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -1465,7 +1465,7 @@ def _update_batch(self, scheduler_output: SchedulerOutput): # Only restore requests that were previously paused if req_id in self.paused_req_ids and req_id in self.requests: req_state = self.requests[req_id] - self.input_batch.resume_request(req_id, req_state) + self.input_batch.resume_request(req_state) self.paused_req_ids.discard(req_id) self.input_batch.refresh_metadata() need_metadata_refresh = False From e2f7ecf3e8efe0c31b8a61116ab85ea5ce2b171f Mon Sep 17 00:00:00 2001 From: Max de Bayser Date: Fri, 12 Jun 2026 16:39:10 -0400 Subject: [PATCH 062/106] remove cond2 Signed-off-by: Max de Bayser --- sendnn_inference/v1/core/scheduler.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 684db62bc..f9cccfb24 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -557,14 +557,7 @@ def _satisfies_last_chunk_constraints(self, request: Request) -> bool: # check that there is space in the current decode batch num_running = len(decoding_requests) - cond1 = num_running + len(self.waiting) < self.max_num_running_reqs - - # Check that the current decode batch is not about to have requests paused. - # This avoids adding more request to be paused and seems to slightly improve - # metrics. - cond2 = lambda: self._can_decode_all_requests(self.running) - - return cond1 and cond2() + return num_running + len(self.waiting) < self.max_num_running_reqs def _has_scheduling_priority(self, request): decoding_requests = [r for r in self.running if r not in self.ongoing_prefills] From 11f4051d81760723a3a4a40c8f7853484847e5c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 12 Jun 2026 22:03:51 +0200 Subject: [PATCH 063/106] display prefix cache hit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/benchmarks/spyre_bench_serve.py | 8 ++++++++ sendnn_inference/v1/core/scheduler.py | 7 ++++++- tests/benchmarks/test_bench_metrics.py | 6 +++++- 3 files changed, 19 insertions(+), 2 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index ee30e9e1a..71663370f 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -160,10 +160,15 @@ def _section(header: str, values: list[float], label: str) -> None: ) ) + cache_hit_pcts = [ + m["prefix_cache_hit_pct"] * 100 for m in metrics_list if "prefix_cache_hit_pct" in m + ] + _section("Queue Wait Time", queue_times_ms, "Queue Wait Time (ms)") _section("Chunked Prefill Count", num_chunks_list, "Num Chunked Prefills") _section("Chunked Prefill Latency", chunk_lats_ms, "Chunk Prefill Latency (ms)") _section("Decode Step Latency", decode_lats_ms, "Decode Step Latency (ms)") + _section("Prefix Cache Hit", cache_hit_pcts, "Prefix Cache Hit (%)") print("=" * 50) @@ -237,6 +242,9 @@ def _inject_spyre_metrics_into_result_file( result["spyre_decode_latencies_s"] = [m.get("decode_latencies_s", []) for m in metrics_list] result["spyre_decode_start_times_s"] = [m.get("decode_start_times_s", []) for m in metrics_list] result["spyre_tkvs"] = [m.get("tkvs", []) for m in metrics_list] + result["spyre_prefix_cache_hit_pct"] = [ + m.get("prefix_cache_hit_pct", 0.0) for m in metrics_list + ] try: with open(file_path, "w", encoding="utf-8") as fh: diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 26d953a10..bf0b6fe66 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -334,9 +334,13 @@ def _free_request(self, request, delay_free_blocks: bool = False): queued_time_s = ( (first_ts - arrival_ts) if first_ts is not None and arrival_ts is not None else 0.0 ) + num_executed = chunk_stats["num_chunked_prefills"] if chunk_stats else 0 + num_expected = math.ceil(request.num_prompt_tokens / self.chunk_size) + num_skipped = max(0, num_expected - num_executed) + cache_hit_pct = num_skipped / num_expected if num_expected > 0 else 0.0 spyre_data = { "queued_time_s": queued_time_s, - "num_chunked_prefills": chunk_stats["num_chunked_prefills"] if chunk_stats else 0, + "num_chunked_prefills": num_executed, "chunk_prefill_latencies_s": chunk_stats["chunk_prefill_latencies_s"] if chunk_stats else [], @@ -346,6 +350,7 @@ def _free_request(self, request, delay_free_blocks: bool = False): "decode_latencies_s": chunk_stats["decode_latencies_s"] if chunk_stats else [], "decode_start_times_s": chunk_stats["decode_start_times_s"] if chunk_stats else [], "tkvs": chunk_stats["tkvs"] if chunk_stats else [], + "prefix_cache_hit_pct": cache_hit_pct, } if kv_xfer_params is None: kv_xfer_params = {"__spyre__": spyre_data} diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index 230ec71ec..2cad89234 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -38,6 +38,7 @@ "decode_latencies_s": [88888.8, 0.000005, 44444.4], "decode_start_times_s": [5000.0, 5088888.8, 5088888.8], "tkvs": [512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072], + "prefix_cache_hit_pct": 0.25, }, { "queued_time_s": 0.00001, @@ -47,6 +48,7 @@ "decode_latencies_s": [0.000002, 77777.7], "decode_start_times_s": [1112345.5, 1112345.5], "tkvs": [256, 512, 1024, 2048, 4096], + "prefix_cache_hit_pct": 0.0, }, ] @@ -192,17 +194,19 @@ def test_print_spyre_section_output(capsys): out.assert_contains("Total prefill chunks processed:") out.assert_contains("10") - # Section separators and mean/median/percentile lines for all four sections + # Section separators and mean/median/percentile lines for all five sections out.assert_contains("Queue Wait Time") out.assert_contains("Chunked Prefill Count") out.assert_contains("Chunked Prefill Latency") out.assert_contains("Decode Step Latency") + out.assert_contains("Prefix Cache Hit") for label in ( "Queue Wait Time (ms)", "Num Chunked Prefills", "Chunk Prefill Latency (ms)", "Decode Step Latency (ms)", + "Prefix Cache Hit (%)", ): out.assert_contains(f"Mean {label}:") out.assert_contains(f"Median {label}:") From 49db7f8b0b95acd06e9d4f27744b59f2820e790c Mon Sep 17 00:00:00 2001 From: Max de Bayser Date: Mon, 15 Jun 2026 11:37:19 -0400 Subject: [PATCH 064/106] address PR review comments Signed-off-by: Max de Bayser --- sendnn_inference/v1/core/scheduler.py | 44 +++++++------------ .../v1/worker/spyre_input_batch.py | 2 +- 2 files changed, 16 insertions(+), 30 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index f9cccfb24..ea890c8a6 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -26,12 +26,6 @@ assert SpyrePlatform.get_block_size() == 64 -def round_up_to_block_size(n: int) -> int: - # Helper function to round up to the nearest block size - # Uses bitwise alignment for better performance - return (n + 63) & ~63 - - class SpyreScheduler(Scheduler): """Base class inheriting from the V1 scheduler to support static and continuous batching respecting AIU Spyre constraints.""" @@ -606,15 +600,12 @@ def _handle_decode_requests_pausing(self) -> None: """ decoding_requests = [r for r in self.running if r not in self.ongoing_prefills] - had_to_remove = False initial_had_requests = len(decoding_requests) > 0 # If we can't decode all requests due to batch TKV limits, iteratively # remove requests with the fewest decoded tokens and pause them until # the remaining batch fits within constraints while not self._can_decode_all_requests(decoding_requests): - had_to_remove = True - # TODO we should test different removal logics: longest request, optimize padding # Remove the request with the fewest decoded tokens # Decoded tokens = num_computed_tokens - num_prompt_tokens @@ -629,26 +620,21 @@ def _handle_decode_requests_pausing(self) -> None: # It shouldn't be possible to remove all requests if we started with some assert not initial_had_requests or len(decoding_requests) > 0 - # If we didn't have to remove any requests, try to add back previously - # paused requests (oldest first) as long as they fit within constraints - if not had_to_remove: - while self.paused_decoding_requests: - # Try adding the oldest paused request (first in list) - request_to_add = self.paused_decoding_requests[0] - test_requests = decoding_requests + [request_to_add] - - if self._can_decode_all_requests(test_requests): - # Can add this request back - self.paused_decoding_requests.pop(0) - self.running.append(request_to_add) - decoding_requests.append(request_to_add) - logger.info( - "Request %s resumed (batch TKV capacity available).", - request_to_add.request_id, - ) - else: - # Can't add any more requests - break + # Check if any paused request can be added back. + for i in range(len(self.paused_decoding_requests) - 1, -1, -1): + # Try adding the oldest paused request (first in list) + request_to_add = self.paused_decoding_requests[i] + test_requests = decoding_requests + [request_to_add] + + if self._can_decode_all_requests(test_requests): + # Can add this request back + self.paused_decoding_requests.pop(i) + self.running.append(request_to_add) + decoding_requests.append(request_to_add) + logger.info( + "Request %s resumed (batch TKV capacity available).", + request_to_add.request_id, + ) def predict_next_decode_tkv(self, running_requests: list[Request]) -> int: """ diff --git a/sendnn_inference/v1/worker/spyre_input_batch.py b/sendnn_inference/v1/worker/spyre_input_batch.py index 88615a241..e0c17d5eb 100644 --- a/sendnn_inference/v1/worker/spyre_input_batch.py +++ b/sendnn_inference/v1/worker/spyre_input_batch.py @@ -490,7 +490,7 @@ def pause_request(self, req_id: str) -> None: ) self.req_output_token_ids.pop(dense_index) - self._unregister_sampling_params(req_id, req_index) + self._unregister_sampling_params(req_id, req_index) def resume_request(self, request: "SamplingRequestState") -> None: """Restore a previously paused request to the active batch. From 6243d3598c453b1c52c3f019cd67f00ed64dd9ec Mon Sep 17 00:00:00 2001 From: Max de Bayser Date: Mon, 15 Jun 2026 14:01:46 -0400 Subject: [PATCH 065/106] address review comment Signed-off-by: Max de Bayser --- sendnn_inference/v1/sample/spyre_logits_processor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sendnn_inference/v1/sample/spyre_logits_processor.py b/sendnn_inference/v1/sample/spyre_logits_processor.py index 3b286232b..ea7f05e2f 100644 --- a/sendnn_inference/v1/sample/spyre_logits_processor.py +++ b/sendnn_inference/v1/sample/spyre_logits_processor.py @@ -167,7 +167,7 @@ def update_state(self, batch_update: BatchUpdate | None) -> None: # Max: I think we can't assume that the request will # be here because it could be a cancelled request that # never made it into the batch. - self._saved.pop(req_id, 0) + self._saved.pop(req_id, None) for adx, bdx, _ in batch_update.moved: update_called[adx], update_called[bdx] = update_called[bdx], update_called[adx] From a19dea29df36539b7967706b6823a078e3591c25 Mon Sep 17 00:00:00 2001 From: Max de Bayser Date: Mon, 15 Jun 2026 14:04:48 -0400 Subject: [PATCH 066/106] fix test Signed-off-by: Max de Bayser --- sendnn_inference/v1/worker/spyre_input_batch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sendnn_inference/v1/worker/spyre_input_batch.py b/sendnn_inference/v1/worker/spyre_input_batch.py index e0c17d5eb..e1c2834f8 100644 --- a/sendnn_inference/v1/worker/spyre_input_batch.py +++ b/sendnn_inference/v1/worker/spyre_input_batch.py @@ -489,7 +489,7 @@ def pause_request(self, req_id: str) -> None: (tmp_dense, tmp_dense + 1, MoveDirectionality.UNIDIRECTIONAL) ) - self.req_output_token_ids.pop(dense_index) + self.req_output_token_ids.pop(dense_index) self._unregister_sampling_params(req_id, req_index) def resume_request(self, request: "SamplingRequestState") -> None: From dcf245957d8a6b9d755abf0dddf35416b52c2aa8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Fri, 12 Jun 2026 22:35:43 +0200 Subject: [PATCH 067/106] display left-padding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_bench_serve.py | 6 +++++- sendnn_inference/v1/core/scheduler.py | 15 ++++++++++++++- tests/benchmarks/test_bench_metrics.py | 18 +++++++++++++++++- 3 files changed, 36 insertions(+), 3 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 71663370f..87bd8a66f 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -147,7 +147,7 @@ def _print_spyre_section( def _section(header: str, values: list[float], label: str) -> None: if not values: - return + values = [0.0] arr = np.array(values) print("{s:{c}^{n}}".format(s=f" {header} ", n=50, c="-")) print("{:<40} {:<10.2f}".format(f"Mean {label}:", float(np.mean(arr)))) @@ -164,11 +164,14 @@ def _section(header: str, values: list[float], label: str) -> None: m["prefix_cache_hit_pct"] * 100 for m in metrics_list if "prefix_cache_hit_pct" in m ] + left_padding_blocks = [v for m in metrics_list for v in m.get("left_padding_blocks", [])] + _section("Queue Wait Time", queue_times_ms, "Queue Wait Time (ms)") _section("Chunked Prefill Count", num_chunks_list, "Num Chunked Prefills") _section("Chunked Prefill Latency", chunk_lats_ms, "Chunk Prefill Latency (ms)") _section("Decode Step Latency", decode_lats_ms, "Decode Step Latency (ms)") _section("Prefix Cache Hit", cache_hit_pcts, "Prefix Cache Hit (%)") + _section("Left Padding Blocks", left_padding_blocks, "Left Padding Blocks") print("=" * 50) @@ -245,6 +248,7 @@ def _inject_spyre_metrics_into_result_file( result["spyre_prefix_cache_hit_pct"] = [ m.get("prefix_cache_hit_pct", 0.0) for m in metrics_list ] + result["spyre_left_padding_blocks"] = [m.get("left_padding_blocks", []) for m in metrics_list] try: with open(file_path, "w", encoding="utf-8") as fh: diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index bf0b6fe66..bd05076bb 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -35,6 +35,7 @@ class SpyreBenchState: decode_latencies: dict[str, list[float]] = field(default_factory=dict) decode_start_times: dict[str, list[float]] = field(default_factory=dict) tkvs: dict[str, list[int]] = field(default_factory=dict) + left_padding_blocks: dict[str, list[int]] = field(default_factory=dict) prefill_step_start: float | None = None decode_step_start: float | None = None @@ -280,10 +281,19 @@ def update_from_output(self, scheduler_output, model_runner_output): assert not self.previous_step_was_prefill and self._bench.prefill_step_start is None t0 = self._bench.decode_step_start duration = now - t0 + tkv = model_runner_output.tkv + max_num_blocks = math.ceil(tkv / self.block_size) + req_by_id = {r.request_id: r for r in self.running} for req_id in scheduler_output.scheduled_cached_reqs.req_ids: self._bench.decode_latencies.setdefault(req_id, []).append(duration) self._bench.decode_start_times.setdefault(req_id, []).append(t0) - self._bench.tkvs.setdefault(req_id, []).append(model_runner_output.tkv) + self._bench.tkvs.setdefault(req_id, []).append(tkv) + req = req_by_id.get(req_id) + if req is not None: + req_num_blocks = math.ceil(req.num_computed_tokens / self.block_size) + self._bench.left_padding_blocks.setdefault(req_id, []).append( + max_num_blocks - req_num_blocks + ) self._bench.prefill_step_start = None self._bench.decode_step_start = None @@ -310,6 +320,7 @@ def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: dec_lats = self._bench.decode_latencies.pop(req_id, None) dec_starts = self._bench.decode_start_times.pop(req_id, None) tkvs = self._bench.tkvs.pop(req_id, None) + left_padding_blocks = self._bench.left_padding_blocks.pop(req_id, None) if lats is None and dec_lats is None: return None return { @@ -319,6 +330,7 @@ def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: "decode_latencies_s": dec_lats or [], "decode_start_times_s": dec_starts or [], "tkvs": tkvs or [], + "left_padding_blocks": left_padding_blocks or [], } def _free_request(self, request, delay_free_blocks: bool = False): @@ -351,6 +363,7 @@ def _free_request(self, request, delay_free_blocks: bool = False): "decode_start_times_s": chunk_stats["decode_start_times_s"] if chunk_stats else [], "tkvs": chunk_stats["tkvs"] if chunk_stats else [], "prefix_cache_hit_pct": cache_hit_pct, + "left_padding_blocks": chunk_stats["left_padding_blocks"] if chunk_stats else [], } if kv_xfer_params is None: kv_xfer_params = {"__spyre__": spyre_data} diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index 2cad89234..9db743814 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -39,6 +39,7 @@ "decode_start_times_s": [5000.0, 5088888.8, 5088888.8], "tkvs": [512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072], "prefix_cache_hit_pct": 0.25, + "left_padding_blocks": [2, 0, 1], }, { "queued_time_s": 0.00001, @@ -49,6 +50,7 @@ "decode_start_times_s": [1112345.5, 1112345.5], "tkvs": [256, 512, 1024, 2048, 4096], "prefix_cache_hit_pct": 0.0, + "left_padding_blocks": [3, 1], }, ] @@ -194,12 +196,13 @@ def test_print_spyre_section_output(capsys): out.assert_contains("Total prefill chunks processed:") out.assert_contains("10") - # Section separators and mean/median/percentile lines for all five sections + # Section separators and mean/median/percentile lines for all six sections out.assert_contains("Queue Wait Time") out.assert_contains("Chunked Prefill Count") out.assert_contains("Chunked Prefill Latency") out.assert_contains("Decode Step Latency") out.assert_contains("Prefix Cache Hit") + out.assert_contains("Left Padding Blocks") for label in ( "Queue Wait Time (ms)", @@ -207,6 +210,7 @@ def test_print_spyre_section_output(capsys): "Chunk Prefill Latency (ms)", "Decode Step Latency (ms)", "Prefix Cache Hit (%)", + "Left Padding Blocks", ): out.assert_contains(f"Mean {label}:") out.assert_contains(f"Median {label}:") @@ -266,6 +270,7 @@ def _make_bare_scheduler(): "decode_latencies": [0.1, 0.2], "decode_start_times": [2000.0, 2000.1], "tkvs": [64, 128, 192, 256], + "left_padding_blocks": [2, 0], "prefill_step_start": 999.0, "decode_step_start": 1999.0, } @@ -279,6 +284,7 @@ def _make_bare_scheduler(): "decode_latencies_s": pytest.approx([0.1, 0.2]), "decode_start_times_s": pytest.approx([2000.0, 2000.1]), "tkvs": [64, 128, 192, 256], + "left_padding_blocks": [2, 0], } @@ -526,6 +532,16 @@ def _capturing_free(self, request, delay_free_blocks=False): assert info["arrival_ts"] is not None, f"req {req_id}: arrival_ts not set" assert info["first_scheduled_ts"] is not None, f"req {req_id}: first_scheduled_ts not set" + # left_padding_blocks: one entry per decode step, matching decode_latencies length + assert len(info["left_padding_blocks"]) == len(info["decode_latencies"]), ( + f"req {req_id}: left_padding_blocks length {len(info['left_padding_blocks'])} " + f"!= decode_latencies length {len(info['decode_latencies'])}" + ) + for blocks in info["left_padding_blocks"]: + assert isinstance(blocks, int) and blocks >= 0, ( + f"req {req_id}: negative left_padding_blocks value {blocks}" + ) + # The two requests must have independent latency lists (no cross-contamination) assert captured["0"]["chunk_latencies"] != captured["1"]["chunk_latencies"] or ( # Allow equal only if prompts produced identical timings by coincidence — From 224c9aaaa8dd653a866bb3617eade38b38a30791 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 16 Jun 2026 11:44:28 +0200 Subject: [PATCH 068/106] don't prioritize any endpoint so that we support --skip-chat-template MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_bench_serve.py | 4 +- .../benchmarks/spyre_request_func.py | 57 ++++++++++++------- 2 files changed, 36 insertions(+), 25 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 87bd8a66f..a0e88060f 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -316,10 +316,8 @@ def main() -> None: parser = _build_parser() args = parser.parse_args(argv) - # Force chat endpoint and our backend. + # Force our custom backend so Spyre metrics are always collected. args.backend = _BACKEND_NAME - if not hasattr(args, "endpoint") or args.endpoint == "/v1/completions": - args.endpoint = "/v1/chat/completions" selected_percentiles = [float(p) for p in args.metric_percentiles.split(",")] diff --git a/sendnn_inference/benchmarks/spyre_request_func.py b/sendnn_inference/benchmarks/spyre_request_func.py index 519f9a9c8..518eb481c 100644 --- a/sendnn_inference/benchmarks/spyre_request_func.py +++ b/sendnn_inference/benchmarks/spyre_request_func.py @@ -37,29 +37,37 @@ async def async_request_spyre_chat( pbar: tqdm | None = None, mm_position: Literal["first", "last"] = "last", ) -> SpyreRequestFuncOutput: - """Chat completions request function that additionally parses the - ``spyre_metrics`` field injected into the final SSE usage chunk.""" + """Chat completions (or raw completions) request function that additionally + parses the ``spyre_metrics`` field injected into the final SSE usage chunk. + + When ``request_func_input.api_url`` targets ``/v1/completions``the prompt + is sent verbatim without wrapping it in a chat message, so the server-side + chat template is not applied a second time.""" api_url = request_func_input.api_url - _validate_api_url(api_url, "OpenAI Chat Completions API", "chat/completions") - - content = _get_chat_content(request_func_input, mm_position=mm_position) - - payload = { - "model": ( - request_func_input.model_name - if request_func_input.model_name - else request_func_input.model - ), - "messages": [ - {"role": "user", "content": content}, - ], - "max_completion_tokens": request_func_input.output_len, - "stream": True, - "stream_options": { - "include_usage": True, - }, - } + use_completions = api_url.endswith("/v1/completions") + + model = request_func_input.model_name or request_func_input.model + + if use_completions: + _validate_api_url(api_url, "OpenAI Completions API", "completions") + payload: dict[str, Any] = { + "model": model, + "prompt": request_func_input.prompt, + "max_tokens": request_func_input.output_len, + "stream": True, + "stream_options": {"include_usage": True}, + } + else: + _validate_api_url(api_url, "OpenAI Chat Completions API", "chat/completions") + content = _get_chat_content(request_func_input, mm_position=mm_position) + payload = { + "model": model, + "messages": [{"role": "user", "content": content}], + "max_completion_tokens": request_func_input.output_len, + "stream": True, + "stream_options": {"include_usage": True}, + } _update_payload_common(payload, request_func_input) headers = _get_headers("application/json") @@ -94,7 +102,12 @@ async def async_request_spyre_chat( data = json.loads(chunk) if choices := data.get("choices"): - content_delta = choices[0]["delta"].get("content") + # Chat completions uses delta.content; raw + # completions uses text directly on the choice. + if use_completions: + content_delta = choices[0].get("text") + else: + content_delta = choices[0]["delta"].get("content") if ttft == 0.0: ttft = timestamp - st output.ttft = ttft From 0cac7812e5753a55e92f97b381d98a02e2d5a8f9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Sat, 13 Jun 2026 00:24:19 +0200 Subject: [PATCH 069/106] display pause time MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_bench_serve.py | 8 +++++++ sendnn_inference/v1/core/scheduler.py | 20 ++++++++++++++++ tests/benchmarks/test_bench_metrics.py | 23 ++++++++++++++++++- 3 files changed, 50 insertions(+), 1 deletion(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 366c1b97e..93d755826 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -165,6 +165,9 @@ def _section(header: str, values: list[float], label: str) -> None: ] left_padding_blocks = [v for m in metrics_list for v in m.get("left_padding_blocks", [])] + pause_lats_ms = [lat * 1000 for m in metrics_list for lat in m.get("pause_latencies_s", [])] + pause_counts = [float(len(m.get("pause_latencies_s", []))) for m in metrics_list] + total_pause_ms = [float(sum(m.get("pause_latencies_s", []))) * 1000 for m in metrics_list] _section("Queue Wait Time", queue_times_ms, "Queue Wait Time (ms)") _section("Chunked Prefill Count", num_chunks_list, "Num Chunked Prefills") @@ -172,6 +175,9 @@ def _section(header: str, values: list[float], label: str) -> None: _section("Decode Step Latency", decode_lats_ms, "Decode Step Latency (ms)") _section("Prefix Cache Hit", cache_hit_pcts, "Prefix Cache Hit (%)") _section("Left Padding Blocks", left_padding_blocks, "Left Padding Blocks") + _section("Pause Latency", pause_lats_ms, "Pause Latency (ms)") + _section("Number of Pauses", pause_counts, "Num Pauses") + _section("Total Time Paused", total_pause_ms, "Total Time Paused (ms)") print("=" * 50) @@ -249,6 +255,8 @@ def _inject_spyre_metrics_into_result_file( m.get("prefix_cache_hit_pct", 0.0) for m in metrics_list ] result["spyre_left_padding_blocks"] = [m.get("left_padding_blocks", []) for m in metrics_list] + result["spyre_pause_latencies_s"] = [m.get("pause_latencies_s", []) for m in metrics_list] + result["spyre_pause_start_times_s"] = [m.get("pause_start_times_s", []) for m in metrics_list] try: with open(file_path, "w", encoding="utf-8") as fh: diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index d94b8e166..652112fbb 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -37,6 +37,8 @@ class SpyreBenchState: decode_start_times: dict[str, list[float]] = field(default_factory=dict) tkvs: dict[str, list[int]] = field(default_factory=dict) left_padding_blocks: dict[str, list[int]] = field(default_factory=dict) + pause_start_times: dict[str, list[float]] = field(default_factory=dict) + pause_latencies: dict[str, list[float]] = field(default_factory=dict) prefill_step_start: float | None = None decode_step_start: float | None = None @@ -334,6 +336,8 @@ def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: dec_starts = self._bench.decode_start_times.pop(req_id, None) tkvs = self._bench.tkvs.pop(req_id, None) left_padding_blocks = self._bench.left_padding_blocks.pop(req_id, None) + pause_lats = self._bench.pause_latencies.pop(req_id, None) + pause_starts = self._bench.pause_start_times.pop(req_id, None) if lats is None and dec_lats is None: return None return { @@ -344,6 +348,8 @@ def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: "decode_start_times_s": dec_starts or [], "tkvs": tkvs or [], "left_padding_blocks": left_padding_blocks or [], + "pause_latencies_s": pause_lats or [], + "pause_start_times_s": pause_starts or [], } def _free_request(self, request, delay_free_blocks: bool = False): @@ -377,6 +383,8 @@ def _free_request(self, request, delay_free_blocks: bool = False): "tkvs": chunk_stats["tkvs"] if chunk_stats else [], "prefix_cache_hit_pct": cache_hit_pct, "left_padding_blocks": chunk_stats["left_padding_blocks"] if chunk_stats else [], + "pause_latencies_s": chunk_stats["pause_latencies_s"] if chunk_stats else [], + "pause_start_times_s": chunk_stats["pause_start_times_s"] if chunk_stats else [], } if kv_xfer_params is None: kv_xfer_params = {"__spyre__": spyre_data} @@ -770,6 +778,11 @@ def _handle_decode_requests_pausing(self) -> None: self.running.remove(request_to_remove) self.paused_decoding_requests.append(request_to_remove) logger.info("Request %s paused due to batch TKV limit ", request_to_remove.request_id) + if self._bench is not None: + pause_ts = time.time() + self._bench.pause_start_times.setdefault(request_to_remove.request_id, []).append( + pause_ts + ) # It shouldn't be possible to remove all requests if we started with some assert not initial_had_requests or len(decoding_requests) > 0 @@ -791,6 +804,13 @@ def _handle_decode_requests_pausing(self) -> None: "Request %s resumed (batch TKV capacity available).", request_to_add.request_id, ) + if self._bench is not None: + starts = self._bench.pause_start_times.get(request_to_add.request_id) + if starts: + duration = time.time() - starts[-1] + self._bench.pause_latencies.setdefault( + request_to_add.request_id, [] + ).append(duration) else: # Can't add any more requests break diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index 9db743814..78f7577ec 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -40,6 +40,8 @@ "tkvs": [512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072], "prefix_cache_hit_pct": 0.25, "left_padding_blocks": [2, 0, 1], + "pause_latencies_s": [0.5, 1.2], + "pause_start_times_s": [0.0, 1.2], }, { "queued_time_s": 0.00001, @@ -51,6 +53,8 @@ "tkvs": [256, 512, 1024, 2048, 4096], "prefix_cache_hit_pct": 0.0, "left_padding_blocks": [3, 1], + "pause_latencies_s": [0.3], + "pause_start_times_s": [0.5], }, ] @@ -196,13 +200,16 @@ def test_print_spyre_section_output(capsys): out.assert_contains("Total prefill chunks processed:") out.assert_contains("10") - # Section separators and mean/median/percentile lines for all six sections + # Section separators and mean/median/percentile lines for all sections out.assert_contains("Queue Wait Time") out.assert_contains("Chunked Prefill Count") out.assert_contains("Chunked Prefill Latency") out.assert_contains("Decode Step Latency") out.assert_contains("Prefix Cache Hit") out.assert_contains("Left Padding Blocks") + out.assert_contains("Pause Latency") + out.assert_contains("Number of Pauses") + out.assert_contains("Total Time Paused") for label in ( "Queue Wait Time (ms)", @@ -211,6 +218,9 @@ def test_print_spyre_section_output(capsys): "Decode Step Latency (ms)", "Prefix Cache Hit (%)", "Left Padding Blocks", + "Pause Latency (ms)", + "Num Pauses", + "Total Time Paused (ms)", ): out.assert_contains(f"Mean {label}:") out.assert_contains(f"Median {label}:") @@ -271,6 +281,8 @@ def _make_bare_scheduler(): "decode_start_times": [2000.0, 2000.1], "tkvs": [64, 128, 192, 256], "left_padding_blocks": [2, 0], + "pause_start_times": [3000.0, 3005.0], + "pause_latencies": [0.5, 1.2], "prefill_step_start": 999.0, "decode_step_start": 1999.0, } @@ -285,6 +297,7 @@ def _make_bare_scheduler(): "decode_start_times_s": pytest.approx([2000.0, 2000.1]), "tkvs": [64, 128, 192, 256], "left_padding_blocks": [2, 0], + "pause_latencies_s": pytest.approx([0.5, 1.2]), } @@ -532,6 +545,14 @@ def _capturing_free(self, request, delay_free_blocks=False): assert info["arrival_ts"] is not None, f"req {req_id}: arrival_ts not set" assert info["first_scheduled_ts"] is not None, f"req {req_id}: first_scheduled_ts not set" + # pause_latencies: list of pause durations (may be absent/None if no pausing occurred) + pause_lats = info["pause_latencies"] or [] + assert isinstance(pause_lats, list), f"req {req_id}: pause_latencies is not a list" + for lat in pause_lats: + assert isinstance(lat, float) and lat > 0, ( + f"req {req_id}: non-positive pause_latency {lat}" + ) + # left_padding_blocks: one entry per decode step, matching decode_latencies length assert len(info["left_padding_blocks"]) == len(info["decode_latencies"]), ( f"req {req_id}: left_padding_blocks length {len(info['left_padding_blocks'])} " From 28d44566350eb9d0143d98a8922aa9a7a736dfd4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 16 Jun 2026 13:48:03 +0200 Subject: [PATCH 070/106] rename 'previous_step_was_prefill' to 'step_is_prefill' MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index ea890c8a6..29c809b67 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -192,7 +192,7 @@ def __init__(self, *args, **kwargs) -> None: # are interleaved with a decode step. This allows to minimize currently # decoding requests self.do_interleaving: bool = envs_spyre.SENDNN_INFERENCE_CP_INTERLEAVE_STEPS - self.previous_step_was_prefill: bool = False + self.step_is_prefill: bool = False self.tkv = 0 self.block_size = SpyrePlatform.get_block_size() @@ -379,11 +379,11 @@ def schedule(self) -> "SchedulerOutput": if schedule_prefill: running_holdback = [r for r in self.running if r not in self.ongoing_prefills] self.running = self.ongoing_prefills - self.previous_step_was_prefill = True + self.step_is_prefill = True else: self.running = [r for r in self.running if r not in self.ongoing_prefills] running_holdback = self.ongoing_prefills - self.previous_step_was_prefill = False + self.step_is_prefill = False # Check new requests to prefill elif len(self.waiting) > 0: @@ -405,7 +405,7 @@ def schedule(self) -> "SchedulerOutput": # Hide current decodes from the scheduler running_holdback = self.running self.running = [] - self.previous_step_was_prefill = True + self.step_is_prefill = True else: # Grammar not yet initialized for any waiting request. # Return them to holdback so the base scheduler doesn't @@ -413,12 +413,12 @@ def schedule(self) -> "SchedulerOutput": while self.waiting: holdback_queue.appendleft(self.waiting.pop()) running_holdback = [] - self.previous_step_was_prefill = False + self.step_is_prefill = False else: - self.previous_step_was_prefill = False + self.step_is_prefill = False running_holdback = [] - if not self.previous_step_was_prefill: + if not self.step_is_prefill: self._handle_decode_requests_pausing() # delegate to super of SpyreScheduler: base V1 Scheduler @@ -558,7 +558,7 @@ def _has_scheduling_priority(self, request): # If we do interleaving, then two consecutive prefill steps are # forbidden when there are decoding requests - if self.do_interleaving and self.previous_step_was_prefill and len(decoding_requests) > 0: + if self.do_interleaving and self.step_is_prefill and len(decoding_requests) > 0: return False # Requests that are already prefilling are prioritized over new requests From d1dbe75abcc892c6b7eab43a786ff23eb61fc2be Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 16 Jun 2026 14:11:30 +0200 Subject: [PATCH 071/106] address pause/resume logic comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 36 +++++++++++++++++++++------ 1 file changed, 28 insertions(+), 8 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 29c809b67..ef98d8bbe 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -599,15 +599,22 @@ def _handle_decode_requests_pausing(self) -> None: 2. Resumes previously paused requests (oldest first) when capacity is available """ decoding_requests = [r for r in self.running if r not in self.ongoing_prefills] + resumed = self._maybe_resume_decoding_requests(decoding_requests) + if not resumed: + self._maybe_pause_decoding_requests(decoding_requests) + def _maybe_pause_decoding_requests(self, decoding_requests: list[Request]) -> int: + """ + Iteratively pauses requests with the fewest decoded tokens until the batch fits + within TKV constraints. Mutates both decoding_requests and self.running. + + Returns the number of requests paused. + """ initial_had_requests = len(decoding_requests) > 0 + num_paused = 0 - # If we can't decode all requests due to batch TKV limits, iteratively - # remove requests with the fewest decoded tokens and pause them until - # the remaining batch fits within constraints + # TODO we should test different removal logics: longest request, optimize padding while not self._can_decode_all_requests(decoding_requests): - # TODO we should test different removal logics: longest request, optimize padding - # Remove the request with the fewest decoded tokens # Decoded tokens = num_computed_tokens - num_prompt_tokens request_to_remove = min( decoding_requests, key=lambda r: r.num_computed_tokens - r.num_prompt_tokens @@ -616,18 +623,28 @@ def _handle_decode_requests_pausing(self) -> None: self.running.remove(request_to_remove) self.paused_decoding_requests.append(request_to_remove) logger.info("Request %s paused due to batch TKV limit ", request_to_remove.request_id) + num_paused += 1 # It shouldn't be possible to remove all requests if we started with some assert not initial_had_requests or len(decoding_requests) > 0 - # Check if any paused request can be added back. + return num_paused + + def _maybe_resume_decoding_requests(self, decoding_requests: list[Request]) -> int: + """ + Resumes previously paused requests (oldest first) when TKV capacity is available. + Mutates both decoding_requests and self.running. + + Returns the number of requests resumed. + """ + num_resumed = 0 + + # Reverse iteration: pop(i) only shifts indices above i, which are already visited. for i in range(len(self.paused_decoding_requests) - 1, -1, -1): - # Try adding the oldest paused request (first in list) request_to_add = self.paused_decoding_requests[i] test_requests = decoding_requests + [request_to_add] if self._can_decode_all_requests(test_requests): - # Can add this request back self.paused_decoding_requests.pop(i) self.running.append(request_to_add) decoding_requests.append(request_to_add) @@ -635,6 +652,9 @@ def _handle_decode_requests_pausing(self) -> None: "Request %s resumed (batch TKV capacity available).", request_to_add.request_id, ) + num_resumed += 1 + + return num_resumed def predict_next_decode_tkv(self, running_requests: list[Request]) -> int: """ From f83d4f4ca318978978386bf3830fa8438e18d71f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 16 Jun 2026 11:44:28 +0200 Subject: [PATCH 072/106] support multi endpoints MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_bench_serve.py | 4 +- .../benchmarks/spyre_request_func.py | 57 ++++++++----- sendnn_inference/v1/metrics/patch_serving.py | 83 ++++++++++++++----- tests/benchmarks/test_bench_metrics.py | 10 ++- 4 files changed, 105 insertions(+), 49 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 87bd8a66f..a0e88060f 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -316,10 +316,8 @@ def main() -> None: parser = _build_parser() args = parser.parse_args(argv) - # Force chat endpoint and our backend. + # Force our custom backend so Spyre metrics are always collected. args.backend = _BACKEND_NAME - if not hasattr(args, "endpoint") or args.endpoint == "/v1/completions": - args.endpoint = "/v1/chat/completions" selected_percentiles = [float(p) for p in args.metric_percentiles.split(",")] diff --git a/sendnn_inference/benchmarks/spyre_request_func.py b/sendnn_inference/benchmarks/spyre_request_func.py index 519f9a9c8..518eb481c 100644 --- a/sendnn_inference/benchmarks/spyre_request_func.py +++ b/sendnn_inference/benchmarks/spyre_request_func.py @@ -37,29 +37,37 @@ async def async_request_spyre_chat( pbar: tqdm | None = None, mm_position: Literal["first", "last"] = "last", ) -> SpyreRequestFuncOutput: - """Chat completions request function that additionally parses the - ``spyre_metrics`` field injected into the final SSE usage chunk.""" + """Chat completions (or raw completions) request function that additionally + parses the ``spyre_metrics`` field injected into the final SSE usage chunk. + + When ``request_func_input.api_url`` targets ``/v1/completions``the prompt + is sent verbatim without wrapping it in a chat message, so the server-side + chat template is not applied a second time.""" api_url = request_func_input.api_url - _validate_api_url(api_url, "OpenAI Chat Completions API", "chat/completions") - - content = _get_chat_content(request_func_input, mm_position=mm_position) - - payload = { - "model": ( - request_func_input.model_name - if request_func_input.model_name - else request_func_input.model - ), - "messages": [ - {"role": "user", "content": content}, - ], - "max_completion_tokens": request_func_input.output_len, - "stream": True, - "stream_options": { - "include_usage": True, - }, - } + use_completions = api_url.endswith("/v1/completions") + + model = request_func_input.model_name or request_func_input.model + + if use_completions: + _validate_api_url(api_url, "OpenAI Completions API", "completions") + payload: dict[str, Any] = { + "model": model, + "prompt": request_func_input.prompt, + "max_tokens": request_func_input.output_len, + "stream": True, + "stream_options": {"include_usage": True}, + } + else: + _validate_api_url(api_url, "OpenAI Chat Completions API", "chat/completions") + content = _get_chat_content(request_func_input, mm_position=mm_position) + payload = { + "model": model, + "messages": [{"role": "user", "content": content}], + "max_completion_tokens": request_func_input.output_len, + "stream": True, + "stream_options": {"include_usage": True}, + } _update_payload_common(payload, request_func_input) headers = _get_headers("application/json") @@ -94,7 +102,12 @@ async def async_request_spyre_chat( data = json.loads(chunk) if choices := data.get("choices"): - content_delta = choices[0]["delta"].get("content") + # Chat completions uses delta.content; raw + # completions uses text directly on the choice. + if use_completions: + content_delta = choices[0].get("text") + else: + content_delta = choices[0]["delta"].get("content") if ttft == 0.0: ttft = timestamp - st output.ttft = ttft diff --git a/sendnn_inference/v1/metrics/patch_serving.py b/sendnn_inference/v1/metrics/patch_serving.py index 080d37bde..098b4af98 100644 --- a/sendnn_inference/v1/metrics/patch_serving.py +++ b/sendnn_inference/v1/metrics/patch_serving.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -"""Patch OpenAIServingChat to inject Spyre per-request metrics into the final -SSE usage chunk when SENDNN_INFERENCE_BENCH_METRICS_ENABLED is set. +"""Patch OpenAIServingChat and OpenAIServingCompletion to inject Spyre +per-request metrics into the final SSE usage chunk when +SENDNN_INFERENCE_BENCH_METRICS_ENABLED is set. Metrics are carried from the engine process to the API server process via RequestOutput.kv_transfer_params["__spyre__"], which already travels over the -ZMQ IPC channel. The patched generator intercepts the result_generator to -capture the final RequestOutput, then injects the metrics into the final SSE +ZMQ IPC channel. The patched generators intercept result_generator to +capture the final RequestOutput, then inject the metrics into the final SSE usage chunk before yielding it. """ @@ -18,23 +19,43 @@ _patched = False +def _inject_into_chunk(chunk: str, spyre_metrics: dict, request_id: str) -> str: + """Return chunk with spyre_metrics injected into the usage SSE payload, + or the original chunk if injection fails.""" + try: + prefix = "data: " + data_str = chunk.removeprefix(prefix).rstrip("\n") + data = json.loads(data_str) + data["spyre_metrics"] = spyre_metrics + return f"{prefix}{json.dumps(data)}\n\n" + except (json.JSONDecodeError, KeyError, TypeError) as e: + logger.warning("Failed to inject spyre_metrics into SSE chunk for %s: %s", request_id, e) + return chunk + + def patch_serving() -> None: - """Wrap OpenAIServingChat.chat_completion_stream_generator to inject - spyre_metrics into the final SSE usage chunk. Idempotent.""" + """Wrap chat and completions stream generators to inject spyre_metrics into + the final SSE usage chunk. Idempotent.""" global _patched if _patched: return + _patch_chat() + _patch_completions() + _patched = True + logger.debug("Spyre serving patch applied: spyre_metrics will be injected in final SSE chunk") + + +def _patch_chat() -> None: try: from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat except ImportError: - logger.warning("Could not import OpenAIServingChat — serving patch skipped") + logger.warning("Could not import OpenAIServingChat — chat serving patch skipped") return _original = OpenAIServingChat.chat_completion_stream_generator async def _patched_generator(self, request, result_generator, request_id, *args, **kwargs): - # Wrap result_generator to capture the final RequestOutput's kv_transfer_params. spyre_metrics: dict | None = None async def _capturing_generator(): @@ -53,18 +74,40 @@ async def _capturing_generator(): and '"usage"' in chunk and '"choices":[]' in chunk ): - try: - prefix = "data: " - data_str = chunk.removeprefix(prefix).rstrip("\n") - data = json.loads(data_str) - data["spyre_metrics"] = spyre_metrics - chunk = f"{prefix}{json.dumps(data)}\n\n" - except (json.JSONDecodeError, KeyError, TypeError) as e: - logger.warning( - "Failed to inject spyre_metrics into SSE chunk for %s: %s", request_id, e - ) + chunk = _inject_into_chunk(chunk, spyre_metrics, request_id) yield chunk OpenAIServingChat.chat_completion_stream_generator = _patched_generator # ty: ignore[invalid-assignment] - _patched = True - logger.debug("Spyre serving patch applied: spyre_metrics will be injected in final SSE chunk") + + +def _patch_completions() -> None: + try: + from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion + except ImportError: + logger.warning( + "Could not import OpenAIServingCompletion — completions serving patch skipped" + ) + return + + _original = OpenAIServingCompletion.completion_stream_generator + + async def _patched_generator( + self, request, engine_inputs, result_generator, request_id, *args, **kwargs + ): + spyre_metrics: dict | None = None + + async def _capturing_generator(): + nonlocal spyre_metrics + async for prompt_idx, res in result_generator: + if res.finished and res.kv_transfer_params: + spyre_metrics = res.kv_transfer_params.get("__spyre__") + yield prompt_idx, res + + async for chunk in _original( + self, request, engine_inputs, _capturing_generator(), request_id, *args, **kwargs + ): + if spyre_metrics is not None and isinstance(chunk, str) and '"usage"' in chunk: + chunk = _inject_into_chunk(chunk, spyre_metrics, request_id) + yield chunk + + OpenAIServingCompletion.completion_stream_generator = _patched_generator # ty: ignore[invalid-assignment] diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index 9db743814..d7c87293f 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -228,9 +228,10 @@ def test_print_noop_when_empty(capsys): @pytest.mark.cpu -def test_print_missing_keys_tolerated(capsys): +def test_print_missing_keys_show_zeros(capsys): # Metrics without chunk_prefill_latencies_s or decode_latencies_s — those - # sections should be absent, but queue time and chunk count should still print. + # sections still print with 0.00 values as a fallback, alongside the sections + # that do have data. metrics = [ {"queued_time_s": 77777.7, "num_chunked_prefills": 13}, {"queued_time_s": 0.000003, "num_chunked_prefills": 99}, @@ -239,8 +240,9 @@ def test_print_missing_keys_tolerated(capsys): out = capsys.readouterr().out assert "Queue Wait Time" in out assert "Chunked Prefill Count" in out - assert "Chunked Prefill Latency" not in out - assert "Decode Step Latency" not in out + assert "Chunked Prefill Latency" in out + assert "Decode Step Latency" in out + assert "0.00" in out # --------------------------------------------------------------------------- From bb6ce985b0a863a77a48cb50ed8074969bac858f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 16 Jun 2026 15:48:46 +0200 Subject: [PATCH 073/106] extract bench part of update_from_output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 91 ++++++++++++++------------- 1 file changed, 48 insertions(+), 43 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index bd05076bb..d50ad5e10 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -259,49 +259,7 @@ def update_from_output(self, scheduler_output, model_runner_output): # Measure timing durations and accumulate metrics (scheduler-side timing injection) if self._bench is not None: - now = time.time() - - # Prefill duration measurement (if prefill step was scheduled) - if self._bench.prefill_step_start is not None: - assert self.previous_step_was_prefill and self._bench.decode_step_start is None - t0 = self._bench.prefill_step_start - duration = now - t0 - all_prefill_reqs = [ - r.req_id for r in scheduler_output.scheduled_new_reqs - ] + scheduler_output.scheduled_cached_reqs.req_ids - for req_id in all_prefill_reqs: - self._bench.chunk_latencies.setdefault(req_id, []).append(duration) - self._bench.chunk_start_times.setdefault(req_id, []).append(t0) - self._bench.tkvs.setdefault(req_id, []).append(model_runner_output.tkv) - self._bench.prefill_step_start = None - self._bench.decode_step_start = None - - # Decode duration measurement (if decode step was scheduled) - elif self._bench.decode_step_start is not None: - assert not self.previous_step_was_prefill and self._bench.prefill_step_start is None - t0 = self._bench.decode_step_start - duration = now - t0 - tkv = model_runner_output.tkv - max_num_blocks = math.ceil(tkv / self.block_size) - req_by_id = {r.request_id: r for r in self.running} - for req_id in scheduler_output.scheduled_cached_reqs.req_ids: - self._bench.decode_latencies.setdefault(req_id, []).append(duration) - self._bench.decode_start_times.setdefault(req_id, []).append(t0) - self._bench.tkvs.setdefault(req_id, []).append(tkv) - req = req_by_id.get(req_id) - if req is not None: - req_num_blocks = math.ceil(req.num_computed_tokens / self.block_size) - self._bench.left_padding_blocks.setdefault(req_id, []).append( - max_num_blocks - req_num_blocks - ) - self._bench.prefill_step_start = None - self._bench.decode_step_start = None - - # Track first-scheduled time and arrival time for queue-wait calculation - for req in self.ongoing_prefills: - if req.request_id not in self._bench.first_scheduled_ts: - self._bench.first_scheduled_ts[req.request_id] = now - self._bench.arrival_ts[req.request_id] = req.arrival_time + self._bench_update_from_output(scheduler_output, model_runner_output) # Remove completed prefills self.ongoing_prefills = [ @@ -311,6 +269,53 @@ def update_from_output(self, scheduler_output, model_runner_output): self.tkv = model_runner_output.tkv return super(SpyreScheduler, self).update_from_output(scheduler_output, model_runner_output) + def _bench_update_from_output( + self, + scheduler_output: "SchedulerOutput", + model_runner_output: SpyreModelRunnerOutput, + ) -> None: + assert self._bench is not None + now = time.time() + + if self._bench.prefill_step_start is not None: + assert self.previous_step_was_prefill and self._bench.decode_step_start is None + t0 = self._bench.prefill_step_start + duration = now - t0 + all_prefill_reqs = [ + r.req_id for r in scheduler_output.scheduled_new_reqs + ] + scheduler_output.scheduled_cached_reqs.req_ids + for req_id in all_prefill_reqs: + self._bench.chunk_latencies.setdefault(req_id, []).append(duration) + self._bench.chunk_start_times.setdefault(req_id, []).append(t0) + self._bench.tkvs.setdefault(req_id, []).append(model_runner_output.tkv) + self._bench.prefill_step_start = None + self._bench.decode_step_start = None + + elif self._bench.decode_step_start is not None: + assert not self.previous_step_was_prefill and self._bench.prefill_step_start is None + t0 = self._bench.decode_step_start + duration = now - t0 + tkv = model_runner_output.tkv + max_num_blocks = math.ceil(tkv / self.block_size) + req_by_id = {r.request_id: r for r in self.running} + for req_id in scheduler_output.scheduled_cached_reqs.req_ids: + self._bench.decode_latencies.setdefault(req_id, []).append(duration) + self._bench.decode_start_times.setdefault(req_id, []).append(t0) + self._bench.tkvs.setdefault(req_id, []).append(tkv) + req = req_by_id.get(req_id) + if req is not None: + req_num_blocks = math.ceil(req.num_computed_tokens / self.block_size) + self._bench.left_padding_blocks.setdefault(req_id, []).append( + max_num_blocks - req_num_blocks + ) + self._bench.prefill_step_start = None + self._bench.decode_step_start = None + + for req in self.ongoing_prefills: + if req.request_id not in self._bench.first_scheduled_ts: + self._bench.first_scheduled_ts[req.request_id] = now + self._bench.arrival_ts[req.request_id] = req.arrival_time + def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: """Return and clear accumulated chunk timing for a finished request.""" if self._bench is None: From 0b8b170e14b090f006ff12e9507b0964e54990fe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 16 Jun 2026 16:25:14 +0200 Subject: [PATCH 074/106] address md error MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .claude/skills/add-bench-metric/SKILL.md | 75 ++++++++++++------------ 1 file changed, 39 insertions(+), 36 deletions(-) diff --git a/.claude/skills/add-bench-metric/SKILL.md b/.claude/skills/add-bench-metric/SKILL.md index fad1fe184..5da1845ed 100644 --- a/.claude/skills/add-bench-metric/SKILL.md +++ b/.claude/skills/add-bench-metric/SKILL.md @@ -14,7 +14,7 @@ If `$ARGUMENTS` is empty, ask the user what they want to measure before doing an Every custom bench metric travels through five layers (all scheduler-process-side): -``` +```text schedule() → SpyreBenchState (accumulate raw values) → _free_request() → kv_transfer_params["__spyre__"] (ZMQ to API server) → patch_serving.py SSE injection (into final usage SSE chunk) @@ -23,6 +23,7 @@ schedule() → SpyreBenchState (accumulate raw values) ``` **Key architecture facts:** + - All timing is measured in the scheduler. - `SpyreBenchState` is the single source of truth for per-request bench state. It is `None` when `SENDNN_INFERENCE_BENCH_METRICS_ENABLED` is off; every access must be guarded by `if self._bench is not None:`. - Timing works by bracket: at the **end** of `schedule()`, a step-start timestamp (`prefill_step_start` or `decode_step_start`) is written into `_bench`. At the **start** of the next `update_from_output()`, the duration is computed from that timestamp and `time.time()`. @@ -203,52 +204,54 @@ Use a `spyre_` prefix so the key is clearly SenDNN-owned in the vllm result JSON 4. **Update `test_scheduler_bench_metrics_accumulated`**: This is an integration test that runs a real engine and captures the bench state just before `_free_request` clears it. The capture (`_capturing_free`) and the post-run empty-dict check are both dynamic — they iterate over `dataclasses.fields(bench)` and need no changes. What you **do** need to add is a value assertion for your new field in the `for req_id in ("0", "1"):` block, following the pattern of the existing ones. For example, for a new list-per-step field: -```python -# In the for req_id in ("0", "1"): block: -assert len(info["my_new_metric"]) >= 1, ( - f"req {req_id}: expected ≥1 my_new_metric entry, got {info['my_new_metric']}" -) -for val in info["my_new_metric"]: - assert isinstance(val, float) and val > 0, f"req {req_id}: non-positive my_new_metric {val}" -``` + ```python + # In the for req_id in ("0", "1"): block: + assert len(info["my_new_metric"]) >= 1, ( + f"req {req_id}: expected ≥1 my_new_metric entry, got {info['my_new_metric']}" + ) + for val in info["my_new_metric"]: + assert isinstance(val, float) and val > 0, f"req {req_id}: non-positive my_new_metric {val}" + ``` -For a scalar field (`arrival_ts`-style), check `is not None`: -```python -assert info["my_scalar_field"] is not None, f"req {req_id}: my_scalar_field not set" -``` + For a scalar field (`arrival_ts`-style), check `is not None`: + + ```python + assert info["my_scalar_field"] is not None, f"req {req_id}: my_scalar_field not set" + ``` 5. **Update `test_get_and_clear_returns_correct_dict`**: The test is driven by two dicts defined just above it — update both: - **`_BENCH_FIXTURE`** — add an entry for your new `SpyreBenchState` field. This dict must cover **every** field of `SpyreBenchState` (dict and scalar alike); `test_bench_fixture_covers_all_per_req_fields` compares `_BENCH_FIXTURE.keys()` against `dataclasses.fields(bench)` and fails if they diverge. For dict fields the value is used as the per-request payload (`bench.["r0"] = value`); for scalar fields it is set directly (`setattr(bench, field, value)`). - **`_EXPECTED_RESULT`** — add an entry `"": `. `test_get_and_clear_result_keys` asserts `result.keys() == _EXPECTED_RESULT.keys()`, so it will fail if the returned dict has any extra or missing keys. -Example: -```python -# In FAKE_METRICS entry 1: -"my_new_metric": [0.001, 0.002], + Example: -# In test_inject_adds_spyre_keys: -assert "spyre_my_new_metric" in data + ```python + # In FAKE_METRICS entry 1: + "my_new_metric": [0.001, 0.002], -# In test_inject_values_correct: -assert data["spyre_my_new_metric"] == [[0.001, 0.002], [0.003]] + # In test_inject_adds_spyre_keys: + assert "spyre_my_new_metric" in data -# In the sentinel dicts (dict field example): -_BENCH_FIXTURE: dict[str, Any] = { - ..., - "my_new_metric": [0.001, 0.002], # NEW — dict field, keyed by req_id at test time -} -_EXPECTED_RESULT: dict[str, Any] = { - ..., - "my_new_metric_s": pytest.approx([0.001, 0.002]), # NEW — key in returned dict -} + # In test_inject_values_correct: + assert data["spyre_my_new_metric"] == [[0.001, 0.002], [0.003]] -# Scalar field example (e.g. a single float per request stored directly): -_BENCH_FIXTURE: dict[str, Any] = { - ..., - "my_scalar_field": 42.0, # NEW — set directly via setattr -} -``` + # In the sentinel dicts (dict field example): + _BENCH_FIXTURE: dict[str, Any] = { + ..., + "my_new_metric": [0.001, 0.002], # NEW — dict field, keyed by req_id at test time + } + _EXPECTED_RESULT: dict[str, Any] = { + ..., + "my_new_metric_s": pytest.approx([0.001, 0.002]), # NEW — key in returned dict + } + + # Scalar field example (e.g. a single float per request stored directly): + _BENCH_FIXTURE: dict[str, Any] = { + ..., + "my_scalar_field": 42.0, # NEW — set directly via setattr + } + ``` 6. **Update `test_print_spyre_section_output`** (if you added a new `_section()` call): Add `out.assert_contains("")` for the section separator, and add the label string (third argument to `_section()`) to the `for label in (...)` loop so mean/median/percentile lines are asserted. `assert_all_lines_covered()` is called at the end of the test and will fail if any output line was not covered by an assertion — the error message lists the exact uncovered lines. From d1af9855995d191ff275675d041730d1eb9db549 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 16 Jun 2026 16:43:21 +0200 Subject: [PATCH 075/106] address precommit hook error MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/argparse_utils.py | 2 +- sendnn_inference/benchmarks/spyre_bench_serve.py | 2 +- sendnn_inference/benchmarks/spyre_plot.py | 4 ++-- sendnn_inference/platform.py | 4 ++-- sendnn_inference/v1/core/scheduler.py | 2 +- sendnn_inference/v1/worker/spyre_model_runner.py | 4 ++-- 6 files changed, 9 insertions(+), 9 deletions(-) diff --git a/sendnn_inference/argparse_utils.py b/sendnn_inference/argparse_utils.py index fde17fc3d..b3e0ceb52 100644 --- a/sendnn_inference/argparse_utils.py +++ b/sendnn_inference/argparse_utils.py @@ -162,7 +162,7 @@ def patched_parse_args( namespace: argparse.Namespace | None = None, ) -> argparse.Namespace: result = original_parse_args(self, args, namespace) - assert result is not None # type: ignore[redundant-expr] + assert result is not None if args is None or len(args) == 0: # Don't override anything if there were no args parsed diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index a0e88060f..a3b7f82a2 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -359,7 +359,7 @@ def main() -> None: candidates = [] if candidates: - json_path = Path(max(candidates, key=os.path.getmtime)) + json_path = Path(str(max(candidates, key=os.path.getmtime))) html_path = json_path.with_name(json_path.stem + "_detailed_timeline.html") decode_thresholds_str = getattr(args, "decode_thresholds", None) # Parse comma-separated milliseconds and convert to seconds diff --git a/sendnn_inference/benchmarks/spyre_plot.py b/sendnn_inference/benchmarks/spyre_plot.py index 40a07217e..78242b35f 100644 --- a/sendnn_inference/benchmarks/spyre_plot.py +++ b/sendnn_inference/benchmarks/spyre_plot.py @@ -229,8 +229,8 @@ def generate_detailed_timeline_plot( """ try: import pandas as pd - import plotly.express as px - import plotly.io as pio + import plotly.express as px # ty: ignore[unresolved-import] + import plotly.io as pio # ty: ignore[unresolved-import] except ImportError as exc: logger.warning( "Cannot generate detailed timeline plot — missing dependency: %s. " diff --git a/sendnn_inference/platform.py b/sendnn_inference/platform.py index 012d31f74..8f0ea400d 100644 --- a/sendnn_inference/platform.py +++ b/sendnn_inference/platform.py @@ -348,7 +348,7 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: scheduler_config.max_num_batched_tokens = ( model_config.max_model_len * scheduler_config.max_num_seqs ) - cache_config.block_size = model_config.max_model_len # ty: ignore[invalid-assignment] + cache_config.block_size = model_config.max_model_len vllm_config.cache_config.enable_prefix_caching = False else: @@ -824,7 +824,7 @@ def maybe_ensure_sendnn_configured(cls, model_config: ModelConfig) -> None: @classmethod def _set_batch_tkv_limit_from_env(cls) -> None: try: - cls._max_batch_tkv_limit = int(os.getenv("VLLM_DT_MAX_BATCH_TKV_LIMIT", "-1")) # ty: ignore + cls._max_batch_tkv_limit = int(os.getenv("VLLM_DT_MAX_BATCH_TKV_LIMIT", "-1")) except ValueError as e: raise ValueError("VLLM_DT_MAX_BATCH_TKV_LIMIT must be an integer") from e diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 637b8ff06..bd0af0807 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -545,7 +545,7 @@ def schedule(self) -> "SchedulerOutput": ready_to_prefill = [ r for r in self.waiting - if r.status != RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR # type: ignore[attr-defined] + if r.status != RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR ] if ready_to_prefill: new_prefill_candidates = list(self.waiting) diff --git a/sendnn_inference/v1/worker/spyre_model_runner.py b/sendnn_inference/v1/worker/spyre_model_runner.py index 5890ebfde..6e40a956d 100644 --- a/sendnn_inference/v1/worker/spyre_model_runner.py +++ b/sendnn_inference/v1/worker/spyre_model_runner.py @@ -393,10 +393,10 @@ def vocab_size(self) -> int: # self.model here is probably a transformers model class if self.model_config.architecture in FMS_POOLING_MODEL_LIST: assert isinstance(self.model.config.src_vocab_size, int) - return self.model.config.src_vocab_size # ty: ignore[invalid-return-type] + return self.model.config.src_vocab_size else: assert isinstance(self.model.config.vocab_size, int) - return self.model.config.vocab_size # ty: ignore[invalid-return-type] + return self.model.config.vocab_size def _prepare_pad_input_ids( self, From 7722d20faa2515a0e67c44fc7872821a420b97b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 16 Jun 2026 17:13:02 +0200 Subject: [PATCH 076/106] fix tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .github/workflows/test.yml | 3 +++ tests/v1/core/test_scheduler_structured_outputs.py | 1 + 2 files changed, 4 insertions(+) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 1f2dc0964..f8b6d0715 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -53,6 +53,9 @@ jobs: - name: "scoring" markers: "cpu and scoring" flags: "--timeout=300" + - name: "benchmarks" + markers: "cpu" + flags: "--timeout=300 tests/benchmarks/" - name: "worker and utils" markers: "not e2e and not quantized and not spyre and not multimodal" flags: "--timeout=300" diff --git a/tests/v1/core/test_scheduler_structured_outputs.py b/tests/v1/core/test_scheduler_structured_outputs.py index f85f44db0..b15cedb88 100644 --- a/tests/v1/core/test_scheduler_structured_outputs.py +++ b/tests/v1/core/test_scheduler_structured_outputs.py @@ -57,6 +57,7 @@ def mocked_scheduler(): scheduler.available_blocks = 1 scheduler.total_reserved_blocks = 0 scheduler.reserved_blocks = dict[str, int]() + scheduler._bench = None scheduler._get_required_blocks = lambda x, *args, **kwargs: (0, 0) scheduler._get_free_blocks = lambda *args, **kwargs: 1 From 6e455fabbf1a05d4adfcc9a2ac028452b893dc8d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 16 Jun 2026 17:47:41 +0200 Subject: [PATCH 077/106] timestamp right before return MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index bd0af0807..73f9fca0c 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -604,16 +604,6 @@ def schedule(self) -> "SchedulerOutput": # collection for better performance. outputs._spyre_grammar_output = self.get_grammar_bitmask(outputs) # type: ignore[attr-defined] - # Inject scheduler-side step-start timestamps for accurate timing measurement - if self._bench is not None: - now = time.time() - if self.previous_step_was_prefill: - self._bench.prefill_step_start = now - self._bench.decode_step_start = None - else: - self._bench.decode_step_start = now - self._bench.prefill_step_start = None - # As blocks are allocated, we discount them from the reserved blocks. # For prefill blocks we must first subtract the cached blocks. free_blocks = self._get_free_blocks() @@ -640,6 +630,16 @@ def schedule(self) -> "SchedulerOutput": assert 0 <= self.total_reserved_blocks <= free_blocks + # Inject step-start timestamps for timing measurement + if self._bench is not None: + now = time.time() + if self.previous_step_was_prefill: + self._bench.prefill_step_start = now + self._bench.decode_step_start = None + else: + self._bench.decode_step_start = now + self._bench.prefill_step_start = None + return outputs def can_schedule_prefill(self, request: Request) -> bool: From 2a262b15306eb3ee4b11d16b76b53abdefa775be Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 16 Jun 2026 19:31:22 +0200 Subject: [PATCH 078/106] rename prev_was_prefill to is_prefill MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 4b1254185..c33b7b657 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -270,7 +270,7 @@ def _bench_update_from_output( now = time.time() if self._bench.prefill_step_start is not None: - assert self.previous_step_was_prefill and self._bench.decode_step_start is None + assert self.step_is_prefill and self._bench.decode_step_start is None t0 = self._bench.prefill_step_start duration = now - t0 all_prefill_reqs = [ @@ -284,7 +284,7 @@ def _bench_update_from_output( self._bench.decode_step_start = None elif self._bench.decode_step_start is not None: - assert not self.previous_step_was_prefill and self._bench.prefill_step_start is None + assert not self.step_is_prefill and self._bench.prefill_step_start is None t0 = self._bench.decode_step_start duration = now - t0 tkv = model_runner_output.tkv @@ -638,7 +638,7 @@ def schedule(self) -> "SchedulerOutput": # Inject step-start timestamps for timing measurement if self._bench is not None: now = time.time() - if self.previous_step_was_prefill: + if self.step_is_prefill: self._bench.prefill_step_start = now self._bench.decode_step_start = None else: From 7a3ec53b45f053b77ff4b043e02322be9edceba0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 16 Jun 2026 19:39:20 +0200 Subject: [PATCH 079/106] don't put dedicated test workflow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .github/workflows/test.yml | 3 --- 1 file changed, 3 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index f8b6d0715..1f2dc0964 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -53,9 +53,6 @@ jobs: - name: "scoring" markers: "cpu and scoring" flags: "--timeout=300" - - name: "benchmarks" - markers: "cpu" - flags: "--timeout=300 tests/benchmarks/" - name: "worker and utils" markers: "not e2e and not quantized and not spyre and not multimodal" flags: "--timeout=300" From 192af4bbb6ea535002c387ceaaa3fb8f4ad75284 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 17 Jun 2026 22:15:58 +0200 Subject: [PATCH 080/106] count the number of requests blocked by missing kv blocks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_bench_serve.py | 7 ++++++- sendnn_inference/v1/core/scheduler.py | 6 ++++++ tests/benchmarks/test_bench_metrics.py | 18 +++++++++++++++--- 3 files changed, 27 insertions(+), 4 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 7c1abc9b2..dc7a8d795 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -142,8 +142,11 @@ def _print_spyre_section( decode_lats_ms = [lat * 1000 for m in metrics_list for lat in m.get("decode_latencies_s", [])] total_prefill_chunks = sum(num_chunks_list) - # Scalar summary line (mirrors vllm's plain-count header section) + total_missing_blocks = sum(1 for m in metrics_list if m.get("was_missing_blocks", False)) + + # Scalar summary lines (mirrors vllm's plain-count header section) print("{:<40} {:<10}".format("Total prefill chunks processed:", total_prefill_chunks)) + print("{:<40} {:<10}".format("Requests blocked by missing KV blocks:", total_missing_blocks)) def _section(header: str, values: list[float], label: str) -> None: if not values: @@ -257,6 +260,8 @@ def _inject_spyre_metrics_into_result_file( result["spyre_left_padding_blocks"] = [m.get("left_padding_blocks", []) for m in metrics_list] result["spyre_pause_latencies_s"] = [m.get("pause_latencies_s", []) for m in metrics_list] result["spyre_pause_start_times_s"] = [m.get("pause_start_times_s", []) for m in metrics_list] + result["spyre_was_missing_blocks"] = [m.get("was_missing_blocks", False) for m in metrics_list] + result["spyre_num_requests_missing_blocks"] = sum(result["spyre_was_missing_blocks"]) try: with open(file_path, "w", encoding="utf-8") as fh: diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index c33b7b657..1b42a73ed 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -39,6 +39,7 @@ class SpyreBenchState: left_padding_blocks: dict[str, list[int]] = field(default_factory=dict) pause_start_times: dict[str, list[float]] = field(default_factory=dict) pause_latencies: dict[str, list[float]] = field(default_factory=dict) + blocks_lacking: dict[str, bool] = field(default_factory=dict) prefill_step_start: float | None = None decode_step_start: float | None = None @@ -320,6 +321,7 @@ def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: left_padding_blocks = self._bench.left_padding_blocks.pop(req_id, None) pause_lats = self._bench.pause_latencies.pop(req_id, None) pause_starts = self._bench.pause_start_times.pop(req_id, None) + was_missing_blocks = self._bench.blocks_lacking.pop(req_id, False) if lats is None and dec_lats is None: return None return { @@ -332,6 +334,7 @@ def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: "left_padding_blocks": left_padding_blocks or [], "pause_latencies_s": pause_lats or [], "pause_start_times_s": pause_starts or [], + "was_missing_blocks": was_missing_blocks, } def _free_request(self, request, delay_free_blocks: bool = False): @@ -367,6 +370,7 @@ def _free_request(self, request, delay_free_blocks: bool = False): "left_padding_blocks": chunk_stats["left_padding_blocks"] if chunk_stats else [], "pause_latencies_s": chunk_stats["pause_latencies_s"] if chunk_stats else [], "pause_start_times_s": chunk_stats["pause_start_times_s"] if chunk_stats else [], + "was_missing_blocks": chunk_stats["was_missing_blocks"] if chunk_stats else False, } if kv_xfer_params is None: kv_xfer_params = {"__spyre__": spyre_data} @@ -482,6 +486,8 @@ def schedule(self) -> "SchedulerOutput": new_request = holdback_queue[0] cached, blocks = self._get_required_blocks(new_request, True) if blocks > available_blocks: + if self._bench is not None: + self._bench.blocks_lacking[new_request.request_id] = True break if self.can_schedule_prefill(new_request): diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index a9f0d5094..9ecb97cee 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -42,6 +42,7 @@ "left_padding_blocks": [2, 0, 1], "pause_latencies_s": [0.5, 1.2], "pause_start_times_s": [0.0, 1.2], + "was_missing_blocks": True, }, { "queued_time_s": 0.00001, @@ -55,6 +56,7 @@ "left_padding_blocks": [3, 1], "pause_latencies_s": [0.3], "pause_start_times_s": [0.5], + "was_missing_blocks": False, }, ] @@ -86,7 +88,10 @@ def test_inject_adds_spyre_keys(tmp_path): result_file = _write_fake_result(tmp_path) _inject_spyre_metrics_into_result_file(_make_args(tmp_path), FAKE_METRICS, time.time() - 1) data = json.loads(result_file.read_text()) - expected_keys = {"spyre_" + k for k in FAKE_METRICS[0]} | {"spyre_total_prefill_chunks"} + expected_keys = {"spyre_" + k for k in FAKE_METRICS[0]} | { + "spyre_total_prefill_chunks", + "spyre_num_requests_missing_blocks", + } for key in expected_keys: assert key in data, f"expected key {key!r} missing from result JSON" @@ -103,10 +108,13 @@ def test_inject_values_correct(tmp_path): continue expected = [m[key] for m in FAKE_METRICS] assert data["spyre_" + key] == expected - # Derived run-level scalar + # Derived run-level scalars assert data["spyre_total_prefill_chunks"] == sum( m["num_chunked_prefills"] for m in FAKE_METRICS ) + assert data["spyre_num_requests_missing_blocks"] == sum( + 1 for m in FAKE_METRICS if m.get("was_missing_blocks", False) + ) # Original keys preserved assert data["backend"] == "spyre-chat" @@ -196,9 +204,10 @@ def test_print_spyre_section_output(capsys): _print_spyre_section(FAKE_METRICS, SELECTED_PERCENTILES) out = _TrackedOutput(capsys.readouterr().out) - # Run-level scalar + # Run-level scalars out.assert_contains("Total prefill chunks processed:") out.assert_contains("10") + out.assert_contains("Requests blocked by missing KV blocks: 1") # Section separators and mean/median/percentile lines for all sections out.assert_contains("Queue Wait Time") @@ -285,6 +294,7 @@ def _make_bare_scheduler(): "left_padding_blocks": [2, 0], "pause_start_times": [3000.0, 3005.0], "pause_latencies": [0.5, 1.2], + "blocks_lacking": True, "prefill_step_start": 999.0, "decode_step_start": 1999.0, } @@ -300,6 +310,8 @@ def _make_bare_scheduler(): "tkvs": [64, 128, 192, 256], "left_padding_blocks": [2, 0], "pause_latencies_s": pytest.approx([0.5, 1.2]), + "pause_start_times_s": pytest.approx([3000.0, 3005.0]), + "was_missing_blocks": True, } From 5260e9c2191bac118603e8c6a7e6adfa23d1c6f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Thu, 18 Jun 2026 13:30:41 +0200 Subject: [PATCH 081/106] restore cond2 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 1b42a73ed..6e9172348 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -725,7 +725,16 @@ def _satisfies_last_chunk_constraints(self, request: Request) -> bool: # check that there is space in the current decode batch num_running = len(decoding_requests) - return num_running + len(self.waiting) < self.max_num_running_reqs + cond1 = num_running + len(self.waiting) < self.max_num_running_reqs + + if request not in decoding_requests: + decoding_requests += [request] + + # check that we can prefill and add the new request to the decode batch without + # getting it paused + cond2 = lambda: self._can_decode_all_requests(decoding_requests) + + return cond1 and cond2() def _has_scheduling_priority(self, request): decoding_requests = [r for r in self.running if r not in self.ongoing_prefills] From db632e13ed11d0bd3f7c8b45ce028f032a8fbe5f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Thu, 18 Jun 2026 17:04:04 +0200 Subject: [PATCH 082/106] per request metrics json file save: one line per metric MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/benchmarks/spyre_bench_serve.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index a3b7f82a2..50bb557bc 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -252,7 +252,12 @@ def _inject_spyre_metrics_into_result_file( try: with open(file_path, "w", encoding="utf-8") as fh: - json.dump(result, fh) + fh.write("{\n") + items = list(result.items()) + for i, (k, v) in enumerate(items): + comma = "," if i < len(items) - 1 else "" + fh.write(f" {json.dumps(k)}: {json.dumps(v)}{comma}\n") + fh.write("}\n") logger.info("Spyre metrics injected into %s", file_path) except Exception as exc: logger.warning("Failed to write Spyre metrics into result JSON %s: %s", file_path, exc) From dc53819fa0a2dfe07f97082e739a5793cd26e361 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 24 Jun 2026 09:01:43 +0200 Subject: [PATCH 083/106] remove dead code MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 3 -- sendnn_inference/v1/metrics/stats_logger.py | 59 --------------------- 2 files changed, 62 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 73f9fca0c..c55eaa40e 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -228,9 +228,6 @@ def __init__(self, *args, **kwargs) -> None: if envs_spyre.SENDNN_INFERENCE_BENCH_METRICS_ENABLED: self._bench = SpyreBenchState() - from sendnn_inference.v1.metrics.stats_logger import register_scheduler - - register_scheduler(self) assert self.max_batch_tkv_limit != -1, ( "Expecting the env var VLLM_DT_MAX_BATCH_TKV_LIMIT to be set in platform.py" diff --git a/sendnn_inference/v1/metrics/stats_logger.py b/sendnn_inference/v1/metrics/stats_logger.py index e1c204970..239706422 100644 --- a/sendnn_inference/v1/metrics/stats_logger.py +++ b/sendnn_inference/v1/metrics/stats_logger.py @@ -1,11 +1,9 @@ import dataclasses import json -import threading import time from datetime import datetime from functools import wraps from pathlib import Path -from typing import Any from vllm.config import VllmConfig from vllm.logger import init_logger @@ -23,63 +21,6 @@ logger = init_logger(__name__) -# --------------------------------------------------------------------------- -# Bench metrics registry — only active when SENDNN_INFERENCE_BENCH_METRICS_ENABLED -# --------------------------------------------------------------------------- - - -@dataclasses.dataclass -class SpyreRequestMetrics: - """Per-request Spyre-specific metrics surfaced to benchmark clients.""" - - request_id: str - queued_time_s: float - num_chunked_prefills: int - chunk_prefill_latencies_s: list[float] - chunk_prefill_start_times_s: list[float] = dataclasses.field(default_factory=list) - decode_latencies_s: list[float] = dataclasses.field(default_factory=list) - decode_start_times_s: list[float] = dataclasses.field(default_factory=list) - tkvs: list[int] = dataclasses.field(default_factory=list) - - -class SpyreMetricsRegistry: - """Thread-safe store of per-request SpyreRequestMetrics, cleared on read.""" - - def __init__(self) -> None: - self._store: dict[str, SpyreRequestMetrics] = {} - self._lock = threading.Lock() - - def put(self, metrics: SpyreRequestMetrics) -> None: - with self._lock: - self._store[metrics.request_id] = metrics - - def get_and_clear(self, request_id: str) -> SpyreRequestMetrics | None: - with self._lock: - return self._store.pop(request_id, None) - - -_REGISTRY: SpyreMetricsRegistry | None = None -_SCHEDULER: Any = None # set by ChunkedPrefillSpyreScheduler.__init__ - - -def enable_registry() -> SpyreMetricsRegistry: - global _REGISTRY - _REGISTRY = SpyreMetricsRegistry() - return _REGISTRY - - -def get_registry() -> SpyreMetricsRegistry | None: - return _REGISTRY - - -def register_scheduler(scheduler: Any) -> None: - """Called by ChunkedPrefillSpyreScheduler at init time.""" - global _SCHEDULER, _REGISTRY - _SCHEDULER = scheduler - if _REGISTRY is None: - _REGISTRY = SpyreMetricsRegistry() - - @dataclasses.dataclass class PerfRecord: """A record for request_metrics.jsonl. From c6f9351c49793507e2469bcc31763f7553425232 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 30 Jun 2026 11:54:45 +0200 Subject: [PATCH 084/106] add explanation comment for queued_time_s MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index c55eaa40e..dde522737 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -341,6 +341,13 @@ def _free_request(self, request, delay_free_blocks: bool = False): chunk_stats = self.get_and_clear_chunk_stats(req_id) first_ts = self._bench.first_scheduled_ts.pop(req_id, None) arrival_ts = self._bench.arrival_ts.pop(req_id, None) + # NOTE: queued_time_s looks like a duplicate of FinishedRequestStats.queued_time, but + # it is not. FinishedRequestStats.queued_time is (scheduled_ts - queued_ts): both + # timestamps are recorded inside the engine core, so it misses the time the request + # spent in transit from the API server to the engine (IPC hop). Here we use the stamp + # of the API server on receipt (request.arrival_time), which gives a complete client- + # visible queue wait that adds up cleanly with prefill and decode latencies when + # reconstructing TTFT. queued_time_s = ( (first_ts - arrival_ts) if first_ts is not None and arrival_ts is not None else 0.0 ) From 3f6ef78341ec3b5350b6fc3f2b051be11d8b06fd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Mon, 10 Aug 2026 09:36:40 +0000 Subject: [PATCH 085/106] restore diverging to main MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 225 +++++++++----------------- 1 file changed, 75 insertions(+), 150 deletions(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 65a5076af..f2f819ec1 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -834,9 +834,14 @@ def _satisfies_last_chunk_constraints(self, request: Request) -> bool: n_blocks = math.floor(max(self.tkv, prompt_len) / self.block_size) new_req_tkv = n_blocks * self.block_size + prompt_len % self.block_size - # check that we can prefill and add the new request to the decode batch without - # getting it paused - cond2 = lambda: self._can_decode_all_requests(decoding_requests) + # check that batch size x tkv is smaller than the max supported number + # Note: using max_tkv is a conservative upper bound here. For the + # optimal check we need model runner to return per sequence tkvs + cond2 = lambda: self.check_batch_tkv_limit_cp( + request=request, + new_req_tkv=new_req_tkv, + running=decoding_requests, + ) return cond1 and cond2() @@ -858,150 +863,77 @@ def _has_scheduling_priority(self, request): num_prefills = len(self.waiting) + len(self.ongoing_prefills) return num_prefills < max_concurrent_prefills - def _can_decode_all_requests(self, decoding_requests: list[Request]) -> bool: + def check_batch_tkv_limit_cp(self, request: Request, new_req_tkv: int, running) -> bool: """ - Check if all decoding requests can be decoded in the next step without - violating the max batch TKV limit. + Check whether adding a new sequence to the decode batch would violate + Spyre's maximum batch volume constraint for chunked prefill. + + In Spyre, the product of `batch_size` and the current `tkv` + (tokens-per-sequence) must not exceed the limit defined by + `VLLM_DT_MAX_BATCH_TKV_LIMIT`. Before scheduling a new sequence, + we must ensure that this constraint will hold for all decoding + steps that result from combining the new sequence with the currently + running decode batch. + + This implementation: + 1. Computes the maximum possible `tkv` for each sequence in the + decode batch. + 2. Sorts these values in ascending order. + 3. Iterates through them, stopping once the `tkv` of the new sequence. + is reached. Remaining sequences do not need to be checked explicitly, + since they were validated when they were added (by inductive reasoning). + + Note: drawing explaining the algorithm in more detail uploaded here: + https://github.com/torch-spyre/sendnn-inference/pull/363#issuecomment-3173605517 """ - if not decoding_requests: - return True - - next_predicted_tkv = self.predict_next_decode_tkv(decoding_requests) - - # the tkv should never get beyond max_model_len - assert next_predicted_tkv <= self.max_model_len - - # check batch tkv limit: batch_size * predicted_tkv must not exceed limit - batch_size = len(decoding_requests) - predicted_batch_tkv = batch_size * next_predicted_tkv - - return predicted_batch_tkv <= self.max_batch_tkv_limit - - def _handle_decode_requests_pausing(self) -> None: - """ - Manage pausing and resuming of decode requests based on batch TKV constraints. - - This method: - 1. Pauses requests with the fewest decoded tokens when batch TKV limit is exceeded - 2. Resumes previously paused requests (oldest first) when capacity is available - """ - decoding_requests = [r for r in self.running if r not in self.ongoing_prefills] - resumed = self._maybe_resume_decoding_requests(decoding_requests) - if not resumed: - self._maybe_pause_decoding_requests(decoding_requests) - def _maybe_pause_decoding_requests(self, decoding_requests: list[Request]) -> int: - """ - Iteratively pauses requests with the fewest decoded tokens until the batch fits - within TKV constraints. Mutates both decoding_requests and self.running. - - Returns the number of requests paused. - """ - initial_had_requests = len(decoding_requests) > 0 - num_paused = 0 - - # TODO we should test different removal logics: longest request, optimize padding - while not self._can_decode_all_requests(decoding_requests): - # Decoded tokens = num_computed_tokens - num_prompt_tokens - request_to_remove = min( - decoding_requests, key=lambda r: r.num_computed_tokens - r.num_prompt_tokens + # Compute the effective token length of the new request + # Rounded up to the nearest block size to account for potential padding + new_req_max_tkv = round_up_to_block_size(new_req_tkv + request.max_tokens - 1) + # Extra block of slack: left-padding can push a sequence's runtime tkv up to + # one block past the scheduler's estimate when the batch re-aligns on admission. + new_req_max_tkv += self.block_size + + # Compute token lengths for all running requests (decode batch) + decode_req_max_tkvs = [] + # Decide new tkv based on max of current tkv or new request prompt tokens + dec_req_tkv = max(self.tkv, request.num_prompt_tokens) + for req in running: + n_generated_output_tokens = req.num_computed_tokens - req.num_prompt_tokens + # Rounded up to the nearest block size to account for potential padding + dec_req_max_tkv = round_up_to_block_size( + dec_req_tkv + (req.max_tokens - n_generated_output_tokens) - 1 ) - decoding_requests.remove(request_to_remove) - self.running.remove(request_to_remove) - self.paused_decoding_requests.append(request_to_remove) - logger.info("Request %s paused due to batch TKV limit ", request_to_remove.request_id) - num_paused += 1 - - # update benchmark pausing statistics - if self._bench is not None: - pause_ts = time.time() - self._bench.pause_start_times.setdefault(request_to_remove.request_id, []).append( - pause_ts - ) - - # It shouldn't be possible to remove all requests if we started with some - assert not initial_had_requests or len(decoding_requests) > 0 - - return num_paused - - def _maybe_resume_decoding_requests(self, decoding_requests: list[Request]) -> int: - """ - Resumes previously paused requests (oldest first) when TKV capacity is available. - Mutates both decoding_requests and self.running. - - Returns the number of requests resumed. - """ - num_resumed = 0 - - # Reverse iteration: pop(i) only shifts indices above i, which are already visited. - for i in range(len(self.paused_decoding_requests) - 1, -1, -1): - request_to_add = self.paused_decoding_requests[i] - test_requests = decoding_requests + [request_to_add] - - if self._can_decode_all_requests(test_requests): - self.paused_decoding_requests.pop(i) - self.running.append(request_to_add) - decoding_requests.append(request_to_add) - logger.info( - "Request %s resumed (batch TKV capacity available).", - request_to_add.request_id, - ) - num_resumed += 1 - - # update benchmark pausing statistics - if self._bench is not None: - starts = self._bench.pause_start_times.get(request_to_add.request_id) - if starts: - duration = time.time() - starts[-1] - self._bench.pause_latencies.setdefault( - request_to_add.request_id, [] - ).append(duration) - - return num_resumed - - def predict_next_decode_tkv(self, running_requests: list[Request]) -> int: - """ - Predicts the TKV after the next decode step for a given batch of running - requests. - - This method replicates the TKV calculation logic from the model runner's - _prepare_decode method, accounting for: - - Block alignment (left-padding to make batch rectangular) - - The next token that will be generated (+1) - - Maximum TKV across all requests in the batch - - Args: - running_requests: List of Request objects currently in the decode batch - - Returns: - The predicted TKV value after the next decode step - """ - if not running_requests: - return 0 - - # Step 1: Find the maximum number of blocks across all requests - # Account for requests that will need a new block after the next token - max_n_blocks = 0 - num_blocks_per_req: list[int] = [] - for request in running_requests: - num_blocks = math.ceil((request.num_computed_tokens + 1) / self.block_size) - num_blocks_per_req.append(num_blocks) - max_n_blocks = max(max_n_blocks, num_blocks) - - # Step 2: Calculate TKV for each request and find the maximum - max_tkv = 0 - for request, num_blocks in zip(running_requests, num_blocks_per_req): - # Calculate left padding blocks needed for alignment - left_pad_blocks_count = max_n_blocks - num_blocks - left_padding = left_pad_blocks_count * self.block_size - - # Calculate TKV for this request (including the next token) - req_tkv = left_padding + request.num_computed_tokens + 1 - - # Track the maximum TKV - max_tkv = max(max_tkv, req_tkv) + # Extra block of slack: left-padding can push a sequence's runtime tkv up to + # one block past the scheduler's estimate when the batch re-aligns on admission. + dec_req_max_tkv += self.block_size + + decode_req_max_tkvs.append(dec_req_max_tkv) + + # Sort decode requests token lengths in ascending order + decode_req_max_tkvs.sort() + + # Initialize values + # The request is already in the running queue if it has done a first + # chunked prefill + batch_size = len(running) + if request not in running: + batch_size += 1 + max_batch_tkv = 0 + + # Try adding the new request to the batch and check the max volume + for decode_req_max_tkv in decode_req_max_tkvs: + if new_req_max_tkv <= decode_req_max_tkv: + # If the new request is shorter, it limits the batch volume + max_batch_tkv = max(max_batch_tkv, batch_size * new_req_max_tkv) + break + else: + # Otherwise, use the current (longer) request's volume + max_batch_tkv = max(max_batch_tkv, batch_size * decode_req_max_tkv) + # decrease batch_size by 1 as the current request finished + batch_size -= 1 - return max_tkv + return max_batch_tkv <= self.max_batch_tkv_limit def _can_decode_all_requests(self, decoding_requests: list[Request]) -> bool: """ @@ -1165,13 +1097,6 @@ def finish_requests( else [r for r in self.paused_decoding_requests if r.request_id not in request_ids] ) - # Also remove from paused_decoding_requests - self.paused_decoding_requests = ( - [] - if request_ids is None - else [r for r in self.paused_decoding_requests if r.request_id not in request_ids] - ) - return aborted_requests def calc_cached_tokens(self, prompt_len: int) -> tuple[int, int]: @@ -1232,4 +1157,4 @@ def make_stats(self, *args, **kwargs) -> SchedulerStats | None: self.pause_events = 0 self.resume_events = 0 - return base_stats + return base_stats \ No newline at end of file From 76d73e6f5cc24b8d41070fd4d2a26eabe852bdfe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Mon, 10 Aug 2026 09:58:53 +0000 Subject: [PATCH 086/106] bench pausing starting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index f2f819ec1..1adf2bc38 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -981,17 +981,29 @@ def _handle_decode_requests_pausing(self) -> None: self.paused_decoding_requests.clear() decoding_requests.clear() + now = time.time() if self._bench is not None else 0.0 + for req, _ in request_order: if self._can_decode_all_requests(decoding_requests + [req]): decoding_requests.append(req) self.request_last_decode_step[req.request_id] = 0 if req.request_id in was_paused: self.running.append(req) + # bench: pause -> running, close the currently open interval + if self._bench is not None: + starts = self._bench.pause_start_times.get(req.request_id) + if starts: + self._bench.pause_latencies.setdefault(req.request_id, []).append( + now - starts[-1] + ) else: self.paused_decoding_requests.append(req) self.request_last_decode_step[req.request_id] += 1 if req.request_id in was_running: self.running.remove(req) + # bench: running -> paused, open a new interval + if self._bench is not None: + self._bench.pause_start_times.setdefault(req.request_id, []).append(now) pause_inc = len(self.paused_decoding_requests) - len(was_paused) if pause_inc >= 0: From 92da4e9d024b0294137a40fbdfc9142b65251b6d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Mon, 10 Aug 2026 10:02:45 +0000 Subject: [PATCH 087/106] restore round_up function MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 1adf2bc38..012fc1754 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -59,6 +59,12 @@ class MMEncodeRequest: assert SpyrePlatform.get_block_size() == 64 +def round_up_to_block_size(n: int) -> int: + # Helper function to round up to the nearest block size + # Uses bitwise alignment for better performance + return (n + 63) & ~63 + + class SpyreScheduler(Scheduler): """Base class inheriting from the V1 scheduler to support static and continuous batching respecting AIU Spyre constraints.""" From 72507da98923758a87bfcf01a60d4ed131a52631 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Mon, 10 Aug 2026 10:30:55 +0000 Subject: [PATCH 088/106] format MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/benchmarks/spyre_bench_serve.py | 5 +++-- sendnn_inference/v1/core/scheduler.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 942f355d1..c98c24ee5 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -24,6 +24,7 @@ RequestFuncInput, ) from vllm.benchmarks.serve import add_cli_args, main_async +from vllm.utils.argparse_utils import FlexibleArgumentParser from sendnn_inference.benchmarks.spyre_request_func import async_request_spyre_chat @@ -81,9 +82,9 @@ def _register_backend() -> None: OPENAI_COMPATIBLE_BACKENDS.append(_BACKEND_NAME) -def _build_parser() -> argparse.ArgumentParser: +def _build_parser() -> FlexibleArgumentParser: """Build an arg parser based on vllm's but with spyre-chat as default backend.""" - parser = argparse.ArgumentParser( + parser = FlexibleArgumentParser( description="Spyre-extended vllm bench serve", formatter_class=argparse.RawDescriptionHelpFormatter, ) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 012fc1754..7525a5841 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -43,8 +43,9 @@ class SpyreBenchState: blocks_lacking: dict[str, bool] = field(default_factory=dict) prefill_step_start: float | None = None decode_step_start: float | None = None - + +@dataclass class MMEncodeRequest: """Lightweight descriptor for a waiting MM request that should be pre-encoded before its Spyre prefill step begins.""" @@ -1175,4 +1176,4 @@ def make_stats(self, *args, **kwargs) -> SchedulerStats | None: self.pause_events = 0 self.resume_events = 0 - return base_stats \ No newline at end of file + return base_stats From 988c4820ae9e99fab98f0f9d3faa26db01e3b21c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Mon, 10 Aug 2026 11:23:45 +0000 Subject: [PATCH 089/106] test fix multimodal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- tests/v1/core/test_scheduler_mm_encoding.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/v1/core/test_scheduler_mm_encoding.py b/tests/v1/core/test_scheduler_mm_encoding.py index 7e6a80d26..9291d24e8 100644 --- a/tests/v1/core/test_scheduler_mm_encoding.py +++ b/tests/v1/core/test_scheduler_mm_encoding.py @@ -68,6 +68,7 @@ def scheduler(): sched.paused_decoding_requests = [] sched.request_last_decode_step = {} sched.long_output_prio = False + sched._bench = None sched.pause_events = 0 sched.resume_events = 0 sched._get_required_blocks = lambda req, *a, **k: (0, 0) From 0eaf86d87f9967b904f7f3b09cad70dbdde0dc02 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 11 Aug 2026 10:07:44 +0000 Subject: [PATCH 090/106] add metrics descriptions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .claude/skills/add-bench-metric/SKILL.md | 57 ++++++++--- .../benchmarks/spyre_bench_serve.py | 96 +++++++++++++++++++ sendnn_inference/v1/core/scheduler.py | 25 ++++- 3 files changed, 162 insertions(+), 16 deletions(-) diff --git a/.claude/skills/add-bench-metric/SKILL.md b/.claude/skills/add-bench-metric/SKILL.md index 5da1845ed..4684d77b4 100644 --- a/.claude/skills/add-bench-metric/SKILL.md +++ b/.claude/skills/add-bench-metric/SKILL.md @@ -20,6 +20,7 @@ schedule() → SpyreBenchState (accumulate raw values) → patch_serving.py SSE injection (into final usage SSE chunk) → async_request_spyre_chat() → output.custom_metrics_dict → _print_spyre_section() + _inject_spyre_metrics_into_result_file() + → _METRIC_DESCRIPTIONS (--describe-metrics explanatory block) ``` **Key architecture facts:** @@ -47,21 +48,18 @@ Before writing any code, answer these questions (ask the user if unclear): **File**: `sendnn_inference/v1/core/scheduler.py` -Add your field to the dataclass at the top of the file: +Add your field to the dataclass at the top of the file. **Every field carries a one- or two-line comment** saying what it holds — match that convention; a bare field is a review comment waiting to happen. State the unit and, for per-step lists, what one entry corresponds to (one prefill step / one decode step / one pause interval). Read the neighbouring fields in the real file rather than copying the abridged sketch below. ```python @dataclass class SpyreBenchState: - chunk_latencies: dict[str, list[float]] = field(default_factory=dict) - chunk_start_times: dict[str, list[float]] = field(default_factory=dict) - arrival_ts: dict[str, float] = field(default_factory=dict) - first_scheduled_ts: dict[str, float] = field(default_factory=dict) + # ... existing fields, each with its own comment ... + + # Duration of each decode step the request took part in. decode_latencies: dict[str, list[float]] = field(default_factory=dict) - decode_start_times: dict[str, list[float]] = field(default_factory=dict) - tkvs: dict[str, list[int]] = field(default_factory=dict) - prefill_step_start: float | None = None - decode_step_start: float | None = None - # NEW — keyed by request_id, choose the container type for your cardinality: + + # NEW — keyed by request_id, choose the container type for your cardinality. + # One entry per ; . my_new_metric: dict[str, ] = field(default_factory=dict) ``` @@ -174,7 +172,35 @@ print("{:<40} {:<10}".format("My run-level total:", total)) --- -## Step 7 — Inject into the result JSON +## Step 7 — Add a description entry + +**File**: `sendnn_inference/benchmarks/spyre_bench_serve.py` + +Every printed metric has a short explanation in the module-level `_METRIC_DESCRIPTIONS` list, printed by `_print_metric_descriptions()` when the user passes `--describe-metrics`. Add an entry for your new metric: + +```python +_METRIC_DESCRIPTIONS: list[tuple[str, str]] = [ + # ... existing entries ... + ( + "My New Metric", # must match the _section() header (or the run-level print label) + "What it measures, in one or two sentences. State the sample granularity explicitly: " + "one sample per request / per prefill chunk / per decode step / per pause interval.", + ), +] +``` + +Rules for a good entry: + +- **The first element must match the printed header exactly** — the `_section()` header string from Step 6, or the label of the run-level print line. This is what lets a reader map a description back to the table above it. +- **Keep the list in printed order** so the descriptions read in the same sequence as the table (run-level scalar lines first, then the `_section()` metrics). +- **Always state the sample granularity.** It changes how the mean should be read: a per-chunk metric is dominated by long-prompt requests, and a per-step metric batched across requests records the same value for every participant in that step. +- **Say what the measurement brackets**, not just its name — e.g. that a step latency is measured from the end of `schedule()` to `update_from_output()` and therefore covers the whole scheduler → executor → model round trip, not just the forward pass. +- **Mention a non-obvious exclusion or convention in one sentence** if there is one (e.g. pause metrics exclude an interval still open when the request finishes; prefix-cache hit is measured in chunks rather than tokens). +- Avoid characters that wrap badly: `textwrap.wrap()` will break on a lone ` - `, leaving a dangling hyphen at end of line. Write `(1 - a/b)` rather than `1 - a / b`. + +--- + +## Step 8 — Inject into the result JSON **File**: `sendnn_inference/benchmarks/spyre_bench_serve.py` @@ -192,7 +218,7 @@ Use a `spyre_` prefix so the key is clearly SenDNN-owned in the vllm result JSON --- -## Step 8 — Add tests +## Step 9 — Add tests **File**: `tests/benchmarks/test_bench_metrics.py` @@ -255,19 +281,22 @@ Use a `spyre_` prefix so the key is clearly SenDNN-owned in the vllm result JSON 6. **Update `test_print_spyre_section_output`** (if you added a new `_section()` call): Add `out.assert_contains("")` for the section separator, and add the label string (third argument to `_section()`) to the `for label in (...)` loop so mean/median/percentile lines are asserted. `assert_all_lines_covered()` is called at the end of the test and will fail if any output line was not covered by an assertion — the error message lists the exact uncovered lines. +7. **Cover the description entry.** `_METRIC_DESCRIPTIONS` must stay in sync with what `_print_spyre_section()` actually prints, and nothing enforces that automatically yet. Add (or extend) a test asserting that every printed section header has a matching description entry — walking the `_section()` header strings and checking each appears as a `_METRIC_DESCRIPTIONS[i][0]`. Without this, a new metric silently ships with no description. + --- ## Checklist Before declaring done: -- [ ] Field added to `SpyreBenchState` with the correct container type (scalar, list, dict-of-list…) +- [ ] Field added to `SpyreBenchState` with the correct container type (scalar, list, dict-of-list…) and a short comment stating unit and what one entry corresponds to - [ ] Populated in `update_from_output()` (or `schedule()`) under `if self._bench is not None:` guard - [ ] Retrieved via `.pop()` in `get_and_clear_chunk_stats()` with a safe default if absent - [ ] Key added to the `spyre_data` dict in `_free_request()` with a safe fallback when `chunk_stats is None` - [ ] `_print_spyre_section()` updated with a new `_section()` call or run-level print line +- [ ] `_METRIC_DESCRIPTIONS` entry added, header matching the printed section exactly, in printed order, stating sample granularity - [ ] `_inject_spyre_metrics_into_result_file()` updated with a new `result["spyre_..."]` key -- [ ] Tests updated: `FAKE_METRICS`, `test_inject_adds_spyre_keys`, `test_inject_values_correct`, `_BENCH_FIXTURE`, `_EXPECTED_RESULT` (drives `test_get_and_clear_returns_correct_dict`), `test_scheduler_bench_metrics_accumulated` adapted accordingly, `test_print_spyre_section_output` updated if a new `_section()` was added +- [ ] Tests updated: `FAKE_METRICS`, `test_inject_adds_spyre_keys`, `test_inject_values_correct`, `_BENCH_FIXTURE`, `_EXPECTED_RESULT` (drives `test_get_and_clear_returns_correct_dict`), `test_scheduler_bench_metrics_accumulated` adapted accordingly, `test_print_spyre_section_output` updated if a new `_section()` was added, and a check that the new metric has a `_METRIC_DESCRIPTIONS` entry - [ ] No changes to `patch_serving.py`, `spyre_request_func.py`, or any model runner file - [ ] No changes to `SpyreRequestMetrics` in `stats_logger.py` (not part of the active pipeline) - [ ] All new tracking is gated by `self._bench is not None` (enforced by `SpyreBenchState` being `None` when env var is off) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index c98c24ee5..9f6882f28 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -109,6 +109,15 @@ def _build_parser() -> FlexibleArgumentParser: "Requires --save-result and SENDNN_INFERENCE_BENCH_METRICS_ENABLED on the server." ), ) + parser.add_argument( + "--describe-metrics", + action="store_true", + default=False, + help=( + "After the SenDNN metrics table, print a short description of every " + "metric section explaining what it measures and its sample granularity." + ), + ) parser.add_argument( "--decode-thresholds", type=str, @@ -125,6 +134,90 @@ def _build_parser() -> FlexibleArgumentParser: return parser +# Short explanation of every metric printed by _print_spyre_section, shown when +# --describe-metrics is passed. Ordered to match the printed output. +_METRIC_DESCRIPTIONS: list[tuple[str, str]] = [ + ( + "Total prefill chunks processed", + "Total number of chunked prefills executed by the server across all requests.", + ), + ( + "Requests blocked by missing KV blocks", + "Number of requests that were held back at least once because not enough free " + "KV-cache blocks were available. A non-zero count signals KV-cache pressure, " + "not an error.", + ), + ( + "Queue Wait Time", + "Time from when the API server received the request until its first prefill step " + "started. Unlike vLLM's own queue time this includes the API-server -> engine-core " + "hop, so it adds up with prefill and decode latencies to reconstruct TTFT. " + "One sample per request.", + ), + ( + "Chunked Prefill Count", + "Number of prefill chunks actually executed for a request. Lower than " + "ceil(prompt_len / chunk_size) when a prefix-cache hit lets whole chunks be " + "skipped. One sample per request.", + ), + ( + "Chunked Prefill Latency", + "Wall-clock duration of a single prefill step. One sample per chunk.", + ), + ( + "Decode Step Latency", + "Wall-clock duration of a single decode step. One sample per decode step per request.", + ), + ( + "Prefix Cache Hit", + "Fraction of a request's prefill chunks that were skipped thanks to the prefix " + "cache, computed as (1 - executed_chunks/expected_chunks). Measured in chunks " + "rather than tokens because Spyre only ever skips whole chunks. " + "One sample per request.", + ), + ( + "Left Padding Blocks", + "Per decode step, the number of KV-cache blocks of left padding a request carries " + "because it is shorter than the batch's longest sequence " + "(ceil(tkv/block_size) - ceil(computed_tokens/block_size)). Zero for the longest " + "request; high values mean the batch mixes very different sequence lengths and " + "wastes compute on padding.", + ), + ( + "Pause Latency", + "Duration of a single pause interval - the time a decoding request spent evicted " + "from the running batch because the batch TKV limit could not accommodate it. " + "One sample per pause interval. A request still paused when it finishes has that " + "final open interval excluded from all three pause metrics.", + ), + ( + "Number of Pauses", + "How many times a request was paused and later resumed. One sample per request " + "(including 0). A request still paused when it finishes has that final open " + "interval excluded from all three pause metrics.", + ), + ( + "Total Time Paused", + "Total time a request spent paused over its lifetime (the sum of its pause " + "intervals). One sample per request, including 0. Read alongside Pause Latency to " + "distinguish many short pauses from one long one. A request still paused when it " + "finishes has that final open interval excluded from all three pause metrics.", + ), +] + + +def _print_metric_descriptions(width: int = 100) -> None: + """Print a short explanation of every metric section printed above.""" + import textwrap + + print("{s:{c}^{n}}".format(s=" Metric Descriptions ", n=50, c="=")) + for header, description in _METRIC_DESCRIPTIONS: + print(f"{header}:") + for line in textwrap.wrap(description, width=width - 2): + print(f" {line}") + print() + + def _print_spyre_section( metrics_list: list[dict[str, Any]], selected_percentiles: list[float], @@ -349,6 +442,9 @@ def main() -> None: print("{s:{c}^{n}}".format(s=" SenDNN Metrics ", n=50, c="=")) _print_spyre_section(_spyre_metrics_collected, selected_percentiles) + if getattr(args, "describe_metrics", False): + _print_metric_descriptions() + _inject_spyre_metrics_into_result_file(args, _spyre_metrics_collected, run_started_at) if getattr(args, "detailed_timeline", False): diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 7525a5841..347c9c77d 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -30,17 +30,38 @@ class SpyreBenchState: """Bench-metrics-only state for tracking per-request and per-step timing. Only instantiated when SENDNN_INFERENCE_BENCH_METRICS_ENABLED is set.""" - chunk_latencies: dict[str, list[float]] = field(default_factory=dict) + # All dicts below are keyed by request_id. Per-step lists grow by one entry per + # step the request took part in. + + # Timestamp the API server stamped on the request on receipt (request.arrival_time). arrival_ts: dict[str, float] = field(default_factory=dict) + # Timestamp of the request's first prefill step. first_scheduled_ts: dict[str, float] = field(default_factory=dict) + # Duration of each prefill step the request took part in. + chunk_latencies: dict[str, list[float]] = field(default_factory=dict) + # Start timestamp of each prefill step. chunk_start_times: dict[str, list[float]] = field(default_factory=dict) + # Duration of each decode step the request took part in. decode_latencies: dict[str, list[float]] = field(default_factory=dict) + # Start timestamp of each decode step. decode_start_times: dict[str, list[float]] = field(default_factory=dict) + # Batch-wide sequence length (tkv) at each step the request took part in. tkvs: dict[str, list[int]] = field(default_factory=dict) + # Per decode step, KV blocks of pure left padding the request carries because it is + # shorter than the batch's longest sequence. Zero for the longest request. left_padding_blocks: dict[str, list[int]] = field(default_factory=dict) - pause_start_times: dict[str, list[float]] = field(default_factory=dict) + # Duration of each *completed* pause interval (paused -> running). An interval still + # open when the request finishes is never recorded here. pause_latencies: dict[str, list[float]] = field(default_factory=dict) + # Start timestamp of each pause interval (running -> paused). + pause_start_times: dict[str, list[float]] = field(default_factory=dict) + # True if the request was ever held back because too few free KV blocks were + # available. Signals KV-cache pressure, not an error. blocks_lacking: dict[str, bool] = field(default_factory=dict) + + # Step-start timestamps for the in-flight step, set at the end of schedule() and + # consumed in update_from_output() to close the duration. Exactly one is non-None + # at a time, depending on whether the current step is a prefill or a decode. prefill_step_start: float | None = None decode_step_start: float | None = None From f88f697c3343e4833a282139728aefcac7f93bc6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 11 Aug 2026 12:54:08 +0000 Subject: [PATCH 091/106] add prefilling breakdown metrics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_bench_serve.py | 31 +++++++++++++ sendnn_inference/v1/core/scheduler.py | 14 ++++++ tests/benchmarks/test_bench_metrics.py | 44 +++++++++++++++++++ 3 files changed, 89 insertions(+) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 9f6882f28..b72e81b35 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -164,6 +164,25 @@ def _build_parser() -> FlexibleArgumentParser: "Chunked Prefill Latency", "Wall-clock duration of a single prefill step. One sample per chunk.", ), + ( + "Prefill Phase Time", + "Elapsed time from the start of a request's first prefill step to the end of its " + "last one, which is when its first token is produced. Starts after the request " + "leaves the waiting queue, so 'Queue Wait Time' does not contribute. One sample " + "per request.", + ), + ( + "Time Spent Prefilling", + "Time a request spent actually executing prefill steps, i.e. the sum of its " + "Chunked Prefill Latency samples. One sample per request.", + ), + ( + "Prefill Phase Idle Time", + "Time a request spent inside its prefill phase without prefilling, computed as " + "(Prefill Phase Time - Time Spent Prefilling). Covers pausing and any other " + "reason the request made no prefill progress, but excludes the initial queue wait. " + "One sample per request.", + ), ( "Decode Step Latency", "Wall-clock duration of a single decode step. One sample per decode step per request.", @@ -261,6 +280,12 @@ def _section(header: str, values: list[float], label: str) -> None: m["prefix_cache_hit_pct"] * 100 for m in metrics_list if "prefix_cache_hit_pct" in m ] + prefill_elapsed_ms = [ + m["prefill_elapsed_s"] * 1000 for m in metrics_list if "prefill_elapsed_s" in m + ] + prefill_busy_ms = [m["prefill_busy_s"] * 1000 for m in metrics_list if "prefill_busy_s" in m] + prefill_idle_ms = [m["prefill_idle_s"] * 1000 for m in metrics_list if "prefill_idle_s" in m] + left_padding_blocks = [v for m in metrics_list for v in m.get("left_padding_blocks", [])] pause_lats_ms = [lat * 1000 for m in metrics_list for lat in m.get("pause_latencies_s", [])] pause_counts = [float(len(m.get("pause_latencies_s", []))) for m in metrics_list] @@ -269,6 +294,9 @@ def _section(header: str, values: list[float], label: str) -> None: _section("Queue Wait Time", queue_times_ms, "Queue Wait Time (ms)") _section("Chunked Prefill Count", num_chunks_list, "Num Chunked Prefills") _section("Chunked Prefill Latency", chunk_lats_ms, "Chunk Prefill Latency (ms)") + _section("Prefill Phase Time", prefill_elapsed_ms, "Prefill Phase Time (ms)") + _section("Time Spent Prefilling", prefill_busy_ms, "Time Spent Prefilling (ms)") + _section("Prefill Phase Idle Time", prefill_idle_ms, "Prefill Phase Idle Time (ms)") _section("Decode Step Latency", decode_lats_ms, "Decode Step Latency (ms)") _section("Prefix Cache Hit", cache_hit_pcts, "Prefix Cache Hit (%)") _section("Left Padding Blocks", left_padding_blocks, "Left Padding Blocks") @@ -348,6 +376,9 @@ def _inject_spyre_metrics_into_result_file( result["spyre_decode_latencies_s"] = [m.get("decode_latencies_s", []) for m in metrics_list] result["spyre_decode_start_times_s"] = [m.get("decode_start_times_s", []) for m in metrics_list] result["spyre_tkvs"] = [m.get("tkvs", []) for m in metrics_list] + result["spyre_prefill_elapsed_s"] = [m.get("prefill_elapsed_s", 0.0) for m in metrics_list] + result["spyre_prefill_busy_s"] = [m.get("prefill_busy_s", 0.0) for m in metrics_list] + result["spyre_prefill_idle_s"] = [m.get("prefill_idle_s", 0.0) for m in metrics_list] result["spyre_prefix_cache_hit_pct"] = [ m.get("prefix_cache_hit_pct", 0.0) for m in metrics_list ] diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 347c9c77d..b024577ab 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -437,6 +437,17 @@ def _free_request(self, request, delay_free_blocks: bool = False): num_expected = math.ceil(request.num_prompt_tokens / self.chunk_size) num_skipped = max(0, num_expected - num_executed) cache_hit_pct = num_skipped / num_expected if num_expected > 0 else 0.0 + + # Prefill-phase breakdown, all derived from the per-chunk timings above. + chunk_lats = chunk_stats["chunk_prefill_latencies_s"] if chunk_stats else [] + chunk_starts = chunk_stats["chunk_prefill_start_times_s"] if chunk_stats else [] + prefill_busy_s = sum(chunk_lats) + if chunk_lats and chunk_starts and first_ts is not None: + prefill_elapsed_s = chunk_starts[-1] + chunk_lats[-1] - first_ts + else: + prefill_elapsed_s = 0.0 + # Guard against a negative result from clock jitter between the two samples. + prefill_idle_s = max(0.0, prefill_elapsed_s - prefill_busy_s) spyre_data = { "queued_time_s": queued_time_s, "num_chunked_prefills": num_executed, @@ -449,6 +460,9 @@ def _free_request(self, request, delay_free_blocks: bool = False): "decode_latencies_s": chunk_stats["decode_latencies_s"] if chunk_stats else [], "decode_start_times_s": chunk_stats["decode_start_times_s"] if chunk_stats else [], "tkvs": chunk_stats["tkvs"] if chunk_stats else [], + "prefill_elapsed_s": prefill_elapsed_s, + "prefill_busy_s": prefill_busy_s, + "prefill_idle_s": prefill_idle_s, "prefix_cache_hit_pct": cache_hit_pct, "left_padding_blocks": chunk_stats["left_padding_blocks"] if chunk_stats else [], "pause_latencies_s": chunk_stats["pause_latencies_s"] if chunk_stats else [], diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index 9ecb97cee..e9ac08653 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -38,6 +38,9 @@ "decode_latencies_s": [88888.8, 0.000005, 44444.4], "decode_start_times_s": [5000.0, 5088888.8, 5088888.8], "tkvs": [512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072], + "prefill_elapsed_s": 7777777.7, + "prefill_busy_s": 2250.9, + "prefill_idle_s": 7775526.8, "prefix_cache_hit_pct": 0.25, "left_padding_blocks": [2, 0, 1], "pause_latencies_s": [0.5, 1.2], @@ -52,6 +55,9 @@ "decode_latencies_s": [0.000002, 77777.7], "decode_start_times_s": [1112345.5, 1112345.5], "tkvs": [256, 512, 1024, 2048, 4096], + "prefill_elapsed_s": 33333333.3, + "prefill_busy_s": 0.0000004, + "prefill_idle_s": 33333333.2999996, "prefix_cache_hit_pct": 0.0, "left_padding_blocks": [3, 1], "pause_latencies_s": [0.3], @@ -213,6 +219,9 @@ def test_print_spyre_section_output(capsys): out.assert_contains("Queue Wait Time") out.assert_contains("Chunked Prefill Count") out.assert_contains("Chunked Prefill Latency") + out.assert_contains("Prefill Phase Time") + out.assert_contains("Time Spent Prefilling") + out.assert_contains("Prefill Phase Idle Time") out.assert_contains("Decode Step Latency") out.assert_contains("Prefix Cache Hit") out.assert_contains("Left Padding Blocks") @@ -224,6 +233,9 @@ def test_print_spyre_section_output(capsys): "Queue Wait Time (ms)", "Num Chunked Prefills", "Chunk Prefill Latency (ms)", + "Prefill Phase Time (ms)", + "Time Spent Prefilling (ms)", + "Prefill Phase Idle Time (ms)", "Decode Step Latency (ms)", "Prefix Cache Hit (%)", "Left Padding Blocks", @@ -239,6 +251,38 @@ def test_print_spyre_section_output(capsys): out.assert_all_lines_covered() +@pytest.mark.cpu +def test_every_printed_metric_has_a_description(capsys): + """Every metric _print_spyre_section prints must have a _METRIC_DESCRIPTIONS entry, + and every entry must correspond to something printed. Fails when a new metric ships + without a description (or a description outlives its metric).""" + from sendnn_inference.benchmarks.spyre_bench_serve import _METRIC_DESCRIPTIONS + + _print_spyre_section(FAKE_METRICS, SELECTED_PERCENTILES) + lines = [line.strip() for line in capsys.readouterr().out.splitlines() if line.strip()] + + # Section headers are printed centred in dashes: "----- Queue Wait Time -----". + # Run-level scalars are printed as "Some label: ". + printed: list[str] = [] + for line in lines: + if set(line) == {"="}: + continue + if line.startswith("-"): + printed.append(line.strip("- ")) + elif ":" in line: + label = line.split(":", 1)[0].strip() + # Skip the per-section Mean/Median/PXX rows; they belong to the section above. + if not label.startswith(("Mean ", "Median ", "P")): + printed.append(label) + + described = [header for header, _ in _METRIC_DESCRIPTIONS] + assert printed == described, ( + "Printed metrics and _METRIC_DESCRIPTIONS disagree (order matters).\n" + f" printed : {printed}\n" + f" described: {described}" + ) + + @pytest.mark.cpu def test_print_noop_when_empty(capsys): _print_spyre_section([], SELECTED_PERCENTILES) From 10822872b7303fad91a5c58885cc51b75d6121f3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Tue, 11 Aug 2026 13:20:06 +0000 Subject: [PATCH 092/106] bench metrics description: move it to file instead of print MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .claude/skills/add-bench-metric/SKILL.md | 4 +- .../benchmarks/spyre_bench_serve.py | 39 ++++++++++++------- 2 files changed, 28 insertions(+), 15 deletions(-) diff --git a/.claude/skills/add-bench-metric/SKILL.md b/.claude/skills/add-bench-metric/SKILL.md index 4684d77b4..da159a27f 100644 --- a/.claude/skills/add-bench-metric/SKILL.md +++ b/.claude/skills/add-bench-metric/SKILL.md @@ -20,7 +20,7 @@ schedule() → SpyreBenchState (accumulate raw values) → patch_serving.py SSE injection (into final usage SSE chunk) → async_request_spyre_chat() → output.custom_metrics_dict → _print_spyre_section() + _inject_spyre_metrics_into_result_file() - → _METRIC_DESCRIPTIONS (--describe-metrics explanatory block) + → _METRIC_DESCRIPTIONS (--describe-metrics description file) ``` **Key architecture facts:** @@ -176,7 +176,7 @@ print("{:<40} {:<10}".format("My run-level total:", total)) **File**: `sendnn_inference/benchmarks/spyre_bench_serve.py` -Every printed metric has a short explanation in the module-level `_METRIC_DESCRIPTIONS` list, printed by `_print_metric_descriptions()` when the user passes `--describe-metrics`. Add an entry for your new metric: +Every printed metric has a short explanation in the module-level `_METRIC_DESCRIPTIONS` list, written to `sendnn_bench_metrics_description.txt` by `_write_metric_descriptions()` when the user passes `--describe-metrics`. Add an entry for your new metric: ```python _METRIC_DESCRIPTIONS: list[tuple[str, str]] = [ diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index b72e81b35..f5cddac96 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -114,8 +114,9 @@ def _build_parser() -> FlexibleArgumentParser: action="store_true", default=False, help=( - "After the SenDNN metrics table, print a short description of every " - "metric section explaining what it measures and its sample granularity." + "Write a 'sendnn_bench_metrics_description.txt' file (in --result-dir, or the " + "current directory) describing every metric section: what it measures and its " + "sample granularity." ), ) parser.add_argument( @@ -134,8 +135,9 @@ def _build_parser() -> FlexibleArgumentParser: return parser -# Short explanation of every metric printed by _print_spyre_section, shown when -# --describe-metrics is passed. Ordered to match the printed output. +# Short explanation of every metric printed by _print_spyre_section, written to +# _METRIC_DESCRIPTION_FILENAME when --describe-metrics is passed. Ordered to match +# the printed output. _METRIC_DESCRIPTIONS: list[tuple[str, str]] = [ ( "Total prefill chunks processed", @@ -225,16 +227,27 @@ def _build_parser() -> FlexibleArgumentParser: ] -def _print_metric_descriptions(width: int = 100) -> None: - """Print a short explanation of every metric section printed above.""" - import textwrap +_METRIC_DESCRIPTION_FILENAME = "sendnn_bench_metrics_description.txt" - print("{s:{c}^{n}}".format(s=" Metric Descriptions ", n=50, c="=")) + +def _write_metric_descriptions(result_dir: str | None = None) -> None: + """Write a short explanation of every metric section to a text file. + + Each description is a single unwrapped line so it reflows in any editor.""" + file_path = os.path.join(result_dir or ".", _METRIC_DESCRIPTION_FILENAME) + + lines = ["SenDNN benchmark metric descriptions", ""] for header, description in _METRIC_DESCRIPTIONS: - print(f"{header}:") - for line in textwrap.wrap(description, width=width - 2): - print(f" {line}") - print() + lines.append(f"{header}:") + lines.append(description) + lines.append("") + + try: + with open(file_path, "w", encoding="utf-8") as fh: + fh.write("\n".join(lines)) + logger.info("SenDNN metric descriptions written to %s", file_path) + except OSError as exc: + logger.warning("Failed to write SenDNN metric descriptions to %s: %s", file_path, exc) def _print_spyre_section( @@ -474,7 +487,7 @@ def main() -> None: _print_spyre_section(_spyre_metrics_collected, selected_percentiles) if getattr(args, "describe_metrics", False): - _print_metric_descriptions() + _write_metric_descriptions(getattr(args, "result_dir", None)) _inject_spyre_metrics_into_result_file(args, _spyre_metrics_collected, run_started_at) From dc1a72dcba7c4601b977b75baa60dcc7f47efd10 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 12 Aug 2026 13:19:19 +0000 Subject: [PATCH 093/106] bugfix: time_ts was using end of first chunk instead of start MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic Co-authored-by: Yannick Schnider --- .../benchmarks/spyre_bench_serve.py | 8 ++--- sendnn_inference/v1/core/scheduler.py | 31 +++++++++---------- tests/benchmarks/test_bench_metrics.py | 14 ++++----- 3 files changed, 24 insertions(+), 29 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index f5cddac96..92e6afe4e 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -151,10 +151,10 @@ def _build_parser() -> FlexibleArgumentParser: ), ( "Queue Wait Time", - "Time from when the API server received the request until its first prefill step " - "started. Unlike vLLM's own queue time this includes the API-server -> engine-core " - "hop, so it adds up with prefill and decode latencies to reconstruct TTFT. " - "One sample per request.", + "Time from when the API server received the request until the start of its first " + "prefill step. Unlike vLLM's own queue time this includes the API-server -> " + "engine-core hop, so it adds up with prefill and decode latencies to reconstruct " + "TTFT. One sample per request.", ), ( "Chunked Prefill Count", diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index b024577ab..e5aa217cb 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -35,8 +35,6 @@ class SpyreBenchState: # Timestamp the API server stamped on the request on receipt (request.arrival_time). arrival_ts: dict[str, float] = field(default_factory=dict) - # Timestamp of the request's first prefill step. - first_scheduled_ts: dict[str, float] = field(default_factory=dict) # Duration of each prefill step the request took part in. chunk_latencies: dict[str, list[float]] = field(default_factory=dict) # Start timestamp of each prefill step. @@ -381,8 +379,7 @@ def _bench_update_from_output( self._bench.decode_step_start = None for req in self.ongoing_prefills: - if req.request_id not in self._bench.first_scheduled_ts: - self._bench.first_scheduled_ts[req.request_id] = now + if req.request_id not in self._bench.arrival_ts: self._bench.arrival_ts[req.request_id] = req.arrival_time def get_and_clear_chunk_stats(self, req_id: str) -> dict | None: @@ -421,18 +418,7 @@ def _free_request(self, request, delay_free_blocks: bool = False): if self._bench is not None: req_id = request.request_id chunk_stats = self.get_and_clear_chunk_stats(req_id) - first_ts = self._bench.first_scheduled_ts.pop(req_id, None) arrival_ts = self._bench.arrival_ts.pop(req_id, None) - # NOTE: queued_time_s looks like a duplicate of FinishedRequestStats.queued_time, but - # it is not. FinishedRequestStats.queued_time is (scheduled_ts - queued_ts): both - # timestamps are recorded inside the engine core, so it misses the time the request - # spent in transit from the API server to the engine (IPC hop). Here we use the stamp - # of the API server on receipt (request.arrival_time), which gives a complete client- - # visible queue wait that adds up cleanly with prefill and decode latencies when - # reconstructing TTFT. - queued_time_s = ( - (first_ts - arrival_ts) if first_ts is not None and arrival_ts is not None else 0.0 - ) num_executed = chunk_stats["num_chunked_prefills"] if chunk_stats else 0 num_expected = math.ceil(request.num_prompt_tokens / self.chunk_size) num_skipped = max(0, num_expected - num_executed) @@ -442,12 +428,23 @@ def _free_request(self, request, delay_free_blocks: bool = False): chunk_lats = chunk_stats["chunk_prefill_latencies_s"] if chunk_stats else [] chunk_starts = chunk_stats["chunk_prefill_start_times_s"] if chunk_stats else [] prefill_busy_s = sum(chunk_lats) - if chunk_lats and chunk_starts and first_ts is not None: - prefill_elapsed_s = chunk_starts[-1] + chunk_lats[-1] - first_ts + if chunk_lats and chunk_starts: + prefill_elapsed_s = chunk_starts[-1] + chunk_lats[-1] - chunk_starts[0] else: prefill_elapsed_s = 0.0 # Guard against a negative result from clock jitter between the two samples. prefill_idle_s = max(0.0, prefill_elapsed_s - prefill_busy_s) + + # NOTE: queued_time_s looks like a duplicate of FinishedRequestStats.queued_time, but + # it is not. FinishedRequestStats.queued_time is (scheduled_ts - queued_ts): both + # timestamps are recorded inside the engine core, so it misses the time the request + # spent in transit from the API server to the engine (IPC hop). Here we use the stamp + # of the API server on receipt (request.arrival_time), which gives a complete client- + # visible queue wait that adds up cleanly with prefill and decode latencies when + # reconstructing TTFT. + queued_time_s = ( + (chunk_starts[0] - arrival_ts) if chunk_starts and arrival_ts is not None else 0.0 + ) spyre_data = { "queued_time_s": queued_time_s, "num_chunked_prefills": num_executed, diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index e9ac08653..3954b7f99 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -4,8 +4,8 @@ Tests cover four layers: 1. _inject_spyre_metrics_into_result_file — JSON result file injection 2. _print_spyre_section — stdout output format - 3. ChunkedPrefillSpyreScheduler accumulation — _chunk_latencies / _arrival_ts / - _first_scheduled_ts, and get_and_clear_chunk_stats + 3. ChunkedPrefillSpyreScheduler accumulation — SpyreBenchState fields + (chunk_latencies, arrival_ts, …) and get_and_clear_chunk_stats 4. async_request_spyre_chat — client-side SSE parsing """ @@ -330,7 +330,6 @@ def _make_bare_scheduler(): _BENCH_FIXTURE: dict[str, Any] = { "chunk_latencies": [88888.8, 0.000005], "arrival_ts": 1000.0, - "first_scheduled_ts": 1001.0, "chunk_start_times": [1000.0, 1088888.8], "decode_latencies": [0.1, 0.2], "decode_start_times": [2000.0, 2000.1], @@ -450,9 +449,8 @@ def test_scheduler_bench_metrics_accumulated( available_blocks, ): """Two requests with prompts longer than max_num_batched_tokens each trigger - multiple prefill chunks. Verify that _chunk_latencies, _arrival_ts, and - _first_scheduled_ts are populated correctly, and cleared once the request - finishes via _free_request.""" + multiple prefill chunks. Verify that the SpyreBenchState fields are populated + correctly and cleared once the request finishes via _free_request.""" from llm_cache import get_cached_engine from scheduling_utils import create_request_for_scheduler_test, random_prompt @@ -509,8 +507,8 @@ def _capturing_free(self, request, delay_free_blocks=False): req_id = request.request_id bench = self._bench # Snapshot all dict fields dynamically so new metrics are captured automatically. - # List-valued dicts are copied; scalar-valued dicts (arrival_ts, first_scheduled_ts) - # are stored as-is so truthiness checks work correctly. + # List-valued dicts are copied; scalar-valued dicts (arrival_ts) are stored + # as-is so truthiness checks work correctly. if bench: snap = {} for f in dc_fields(bench): From 65af15433b2ebd7670d54a971f84e22323f2bb66 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 12 Aug 2026 13:20:31 +0000 Subject: [PATCH 094/106] more thorough prefill phase breakdown test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- tests/benchmarks/test_bench_metrics.py | 25 ++++++++++++-- tests/hf_cache.json | 46 ++++++++++++++++++++++++++ 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index 3954b7f99..a31e07be4 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -497,8 +497,10 @@ def test_scheduler_bench_metrics_accumulated( generate_hf_results=False, ) - # Capture metrics just before _free_request clears them + # Capture metrics just before _free_request clears them, plus the derived + # __spyre__ payload it returns (prefill_elapsed_s & co. only exist there). captured: dict[str, dict] = {} + captured_spyre: dict[str, dict] = {} original_free = scheduler.__class__._free_request def _capturing_free(self, request, delay_free_blocks=False): @@ -520,7 +522,10 @@ def _capturing_free(self, request, delay_free_blocks=False): captured[req_id] = snap else: captured[req_id] = {} - return original_free(self, request, delay_free_blocks) + kv_params = original_free(self, request, delay_free_blocks) + if kv_params and "__spyre__" in kv_params: + captured_spyre[req_id] = dict(kv_params["__spyre__"]) + return kv_params scheduler._free_request = _capturing_free.__get__(scheduler) @@ -599,7 +604,21 @@ def _capturing_free(self, request, delay_free_blocks=False): assert isinstance(tkv, int) and tkv > 0, f"req {req_id}: non-positive tkv {tkv}" assert info["arrival_ts"] is not None, f"req {req_id}: arrival_ts not set" - assert info["first_scheduled_ts"] is not None, f"req {req_id}: first_scheduled_ts not set" + + # Derived prefill-phase breakdown (only present in the __spyre__ payload). + # The phase brackets the same steps as prefill_busy_s plus the gaps between + # them, so it can never be shorter. + assert req_id in captured_spyre, f"req {req_id}: no __spyre__ payload captured" + derived = captured_spyre[req_id] + elapsed = derived["prefill_elapsed_s"] + busy = derived["prefill_busy_s"] + assert elapsed >= busy - 1e-6, ( + f"req {req_id}: prefill_elapsed_s {elapsed} < prefill_busy_s {busy} — " + f"elapsed/busy clocks disagree" + ) + assert derived["prefill_idle_s"] == pytest.approx(max(0.0, elapsed - busy), abs=1e-9), ( + f"req {req_id}: prefill_idle_s {derived['prefill_idle_s']} != elapsed - busy" + ) # pause_latencies: list of pause durations (may be absent/None if no pausing occurred) pause_lats = info["pause_latencies"] or [] diff --git a/tests/hf_cache.json b/tests/hf_cache.json index 368ba456a..870ee8a19 100644 --- a/tests/hf_cache.json +++ b/tests/hf_cache.json @@ -634,6 +634,52 @@ "tokens": [ "formed", "URLException" ], "logprobs": [ -4.1111040115356445, -5.85740327835083 ] } + }, + "__tokens__41511_37260_20675_12728_25134_19906_38530_14911_23429_28678_44642_24810_13855_37154_30398_12315_44722_48312_39829_44349_15247_35878_44186_33624_23210": { + "7": { + "text": " Latin-1 Latin-", + "token_ids": [ 19190, 266, 31, 35, 19190, 266, 31 ], + "tokens": [ " Lat", "in", "-", "1", " Lat", "in", "-" ], + "logprobs": [ -3.8131632804870605, -1.6206474304199219, -2.7686638832092285, -2.111767530441284, -2.923048496246338, -0.0068110208958387375, -0.34963783621788025 ] + } + }, + "__tokens__6606_41659_37546_12539_24355_22097_32032_38773_4614_1394_41085_21274_37473_104_21895_35470_11246_46468_44313_1504_1251_26615_46167_18740_10648": { + "6": { + "text": "nique(s)\n\n", + "token_ids": [ 37881, 26, 101, 27, 203, 203 ], + "tokens": [ "nique", "(", "s", ")", "\n", "\n" ], + "logprobs": [ -0.3878551423549652, -3.3403029441833496, -3.473588705062866, -0.7978524565696716, -2.7161076068878174, -0.6159165501594543 ] + } + }, + "__tokens__46997_46594_2780_4173_41072_36179_32923_15148_29788_29830_28571_7786_21171_19346_35542_48904_46671_26751_21869_13187_1766_1350_22854_15656_18681_43839_25845_27554_11608_1173_15984_6720_25082_49094_33157_8940_43927_39168_36102_44567_37502_38823_17392_48223_47286_7924_37066_35156_22682_26072_24089_45463_24621_40877_17399_43400_44228_22663_27908_45242_35580_23921_10904_15960_34390_8164": { + "3": { + "text": "ol -o", + "token_ids": [ 362, 429, 97 ], + "tokens": [ "ol", " -", "o" ], + "logprobs": [ -4.425321102142334, -4.811501979827881, -2.9282565116882324 ] + }, + "10": { + "text": "ol -o -o -o -o -", + "token_ids": [ 362, 429, 97, 429, 97, 429, 97, 429, 97, 429 ], + "tokens": [ "ol", " -", "o", " -", "o", " -", "o", " -", "o", " -" ], + "logprobs": [ -4.425321102142334, -4.811501979827881, -2.9282565116882324, -2.2208919525146484, -0.7492299675941467, -0.3586389422416687, -0.17332826554775238, -0.3304527997970581, -0.11323630064725876, -0.3726454973220825 ] + } + }, + "__tokens__41511_37260_20675_12728_25134_19906_38530_14911_23429_28678_44642_24810_13855_37154_30398": { + "11": { + "text": "enance.\n\n# 1. 2.", + "token_ids": [ 12988, 32, 203, 203, 21, 225, 35, 32, 225, 36, 32 ], + "tokens": [ "enance", ".", "\n", "\n", "#", " ", "1", ".", " ", "2", "." ], + "logprobs": [ -0.7929723858833313, -2.5143587589263916, -3.4940693378448486, -0.8417129516601562, -2.1874706745147705, -3.206106185913086, -1.5070531368255615, -0.6664941906929016, -2.186004638671875, -1.0107399225234985, -0.24572384357452393 ] + } + }, + "__tokens__6606_41659_37546_12539_24355_22097_32032_38773_4614_1394_41085_21274_37473_104_21895": { + "13": { + "text": "'';\n\n//\n// //\n// //", + "token_ids": [ 25, 920, 203, 203, 306, 203, 306, 225, 434, 203, 306, 225, 434 ], + "tokens": [ "'", "';", "\n", "\n", "//", "\n", "//", " ", " //", "\n", "//", " ", " //" ], + "logprobs": [ -3.191380023956299, -2.213517904281616, -1.3364982604980469, -0.595153272151947, -2.902264356613159, -2.850527763366699, -0.6304447054862976, -1.338612675666809, -2.706216812133789, -2.0335094928741455, -0.69097501039505, -0.7241754531860352, -0.8122202157974243 ] + } } } } From f105f25a8b24023f1631bbf0a5c2b4ae9cb358cd Mon Sep 17 00:00:00 2001 From: Yannick Schnider Date: Wed, 12 Aug 2026 15:57:08 +0200 Subject: [PATCH 095/106] fix ordering Signed-off-by: Yannick Schnider --- .../benchmarks/spyre_bench_serve.py | 88 +++++++++++++------ 1 file changed, 60 insertions(+), 28 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 92e6afe4e..105e11541 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -36,6 +36,11 @@ _spyre_metrics_collected: list[dict[str, Any]] = [] _request_outputs_collected: list[dict[str, Any]] = [] +# Per-request join key. _spyre_metrics_collected fills in completion order while +# vllm's result arrays (ttfts/itls/start_times) are in submission order, so they +# are realigned by start_time at injection time rather than zipped positionally. +_SPYRE_START_TIME_KEY = "__spyre_start_time__" + def _make_collecting_func(): """Return a wrapper around async_request_spyre_chat that accumulates @@ -50,7 +55,10 @@ async def _wrapper( output = await async_request_spyre_chat(request_func_input, session, pbar) if output.success: if output.custom_metrics_dict: - _spyre_metrics_collected.append(output.custom_metrics_dict) + # Copy so the join key doesn't leak into the splat below. + _spyre_metrics_collected.append( + {**output.custom_metrics_dict, _SPYRE_START_TIME_KEY: output.start_time} + ) _request_outputs_collected.append( { "start_time": output.start_time, @@ -373,33 +381,57 @@ def _inject_spyre_metrics_into_result_file( logger.warning("Failed to read vllm result JSON %s: %s", file_path, exc) return - result["spyre_queued_time_s"] = [ - m["queued_time_s"] for m in metrics_list if "queued_time_s" in m - ] - result["spyre_num_chunked_prefills"] = [ - m["num_chunked_prefills"] for m in metrics_list if "num_chunked_prefills" in m - ] - result["spyre_chunk_prefill_latencies_s"] = [ - m.get("chunk_prefill_latencies_s", []) for m in metrics_list - ] - result["spyre_chunk_prefill_start_times_s"] = [ - m.get("chunk_prefill_start_times_s", []) for m in metrics_list - ] - result["spyre_total_prefill_chunks"] = sum(result["spyre_num_chunked_prefills"]) - result["spyre_decode_latencies_s"] = [m.get("decode_latencies_s", []) for m in metrics_list] - result["spyre_decode_start_times_s"] = [m.get("decode_start_times_s", []) for m in metrics_list] - result["spyre_tkvs"] = [m.get("tkvs", []) for m in metrics_list] - result["spyre_prefill_elapsed_s"] = [m.get("prefill_elapsed_s", 0.0) for m in metrics_list] - result["spyre_prefill_busy_s"] = [m.get("prefill_busy_s", 0.0) for m in metrics_list] - result["spyre_prefill_idle_s"] = [m.get("prefill_idle_s", 0.0) for m in metrics_list] - result["spyre_prefix_cache_hit_pct"] = [ - m.get("prefix_cache_hit_pct", 0.0) for m in metrics_list - ] - result["spyre_left_padding_blocks"] = [m.get("left_padding_blocks", []) for m in metrics_list] - result["spyre_pause_latencies_s"] = [m.get("pause_latencies_s", []) for m in metrics_list] - result["spyre_pause_start_times_s"] = [m.get("pause_start_times_s", []) for m in metrics_list] - result["spyre_was_missing_blocks"] = [m.get("was_missing_blocks", False) for m in metrics_list] - result["spyre_num_requests_missing_blocks"] = sum(result["spyre_was_missing_blocks"]) + # Reorder metrics to match vllm's submission-ordered rows by start_time, + # emitting one entry per row (missing rows get sentinels) so every spyre_* + # array stays aligned with ttfts/itls/start_times. + ordered: list[dict[str, Any] | None] + start_times = result.get("start_times") + if isinstance(start_times, list): + by_start = { + round(float(m[_SPYRE_START_TIME_KEY]), 9): m + for m in metrics_list + if _SPYRE_START_TIME_KEY in m + } + ordered = [by_start.get(round(float(st), 9)) for st in start_times] + matched = sum(1 for m in ordered if m is not None) + if matched != len(metrics_list): + logger.warning( + "Spyre metrics: matched %d/%d collected metric sets to result rows " + "by start_time; unmatched rows use sentinels.", + matched, + len(metrics_list), + ) + else: + logger.warning( + "Spyre metrics: result JSON has no 'start_times'; falling back to " + "collection order, which may not match vllm's per-request arrays." + ) + ordered = list(metrics_list) + + def _col(key: str, default: Any) -> list[Any]: + return [default if m is None else m.get(key, default) for m in ordered] + + result["spyre_queued_time_s"] = _col("queued_time_s", None) + result["spyre_num_chunked_prefills"] = _col("num_chunked_prefills", None) + result["spyre_chunk_prefill_latencies_s"] = _col("chunk_prefill_latencies_s", []) + result["spyre_chunk_prefill_start_times_s"] = _col("chunk_prefill_start_times_s", []) + result["spyre_total_prefill_chunks"] = sum( + n for n in result["spyre_num_chunked_prefills"] if n is not None + ) + result["spyre_decode_latencies_s"] = _col("decode_latencies_s", []) + result["spyre_decode_start_times_s"] = _col("decode_start_times_s", []) + result["spyre_tkvs"] = _col("tkvs", []) + result["spyre_prefill_elapsed_s"] = _col("prefill_elapsed_s", 0.0) + result["spyre_prefill_busy_s"] = _col("prefill_busy_s", 0.0) + result["spyre_prefill_idle_s"] = _col("prefill_idle_s", 0.0) + result["spyre_prefix_cache_hit_pct"] = _col("prefix_cache_hit_pct", 0.0) + result["spyre_left_padding_blocks"] = _col("left_padding_blocks", []) + result["spyre_pause_latencies_s"] = _col("pause_latencies_s", []) + result["spyre_pause_start_times_s"] = _col("pause_start_times_s", []) + result["spyre_was_missing_blocks"] = _col("was_missing_blocks", False) + result["spyre_num_requests_missing_blocks"] = sum( + 1 for v in result["spyre_was_missing_blocks"] if v + ) try: with open(file_path, "w", encoding="utf-8") as fh: From 2be93307e207ade6ebd964136bcab78194e2e7ef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 12 Aug 2026 15:45:47 +0000 Subject: [PATCH 096/106] test patch_serving compatibility MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- tests/utils/test_upstream_compatibility.py | 99 ++++++++++++++++++++++ 1 file changed, 99 insertions(+) diff --git a/tests/utils/test_upstream_compatibility.py b/tests/utils/test_upstream_compatibility.py index 5f8807aed..d34c298dc 100644 --- a/tests/utils/test_upstream_compatibility.py +++ b/tests/utils/test_upstream_compatibility.py @@ -5,6 +5,105 @@ compatibility code that can be cleaned up. """ +import inspect + import pytest pytestmark = pytest.mark.compat + + +# --------------------------------------------------------------------------- +# sendnn-bench serve custom metrics — sendnn_inference/v1/metrics/patch_serving.py +# +# patch_serving() wraps two upstream stream generators, passing the parameters +# below positionally. A rename/reorder/insert within that prefix, or a switch to +# keyword-only, silently misbinds result_generator and breaks every streaming +# request. Changes after the prefix are absorbed by *args/**kwargs. +# +# On failure: update the wrapper in patch_serving.py, then the prefix here. +# --------------------------------------------------------------------------- + +# Must match the parameters named in patch_serving._patch_chat._patched_generator +CHAT_STREAM_GENERATOR_PREFIX = ("self", "request", "result_generator", "request_id") + +# Must match the parameters named in patch_serving._patch_completions._patched_generator +COMPLETION_STREAM_GENERATOR_PREFIX = ( + "self", + "request", + "engine_inputs", + "result_generator", + "request_id", +) + + +def _assert_positional_prefix(func, expected_prefix: tuple[str, ...]) -> None: + params = list(inspect.signature(func).parameters.values()) + actual_prefix = tuple(p.name for p in params[: len(expected_prefix)]) + + assert actual_prefix == expected_prefix, ( + f"{func.__qualname__} leading parameters changed upstream: " + f"expected {expected_prefix}, got {actual_prefix}. " + f"patch_serving() binds these positionally and must be updated." + ) + + # Positional binding also requires that none of them became keyword-only. + for param in params[: len(expected_prefix)]: + assert param.kind is inspect.Parameter.POSITIONAL_OR_KEYWORD, ( + f"{func.__qualname__} parameter '{param.name}' is now {param.kind.name}; " + f"patch_serving() passes it positionally and must be updated." + ) + + +def test_sendnn_bench__serving_module_paths_unchanged(): + """patch_serving() imports these modules and skips patching (with only a + warning) if they move, so bench metrics would silently go missing.""" + import importlib + + for module_path, cls_name, method_name in [ + ( + "vllm.entrypoints.openai.chat_completion.serving", + "OpenAIServingChat", + "chat_completion_stream_generator", + ), + ( + "vllm.entrypoints.openai.completion.serving", + "OpenAIServingCompletion", + "completion_stream_generator", + ), + ]: + try: + module = importlib.import_module(module_path) + except ImportError as e: # pragma: no cover - only on upstream move + pytest.fail( + f"{module_path} is no longer importable ({e}); patch_serving() " + f"skips patching and bench metrics will be missing from SSE output." + ) + + cls = getattr(module, cls_name, None) + assert cls is not None, ( + f"{cls_name} no longer exists in {module_path}; patch_serving() must be updated." + ) + assert hasattr(cls, method_name), ( + f"{cls_name}.{method_name} no longer exists; patch_serving() patches this " + f"attribute and must be updated." + ) + + +def test_sendnn_bench__chat_stream_generator_signature_unchanged(): + """patch_serving._patch_chat wraps this method and binds its first + parameters positionally.""" + from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat + + _assert_positional_prefix( + OpenAIServingChat.chat_completion_stream_generator, CHAT_STREAM_GENERATOR_PREFIX + ) + + +def test_sendnn_bench__completion_stream_generator_signature_unchanged(): + """patch_serving._patch_completions wraps this method and binds its first + parameters positionally.""" + from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion + + _assert_positional_prefix( + OpenAIServingCompletion.completion_stream_generator, COMPLETION_STREAM_GENERATOR_PREFIX + ) From a0844f065722e716a531f984573a4c98dc545127 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Wed, 12 Aug 2026 15:56:14 +0000 Subject: [PATCH 097/106] test main_async marker compatibility MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_bench_serve.py | 14 +++- tests/utils/test_upstream_compatibility.py | 78 +++++++++++++++++++ 2 files changed, 91 insertions(+), 1 deletion(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 105e11541..ee9b64379 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -32,6 +32,10 @@ _BACKEND_NAME = "spyre-chat" +# Line upstream prints to close its metrics table (`print("=" * 50)` in +# vllm/benchmarks/serve.py). Pinned by tests/utils/test_upstream_compatibility.py. +_VLLM_METRICS_TABLE_END_MARKER = "=" * 50 + # Shared accumulators — populated by the wrapper below during the benchmark run. _spyre_metrics_collected: list[dict[str, Any]] = [] _request_outputs_collected: list[dict[str, Any]] = [] @@ -461,7 +465,7 @@ def _run_vllm_and_capture_trailing(args: Any) -> tuple[str, str]: class _StdoutSplitter: def write(self, s): if not done["v"]: - if s.strip() == "=" * 50: + if s.strip() == _VLLM_METRICS_TABLE_END_MARKER: done["v"] = True else: orig_stdout.write(s) @@ -490,6 +494,14 @@ def flush(self): sys.stdout = orig_stdout sys.stderr = orig_stderr + if not done["v"]: + logger.warning( + "Never saw upstream's end-of-metrics-table marker (%r), so no trailing " + "output was captured. vllm's `bench serve` output format has likely " + "changed; see _VLLM_METRICS_TABLE_END_MARKER.", + _VLLM_METRICS_TABLE_END_MARKER, + ) + return stdout_buf.getvalue(), stderr_buf.getvalue() diff --git a/tests/utils/test_upstream_compatibility.py b/tests/utils/test_upstream_compatibility.py index d34c298dc..e90b302e5 100644 --- a/tests/utils/test_upstream_compatibility.py +++ b/tests/utils/test_upstream_compatibility.py @@ -6,6 +6,7 @@ """ import inspect +import re import pytest @@ -107,3 +108,80 @@ def test_sendnn_bench__completion_stream_generator_signature_unchanged(): _assert_positional_prefix( OpenAIServingCompletion.completion_stream_generator, COMPLETION_STREAM_GENERATOR_PREFIX ) + + +# --------------------------------------------------------------------------- +# sendnn_inference/benchmarks/spyre_bench_serve.py +# +# _run_vllm_and_capture_trailing() calls main_async() and splits its stdout on a +# closing line of specific form, so it relies on that line's exact shape. +# +# Both break modes are silent: never matching captures nothing, matching too early +# swallows upstream's metrics table. +# --------------------------------------------------------------------------- + + +def test_sendnn_bench__main_async_unchanged(): + """spyre_bench_serve imports main_async and calls it as main_async(args).""" + import asyncio + + from vllm.benchmarks.serve import main_async + + assert asyncio.iscoroutinefunction(main_async), ( + "vllm.benchmarks.serve.main_async is no longer a coroutine function; " + "_run_vllm_and_capture_trailing calls it via asyncio.run()." + ) + + params = list(inspect.signature(main_async).parameters.values()) + positional = [ + p + for p in params + if p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) + ] + assert len(positional) == 1, ( + f"vllm.benchmarks.serve.main_async now takes {len(positional)} positional " + f"parameters ({[p.name for p in positional]}); _run_vllm_and_capture_trailing " + f"calls main_async(args) with exactly one." + ) + + +def test_sendnn_bench__metrics_table_end_marker_unchanged(): + """_run_vllm_and_capture_trailing splits stdout on the line closing upstream's + metrics table. It must still be printed, exactly once, as a standalone print().""" + import vllm.benchmarks.serve as upstream_serve + + from sendnn_inference.benchmarks.spyre_bench_serve import _VLLM_METRICS_TABLE_END_MARKER + + source = inspect.getsource(upstream_serve) + + # Matches `print("=" * 50)` / `print('=' * 50)` with flexible inner spacing. + width = len(_VLLM_METRICS_TABLE_END_MARKER) + char = _VLLM_METRICS_TABLE_END_MARKER[0] + pattern = re.compile(rf"""print\(\s*['"]{re.escape(char)}['"]\s*\*\s*{width}\s*\)""") + matches = pattern.findall(source) + + assert len(matches) == 1, ( + f"Expected exactly one `print({char!r} * {width})` in vllm.benchmarks.serve " + f"(found {len(matches)})." + ) + + +def test_sendnn_bench__table_headers_do_not_collide_with_end_marker(): + """Upstream's centered section headers must keep a non-empty title, else one + strips down to the end marker and the splitter swallows the metrics table.""" + import vllm.benchmarks.serve as upstream_serve + + from sendnn_inference.benchmarks.spyre_bench_serve import _VLLM_METRICS_TABLE_END_MARKER + + source = inspect.getsource(upstream_serve) + marker_char = _VLLM_METRICS_TABLE_END_MARKER[0] + + # Upstream centers section titles in a run of '=' or '-', e.g. + # print("{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="=")) + for match in re.finditer(r"""\.format\(\s*s=\s*(['"])(.*?)\1""", source): + title = match.group(2) + assert title.strip(marker_char).strip(), ( + f"Upstream renders a centered header with title {title!r}, which strips to " + f"the end-of-table marker and would flip _run_vllm_and_capture_trailing's " + f"splitter early, swallowing upstream's metrics table." + ) From a178f5cd59b55c3d975f0076c490a913ce98fab1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Thu, 13 Aug 2026 08:32:56 +0000 Subject: [PATCH 098/106] force set explicit name when --save-result or --plot-timeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- .../benchmarks/spyre_bench_serve.py | 192 ++++++++++-------- tests/benchmarks/test_bench_metrics.py | 110 ++++++++-- 2 files changed, 205 insertions(+), 97 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index ee9b64379..5f4f04507 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -5,8 +5,15 @@ sendnn-bench serve --host localhost --port 8000 --model \\ --dataset-name random --num-prompts 20 --request-rate 2 +Result files: + Any flag that makes vllm write a result JSON (--save-result, --plot-timeline, + --detailed-timeline) requires both --result-dir and --result-filename, so the + path to inject Spyre metrics into is unambiguous. --append-result is rejected: + it produces JSONL, compatibility with injected metrics is not yet supported. + Env var: - SENDNN_INFERENCE_BENCH_METRICS_ENABLED=1 (must also be set on the server) + SENDNN_INFERENCE_BENCH_METRICS_ENABLED=1 (must be set on the server; the + client does not read it) """ import argparse @@ -14,7 +21,6 @@ import json import logging import os -import time from typing import Any import numpy as np @@ -117,8 +123,9 @@ def _build_parser() -> FlexibleArgumentParser: default=False, help=( "Write a detailed per-request Gantt-chart timeline HTML alongside the " - "JSON result file (same name with a _detailed.html suffix). " - "Requires --save-result and SENDNN_INFERENCE_BENCH_METRICS_ENABLED on the server." + "JSON result file (same name with a _detailed_timeline.html suffix). " + "Requires --result-dir and --result-filename, and " + "SENDNN_INFERENCE_BENCH_METRICS_ENABLED on the server." ), ) parser.add_argument( @@ -147,6 +154,61 @@ def _build_parser() -> FlexibleArgumentParser: return parser +# Flags that make upstream `vllm bench serve` write a result JSON, plus our own +# flags that need to locate that JSON afterwards. Each entry is (attr, cli_flag). +_RESULT_FILE_FLAGS: tuple[tuple[str, str], ...] = ( + ("save_result", "--save-result"), + ("plot_timeline", "--plot-timeline"), + ("detailed_timeline", "--detailed-timeline"), +) + + +def _validate_result_file_args(args: Any) -> None: + """Require an unambiguous result-file path whenever a JSON will be written. + + Upstream picks the result filename itself (timestamped) unless + ``--result-filename`` is given, which would force us to guess which file it + wrote by scanning for the newest ``.json`` in the directory. Requiring both + ``--result-dir`` and ``--result-filename`` makes the path exact. + + ``--append-result`` is rejected outright: it makes upstream emit one JSON + object per line (JSONL), which cannot be read back and rewritten as a single + document without corrupting the file. + """ + if getattr(args, "append_result", False): + raise ValueError( + "--append-result is not supported by sendnn-bench: it is incompatible " + "with Spyre metrics injection." + ) + + triggered = [flag for attr, flag in _RESULT_FILE_FLAGS if getattr(args, attr, False)] + if not triggered: + return + + missing = [] + if not getattr(args, "result_dir", None): + missing.append("--result-dir") + if not getattr(args, "result_filename", None): + missing.append("--result-filename") + if missing: + raise ValueError( + f"{', '.join(triggered)} requires both --result-dir and --result-filename " + f"in sendnn-bench for Spyre metrics injection (missing: {', '.join(missing)})." + ) + + +def _result_file_path(args: Any) -> str | None: + """Path of the result JSON, or None if no JSON will be written.""" + if not any(getattr(args, attr, False) for attr, _ in _RESULT_FILE_FLAGS): + return None + name = getattr(args, "result_filename", None) + if not name: + return None + if os.path.isabs(name): + return name + return os.path.join(getattr(args, "result_dir", None) or ".", name) + + # Short explanation of every metric printed by _print_spyre_section, written to # _METRIC_DESCRIPTION_FILENAME when --describe-metrics is passed. Ordered to match # the printed output. @@ -335,49 +397,26 @@ def _section(header: str, values: list[float], label: str) -> None: def _inject_spyre_metrics_into_result_file( args: Any, metrics_list: list[dict[str, Any]], - run_started_at: float, ) -> None: - """If vllm wrote a result JSON (--save-result / --append-result / --result-filename), - find it and inject per-request Spyre metric lists alongside vllm's own per-request + """If vllm wrote a result JSON (--save-result / --plot-timeline / --detailed-timeline), + inject per-request Spyre metric lists alongside vllm's own per-request fields (ttfts, itls, …).""" if not metrics_list: return - if not ( - getattr(args, "save_result", False) - or getattr(args, "append_result", False) - or getattr(args, "result_filename", None) - ): - return - # Locate the file vllm just wrote by finding the newest .json in the result dir - # that was modified after we started the run. - result_dir = getattr(args, "result_dir", None) or "." - explicit_name = getattr(args, "result_filename", None) + # None means no result-file flag was passed, i.e. vllm wrote no JSON — nothing to do. + file_path = _result_file_path(args) + if file_path is None: + return - if explicit_name: - candidate = ( - explicit_name - if os.path.isabs(explicit_name) - else os.path.join(result_dir, explicit_name) + if not os.path.isfile(file_path): + logger.warning( + "Expected vllm result JSON at %s but it does not exist; " + "skipping Spyre metric injection.", + file_path, ) - candidates = [candidate] if os.path.isfile(candidate) else [] - else: - try: - candidates = [ - os.path.join(result_dir, f) - for f in os.listdir(result_dir) - if f.endswith(".json") - and os.path.getmtime(os.path.join(result_dir, f)) >= run_started_at - ] - except OSError: - candidates = [] - - if not candidates: - logger.warning("Could not locate vllm result JSON to inject Spyre metrics into.") return - file_path = max(candidates, key=os.path.getmtime) - try: with open(file_path, encoding="utf-8") as fh: result = json.load(fh) @@ -519,12 +558,17 @@ def main() -> None: # Force our custom backend so Spyre metrics are always collected. args.backend = _BACKEND_NAME + # Validate the results JSON path arguments + try: + _validate_result_file_args(args) + except ValueError as exc: + parser.error(str(exc)) + selected_percentiles = [float(p) for p in args.metric_percentiles.split(",")] _spyre_metrics_collected.clear() _request_outputs_collected.clear() - run_started_at = time.time() stdout_trailing, stderr_trailing = _run_vllm_and_capture_trailing(args) print("{s:{c}^{n}}".format(s=" SenDNN Metrics ", n=50, c="=")) @@ -533,60 +577,38 @@ def main() -> None: if getattr(args, "describe_metrics", False): _write_metric_descriptions(getattr(args, "result_dir", None)) - _inject_spyre_metrics_into_result_file(args, _spyre_metrics_collected, run_started_at) + _inject_spyre_metrics_into_result_file(args, _spyre_metrics_collected) if getattr(args, "detailed_timeline", False): from pathlib import Path from sendnn_inference.benchmarks.spyre_plot import generate_detailed_timeline_plot + # the path is always resolvable here (see _validate_result_file_args) + result_json = _result_file_path(args) + assert result_json is not None, "--detailed-timeline implies a known result JSON path" + + json_path = Path(result_json) # Derive the HTML path from the JSON result file: same name, _detailed.html suffix. - result_dir = getattr(args, "result_dir", None) or "." - explicit_name = getattr(args, "result_filename", None) - if explicit_name: - json_candidate = ( - explicit_name - if os.path.isabs(explicit_name) - else os.path.join(result_dir, explicit_name) - ) - candidates = [json_candidate] if os.path.isfile(json_candidate) else [] - else: + html_path = json_path.with_name(json_path.stem + "_detailed_timeline.html") + decode_thresholds_str = getattr(args, "decode_thresholds", None) + # Parse comma-separated milliseconds and convert to seconds + decode_thresholds = None + if decode_thresholds_str: try: - candidates = [ - os.path.join(result_dir, f) - for f in os.listdir(result_dir) - if f.endswith(".json") - and os.path.getmtime(os.path.join(result_dir, f)) >= run_started_at - ] - except OSError: - candidates = [] - - if candidates: - json_path = Path(str(max(candidates, key=os.path.getmtime))) - html_path = json_path.with_name(json_path.stem + "_detailed_timeline.html") - decode_thresholds_str = getattr(args, "decode_thresholds", None) - # Parse comma-separated milliseconds and convert to seconds - decode_thresholds = None - if decode_thresholds_str: - try: - thresholds_ms = [float(x.strip()) for x in decode_thresholds_str.split(",")] - if len(thresholds_ms) != 2: - raise ValueError("Expected exactly 2 comma-separated values") - decode_thresholds = [ms / 1000.0 for ms in thresholds_ms] - except (ValueError, AttributeError) as e: - logger.warning( - "Invalid --decode-thresholds format: %s (expected LOW,HIGH in ms)", - e, - ) - decode_thresholds = None - generate_detailed_timeline_plot( - _request_outputs_collected, html_path, decode_thresholds=decode_thresholds - ) - else: - logger.warning( - "--detailed-timeline requires --save-result so the JSON path is known; " - "no result file found, skipping timeline." - ) + thresholds_ms = [float(x.strip()) for x in decode_thresholds_str.split(",")] + if len(thresholds_ms) != 2: + raise ValueError("Expected exactly 2 comma-separated values") + decode_thresholds = [ms / 1000.0 for ms in thresholds_ms] + except (ValueError, AttributeError) as e: + logger.warning( + "Invalid --decode-thresholds format: %s (expected LOW,HIGH in ms)", + e, + ) + decode_thresholds = None + generate_detailed_timeline_plot( + _request_outputs_collected, html_path, decode_thresholds=decode_thresholds + ) trailing = stdout_trailing + stderr_trailing if trailing.strip(): diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index a31e07be4..a8fa62b57 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -13,7 +13,6 @@ import asyncio import json import pathlib -import time from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -22,6 +21,8 @@ from sendnn_inference.benchmarks.spyre_bench_serve import ( _inject_spyre_metrics_into_result_file, _print_spyre_section, + _result_file_path, + _validate_result_file_args, ) # --------------------------------------------------------------------------- @@ -69,18 +70,34 @@ SELECTED_PERCENTILES = [90.0, 99.0, 100.0] +_RESULT_FILENAME = "result.json" + + def _write_fake_result(tmp_path) -> pathlib.Path: - p = tmp_path / "result.json" + p = tmp_path / _RESULT_FILENAME p.write_text(json.dumps({"backend": "spyre-chat", "num_prompts": 2})) return p -def _make_args(tmp_path, *, save_result: bool = True, result_filename=None): +def _make_args( + tmp_path, + *, + save_result: bool = True, + append_result: bool = False, + plot_timeline: bool = False, + detailed_timeline: bool = False, + result_filename: Any = _RESULT_FILENAME, + result_dir: Any = None, +): + """Build an args namespace. Defaults mirror the validated happy path: + --save-result with both --result-dir and --result-filename set.""" return argparse.Namespace( save_result=save_result, - append_result=False, + append_result=append_result, + plot_timeline=plot_timeline, + detailed_timeline=detailed_timeline, result_filename=str(result_filename) if result_filename else None, - result_dir=str(tmp_path), + result_dir=str(result_dir) if result_dir else str(tmp_path), ) @@ -92,7 +109,7 @@ def _make_args(tmp_path, *, save_result: bool = True, result_filename=None): @pytest.mark.cpu def test_inject_adds_spyre_keys(tmp_path): result_file = _write_fake_result(tmp_path) - _inject_spyre_metrics_into_result_file(_make_args(tmp_path), FAKE_METRICS, time.time() - 1) + _inject_spyre_metrics_into_result_file(_make_args(tmp_path), FAKE_METRICS) data = json.loads(result_file.read_text()) expected_keys = {"spyre_" + k for k in FAKE_METRICS[0]} | { "spyre_total_prefill_chunks", @@ -105,7 +122,7 @@ def test_inject_adds_spyre_keys(tmp_path): @pytest.mark.cpu def test_inject_values_correct(tmp_path): result_file = _write_fake_result(tmp_path) - _inject_spyre_metrics_into_result_file(_make_args(tmp_path), FAKE_METRICS, time.time() - 1) + _inject_spyre_metrics_into_result_file(_make_args(tmp_path), FAKE_METRICS) data = json.loads(result_file.read_text()) # Per-request lists map directly: "spyre_" + key → [m[key] for m in FAKE_METRICS] @@ -129,9 +146,7 @@ def test_inject_values_correct(tmp_path): def test_inject_noop_when_save_result_false(tmp_path): result_file = _write_fake_result(tmp_path) original = result_file.read_text() - _inject_spyre_metrics_into_result_file( - _make_args(tmp_path, save_result=False), FAKE_METRICS, time.time() - 1 - ) + _inject_spyre_metrics_into_result_file(_make_args(tmp_path, save_result=False), FAKE_METRICS) assert result_file.read_text() == original @@ -139,7 +154,7 @@ def test_inject_noop_when_save_result_false(tmp_path): def test_inject_noop_when_metrics_empty(tmp_path): result_file = _write_fake_result(tmp_path) original = result_file.read_text() - _inject_spyre_metrics_into_result_file(_make_args(tmp_path), [], time.time() - 1) + _inject_spyre_metrics_into_result_file(_make_args(tmp_path), []) assert result_file.read_text() == original @@ -147,11 +162,82 @@ def test_inject_noop_when_metrics_empty(tmp_path): def test_inject_explicit_result_filename(tmp_path): result_file = _write_fake_result(tmp_path) args = _make_args(tmp_path, result_filename=str(result_file)) - _inject_spyre_metrics_into_result_file(args, FAKE_METRICS, time.time() - 1) + _inject_spyre_metrics_into_result_file(args, FAKE_METRICS) data = json.loads(result_file.read_text()) assert "spyre_queued_time_s" in data +# --------------------------------------------------------------------------- +# Test 1B — _validate_result_file_args / _result_file_path +# --------------------------------------------------------------------------- + + +@pytest.mark.cpu +def test_validate_rejects_append_result(tmp_path): + """--append-result makes vllm write JSONL, which cannot carry injected metrics.""" + args = _make_args(tmp_path, append_result=True) + with pytest.raises(ValueError, match="--append-result is not supported"): + _validate_result_file_args(args) + + +@pytest.mark.cpu +@pytest.mark.parametrize("flag", ["save_result", "plot_timeline", "detailed_timeline"]) +@pytest.mark.parametrize( + ("result_dir", "result_filename", "expected_missing"), + [ + (None, None, ["--result-dir", "--result-filename"]), + ("set", None, ["--result-filename"]), + (None, "r.json", ["--result-dir"]), + ], +) +def test_validate_requires_both_dir_and_filename( + tmp_path, flag, result_dir, result_filename, expected_missing +): + """Every flag that causes a result JSON to be written requires an explicit path.""" + # Built directly rather than via _make_args, which defaults result_dir to tmp_path + # and so could not express the "--result-dir absent" case. + args = argparse.Namespace( + save_result=False, + append_result=False, + plot_timeline=False, + detailed_timeline=False, + result_dir=str(tmp_path) if result_dir else None, + result_filename=result_filename, + ) + setattr(args, flag, True) + with pytest.raises(ValueError) as exc: + _validate_result_file_args(args) + for missing in expected_missing: + assert missing in str(exc.value) + + +@pytest.mark.cpu +def test_validate_accepts_full_path(tmp_path): + _validate_result_file_args(_make_args(tmp_path)) # must not raise + + +@pytest.mark.cpu +def test_validate_noop_when_no_result_file_requested(tmp_path): + """Without any result-file flag, the path args are irrelevant.""" + args = _make_args(tmp_path, save_result=False, result_filename=None, result_dir="") + _validate_result_file_args(args) # must not raise + + +@pytest.mark.cpu +def test_result_file_path_none_when_nothing_written(tmp_path): + args = _make_args(tmp_path, save_result=False) + assert _result_file_path(args) is None + + +@pytest.mark.cpu +def test_inject_warns_when_file_absent(tmp_path, caplog): + """Path is known but vllm wrote nothing there — warn, don't crash.""" + args = _make_args(tmp_path, result_filename="never_written.json") + with caplog.at_level("WARNING"): + _inject_spyre_metrics_into_result_file(args, FAKE_METRICS) + assert "does not exist" in caplog.text + + # --------------------------------------------------------------------------- # Test 2 — _print_spyre_section # --------------------------------------------------------------------------- From 3e3fd99d6527a7f7958e4c2c6cb6efd43701c360 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Thu, 13 Aug 2026 09:08:04 +0000 Subject: [PATCH 099/106] bugfix: number of left-padding blocks can momentally go negative MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/v1/core/scheduler.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index e5aa217cb..0d172d828 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -372,8 +372,11 @@ def _bench_update_from_output( req = req_by_id.get(req_id) if req is not None: req_num_blocks = math.ceil(req.num_computed_tokens / self.block_size) + # Clamp at 0: num_computed_tokens is read after the step, while + # tkv is the runner's value for that step, so a request that just + # crossed a block boundary can momentarily exceed max_num_blocks. self._bench.left_padding_blocks.setdefault(req_id, []).append( - max_num_blocks - req_num_blocks + max(0, max_num_blocks - req_num_blocks) ) self._bench.prefill_step_start = None self._bench.decode_step_start = None From 68e5288183c0bed21e08858759e9c0cbf925b7dc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Thu, 13 Aug 2026 11:14:42 +0000 Subject: [PATCH 100/106] handle detailed_timeline_plot in warning message MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- sendnn_inference/benchmarks/spyre_bench_serve.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 5f4f04507..b48797522 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -606,9 +606,13 @@ def main() -> None: e, ) decode_thresholds = None - generate_detailed_timeline_plot( - _request_outputs_collected, html_path, decode_thresholds=decode_thresholds - ) + # TODO need to fix this timeline bug on long runs + try: + generate_detailed_timeline_plot( + _request_outputs_collected, html_path, decode_thresholds=decode_thresholds + ) + except Exception: + logger.warning("Failed to generate detailed timeline plot", exc_info=True) trailing = stdout_trailing + stderr_trailing if trailing.strip(): From 8f579b490d298f991559007f4b71998bf3596280 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Thu, 13 Aug 2026 11:50:04 +0000 Subject: [PATCH 101/106] bugfix failing test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- tests/benchmarks/test_bench_metrics.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index a8fa62b57..a8e9a3a19 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -230,12 +230,12 @@ def test_result_file_path_none_when_nothing_written(tmp_path): @pytest.mark.cpu -def test_inject_warns_when_file_absent(tmp_path, caplog): +def test_inject_warns_when_file_absent(tmp_path, caplog_sendnn_inference): """Path is known but vllm wrote nothing there — warn, don't crash.""" args = _make_args(tmp_path, result_filename="never_written.json") - with caplog.at_level("WARNING"): + with caplog_sendnn_inference.at_level("WARNING"): _inject_spyre_metrics_into_result_file(args, FAKE_METRICS) - assert "does not exist" in caplog.text + assert "does not exist" in caplog_sendnn_inference.text # --------------------------------------------------------------------------- From e828990c390d76ffbb3d5cd89b00b6765735f2dc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sophie=20du=20Cou=C3=A9dic?= Date: Thu, 13 Aug 2026 16:16:52 +0000 Subject: [PATCH 102/106] add sendnn-bench docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sophie du Couédic --- docs/.nav.yml | 2 + .../detailed_performance_measurement.md | 106 ++++++++++++++++++ docs/user_guide/performance.md | 4 + 3 files changed, 112 insertions(+) create mode 100644 docs/user_guide/detailed_performance_measurement.md diff --git a/docs/.nav.yml b/docs/.nav.yml index 5b4936ca1..3393a9064 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -17,6 +17,7 @@ nav: - Supported Features: user_guide/supported_features.md - Supported Models: user_guide/supported_models.md - Performance Tuning: user_guide/performance.md + - Detailed Performance Measurement: user_guide/detailed_performance_measurement.md - Developer Guide: - Contributing: contributing/README.md - Maintaining: contributing/maintaining.md @@ -38,6 +39,7 @@ nav: - Supported Features: user_guide/supported_features.md - Supported Models: user_guide/supported_models.md - Performance Tuning: user_guide/performance.md + - Detailed Performance Measurement: user_guide/detailed_performance_measurement.md - Developer Guide: - Contributing: contributing/README.md - Maintaining: contributing/maintaining.md diff --git a/docs/user_guide/detailed_performance_measurement.md b/docs/user_guide/detailed_performance_measurement.md new file mode 100644 index 000000000..4ad8bb431 --- /dev/null +++ b/docs/user_guide/detailed_performance_measurement.md @@ -0,0 +1,106 @@ +# Detailed Performance Measurement + +`sendnn-bench` is a drop-in replacement for [`vllm bench serve`](https://docs.vllm.ai/en/stable/benchmarking/cli/#online-benchmark) that collects **additional Spyre-specific per-request metrics** on top of the usual TTFT/TPOT/ITL/E2EL values described in [Benchmarking and Performance](./performance.md). It reuses the upstream implementation — every `vllm bench serve` flag keeps working — and augments it with scheduler-level information such as queue wait time, per-chunk prefill latencies, per-step decode latencies, prefix cache hit rate, left padding and request pausing. + +## Usage + +1. Start the server with Spyre metrics collection enabled: + +```bash +SENDNN_INFERENCE_BENCH_METRICS_ENABLED=1 vllm serve \ + --model {model} \ + --max-model-len {max-model-len} \ + --max-num-seqs {max-num-seqs} +``` + +!!! warning + + `SENDNN_INFERENCE_BENCH_METRICS_ENABLED=1` must be set **on the server**. The client does not read it. Without it the server returns no Spyre metrics and the client logs a warning. + +1. Run the benchmark client with `sendnn-bench serve` instead of `vllm bench serve`: + +```bash +sendnn-bench serve \ + --model {model} \ + --endpoint /v1/completions \ + --dataset-name {custom/sharegpt/random...} \ + --dataset-path {path to dataset} \ + --num-prompts {num-prompts} \ + --max-concurrency {num-concurrent-users} \ + --save-result \ + --result-dir {path/to/results} \ + --result-filename result.json +``` + +!!! note + + `sendnn-bench` uses its own `spyre-chat` backend by default, which parses the extra metrics out of the streamed response. Do not override it with `--backend`. + + Any flag that makes a result JSON be written (`--save-result`, `--plot-timeline`, `--detailed-timeline`) requires **both** `--result-dir` and `--result-filename`, so that the file to inject the Spyre metrics into is unambiguous. `--append-result` is not supported. + +The Spyre metrics are printed in a `SenDNN Metrics` section appended to the regular benchmark result table: + +```text +============ Serving Benchmark Result ============ +Successful requests: XX +... +----------------End-to-end Latency---------------- +Mean E2EL (ms): XX +... +================= SenDNN Metrics ================= +Total prefill chunks processed: XX +Requests blocked by missing KV blocks: XX +---------------- Queue Wait Time ----------------- +Mean Queue Wait Time (ms): XX +Median Queue Wait Time (ms): XX +P99 Queue Wait Time (ms): XX +P100 Queue Wait Time (ms): XX +------------- Chunked Prefill Count -------------- +... +------------ Chunked Prefill Latency ------------- +... +--------------- Prefill Phase Time --------------- +... +------------- Time Spent Prefilling -------------- +... +------------ Prefill Phase Idle Time ------------- +... +-------------- Decode Step Latency --------------- +... +---------------- Prefix Cache Hit ---------------- +... +-------------- Left Padding Blocks --------------- +... +----------------- Pause Latency ------------------ +... +---------------- Number of Pauses ---------------- +... +--------------- Total Time Paused ---------------- +... +================================================== +``` + +## Additional Flags + +Beyond the upstream `vllm bench serve` flags, `sendnn-bench` adds: + +- **`--describe-metrics`**: writes a `sendnn_bench_metrics_description.txt` file into `--result-dir` (or the current directory), documenting what every printed metric measures and its sample granularity (one sample per request, per prefill chunk, per decode step, …). Recommended whenever you share results with someone else. + +- **`--detailed-timeline`**: writes a `{result-filename}_detailed_timeline.html` Gantt chart next to the result JSON, viewable in any modern web browser. Unlike the upstream `--plot-timeline` — which shows a single TTFT bar per request — it breaks each request down into its queue wait, its individual chunked prefill steps, and its decode steps, which makes it easy to see where the time actually went. + +- **`--decode-thresholds LOW,HIGH`**: two decode latency thresholds in milliseconds used to color the decode steps of the detailed timeline (green below `LOW`, orange between, red above `HIGH`). Only meaningful together with `--detailed-timeline`. + +Combined with `--save-result`, per-request Spyre values are also injected into the result JSON as `spyre_*` arrays (`spyre_queued_time_s`, `spyre_chunk_prefill_latencies_s`, `spyre_decode_latencies_s`, `spyre_prefix_cache_hit_pct`, …), aligned with vLLM's own per-request arrays (`ttfts`, `itls`, `start_times`). + +!!! info + + Plot generation requires the plotting libraries: `uv pip install vllm[bench]` + +## Adding a New Metric + +The set of collected metrics is meant to grow. The [`add-bench-metric`](https://github.com/torch-spyre/sendnn-inference/blob/main/.claude/skills/add-bench-metric/SKILL.md) Claude Code skill walks through every layer that a new per-request metric has to touch — scheduler-side timing, transport to the client, aggregation, printing, result JSON injection and tests. Give it a precise description of the metric and where its value should be computed: + +```text +/add-bench-metric add a metric for the prefix cache hit percent, based on the number of +chunks saved from a cache hit over the expected number of prefill chunks, for each request. +``` diff --git a/docs/user_guide/performance.md b/docs/user_guide/performance.md index 72580e663..0b900cb25 100644 --- a/docs/user_guide/performance.md +++ b/docs/user_guide/performance.md @@ -68,6 +68,10 @@ The following additional flags can help with insights and result interpretation: - `--save-detailed`: saves individual recorded data per request (useful for debugging) - `--result-dir {path/to/results}`: target path for output results +!!! tip + + For Spyre-specific per-request metrics (queue wait time, per-chunk prefill latencies, per-step decode latencies, prefix cache hit rate, left padding, pausing), use `sendnn-bench serve` instead of `vllm bench serve`. See [Detailed Performance Measurement](./detailed_performance_measurement.md). + ### `--custom-output-len -1` When running benchmarks, all requests typically use the same `max-tokens` value (the maximum number of output tokens for a request). This value can be set using [`--output-len`](https://docs.vllm.ai/en/stable/cli/bench/serve/#-output-len). For the `custom` dataset (`--dataset-name custom`), if the dataset contains per-request output token counts as shown in the [Custom dataset documentation](https://docs.vllm.ai/en/stable/api/vllm/benchmarks/datasets/#vllm.benchmarks.datasets.CustomDataset), you can load the per-request `max-tokens` using `--custom-output-len -1`. Paired with `--ignore-eos` (which tells the model to ignore the EOS token and always generate exactly `max-tokens` tokens), this makes benchmarks more stable and reproducible, since the number of output tokens is fixed across runs. Without this, output length varies across runs — even at temperature 0.0, unless using [batch invariance](https://docs.vllm.ai/en/latest/features/batch_invariance/#batch-invariance) — making results more variable and difficult to interpret. From f9e2767de2913849f2c6aabbee423613a0bcb0c5 Mon Sep 17 00:00:00 2001 From: Yannick Schnider Date: Fri, 14 Aug 2026 14:54:55 +0200 Subject: [PATCH 103/106] add test to pin _free_request return contract and kv_transfer_params seams in upstream compat suite Signed-off-by: Yannick Schnider --- tests/utils/test_upstream_compatibility.py | 77 ++++++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/tests/utils/test_upstream_compatibility.py b/tests/utils/test_upstream_compatibility.py index e90b302e5..281ba9d2c 100644 --- a/tests/utils/test_upstream_compatibility.py +++ b/tests/utils/test_upstream_compatibility.py @@ -185,3 +185,80 @@ def test_sendnn_bench__table_headers_do_not_collide_with_end_marker(): f"the end-of-table marker and would flip _run_vllm_and_capture_trailing's " f"splitter early, swallowing upstream's metrics table." ) + + +# --------------------------------------------------------------------------- +# sendnn-bench serve custom metrics — cross-process metric hop +# +# The scheduler-side override packs its per-request metrics into the dict that +# Scheduler._free_request returns, and relies on that dict riding +# kv_transfer_params all the way to the API server: +# +# SpyreScheduler._free_request (mutates super()'s returned kv_xfer_params) +# -> EngineCoreOutput.kv_transfer_params (ZMQ hop to API server process) +# -> RequestOutput.kv_transfer_params (read in patch_serving, gated on +# RequestOutput.finished) +# +# If upstream stops returning the dict from _free_request, or drops +# kv_transfer_params from either struct, or renames RequestOutput.finished, the +# metrics silently vanish with a green suite. These tests pin each seam. +# +# On failure: the metric plumbing in scheduler.py / patch_serving.py must be +# reworked around the new upstream shape. +# --------------------------------------------------------------------------- + + +def test_sendnn_bench__free_request_returns_kv_xfer_dict(): + """SpyreScheduler._free_request calls super()._free_request(request, + delay_free_blocks), mutates the returned dict with a '__spyre__' key, and + returns it. Upstream must keep the (request, delay_free_blocks) signature and + a dict|None return.""" + from vllm.v1.core.sched.scheduler import Scheduler + + sig = inspect.signature(Scheduler._free_request) + params = list(sig.parameters.values()) + names = [p.name for p in params] + + assert names[:3] == ["self", "request", "delay_free_blocks"], ( + f"Scheduler._free_request parameters changed upstream: expected " + f"(self, request, delay_free_blocks, ...), got {names}. " + f"SpyreScheduler._free_request delegates positionally and must be updated." + ) + + delay_free_blocks = sig.parameters["delay_free_blocks"] + assert delay_free_blocks.kind is inspect.Parameter.POSITIONAL_OR_KEYWORD, ( + f"Scheduler._free_request 'delay_free_blocks' is now {delay_free_blocks.kind.name}; " + f"SpyreScheduler._free_request passes it positionally and must be updated." + ) + + # The override packs metrics into this return value and hands it back; a None + # or non-dict-bearing return would drop the metrics on the floor. + ret = sig.return_annotation + assert ret is not inspect.Signature.empty and "dict" in str(ret), ( + f"Scheduler._free_request return annotation is {ret!r}, no longer a dict|None; " + f"SpyreScheduler._free_request stashes metrics under a '__spyre__' key in this " + f"dict and must be updated." + ) + + +def test_sendnn_bench__kv_transfer_params_survives_to_request_output(): + """The '__spyre__' metrics ride kv_transfer_params from EngineCoreOutput to + RequestOutput, and patch_serving reads them gated on RequestOutput.finished.""" + from vllm.outputs import RequestOutput + from vllm.v1.engine import EngineCoreOutput + + assert "kv_transfer_params" in EngineCoreOutput.__struct_fields__, ( + "EngineCoreOutput no longer carries kv_transfer_params; the scheduler-side " + "'__spyre__' metrics cannot cross to the API server process. patch_serving " + "and scheduler._free_request must be updated." + ) + + ro_params = inspect.signature(RequestOutput.__init__).parameters + assert "kv_transfer_params" in ro_params, ( + "RequestOutput.__init__ no longer accepts kv_transfer_params; patch_serving " + "reads res.kv_transfer_params.get('__spyre__') and must be updated." + ) + assert "finished" in ro_params, ( + "RequestOutput.__init__ no longer accepts 'finished'; patch_serving gates the " + "'__spyre__' read on res.finished and must be updated." + ) From 2603eb781964378a4007f36dc5aa26dc188926d7 Mon Sep 17 00:00:00 2001 From: Yannick Schnider Date: Fri, 14 Aug 2026 15:56:28 +0200 Subject: [PATCH 104/106] fix(bench): skip per-request metric injection when result JSON lacks start_times, and cover the reordering + sentinel-pad paths Signed-off-by: Yannick Schnider --- .../benchmarks/spyre_bench_serve.py | 10 ++- tests/benchmarks/test_bench_metrics.py | 89 +++++++++++++++++-- 2 files changed, 91 insertions(+), 8 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index b48797522..9195f795d 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -445,11 +445,15 @@ def _inject_spyre_metrics_into_result_file( len(metrics_list), ) else: + # No start_times means no key to align against vllm's submission-ordered + # arrays, so any failed/reordered request would misalign spyre_*[i] from + # ttfts[i]. Refuse to write the columns rather than emit corrupt data. logger.warning( - "Spyre metrics: result JSON has no 'start_times'; falling back to " - "collection order, which may not match vllm's per-request arrays." + "Spyre metrics: result JSON has no 'start_times' to align against; " + "skipping per-request Spyre metric injection to avoid emitting arrays " + "misaligned with vllm's ttfts/itls/start_times." ) - ordered = list(metrics_list) + return def _col(key: str, default: Any) -> list[Any]: return [default if m is None else m.get(key, default) for m in ordered] diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index a8e9a3a19..c93cbf738 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -19,6 +19,7 @@ import pytest from sendnn_inference.benchmarks.spyre_bench_serve import ( + _SPYRE_START_TIME_KEY, _inject_spyre_metrics_into_result_file, _print_spyre_section, _result_file_path, @@ -29,7 +30,14 @@ # Shared test data # --------------------------------------------------------------------------- -# Synthetic values +# Per-request start_time stamps carried on each collected metric dict. vllm's +# result arrays (ttfts/itls/start_times) are submission-ordered, so injection +# realigns metrics against them by this key. The values are distinct so the +# reorder is observable. +FAKE_START_TIMES = [111.111, 222.222] + +# Synthetic values. FAKE_METRICS[i] carries FAKE_START_TIMES[i] under +# _SPYRE_START_TIME_KEY so the injection reordering path can match it to a row. FAKE_METRICS: list[dict[str, Any]] = [ { "queued_time_s": 42.0, @@ -47,6 +55,7 @@ "pause_latencies_s": [0.5, 1.2], "pause_start_times_s": [0.0, 1.2], "was_missing_blocks": True, + _SPYRE_START_TIME_KEY: FAKE_START_TIMES[0], }, { "queued_time_s": 0.00001, @@ -64,18 +73,29 @@ "pause_latencies_s": [0.3], "pause_start_times_s": [0.5], "was_missing_blocks": False, + _SPYRE_START_TIME_KEY: FAKE_START_TIMES[1], }, ] +# Metric keys that map 1:1 to a spyre_ result column. The start-time stamp +# is an alignment key only — it is never written back out. +FAKE_METRIC_KEYS = [k for k in FAKE_METRICS[0] if k != _SPYRE_START_TIME_KEY] + SELECTED_PERCENTILES = [90.0, 99.0, 100.0] _RESULT_FILENAME = "result.json" -def _write_fake_result(tmp_path) -> pathlib.Path: +def _write_fake_result(tmp_path, *, start_times: Any = FAKE_START_TIMES) -> pathlib.Path: + """Write a stand-in vllm result JSON. By default it carries submission-ordered + start_times aligned with FAKE_METRICS so injection takes the reordering path; + pass start_times=None to simulate a result file lacking that key.""" p = tmp_path / _RESULT_FILENAME - p.write_text(json.dumps({"backend": "spyre-chat", "num_prompts": 2})) + payload: dict[str, Any] = {"backend": "spyre-chat", "num_prompts": 2} + if start_times is not None: + payload["start_times"] = start_times + p.write_text(json.dumps(payload)) return p @@ -111,12 +131,14 @@ def test_inject_adds_spyre_keys(tmp_path): result_file = _write_fake_result(tmp_path) _inject_spyre_metrics_into_result_file(_make_args(tmp_path), FAKE_METRICS) data = json.loads(result_file.read_text()) - expected_keys = {"spyre_" + k for k in FAKE_METRICS[0]} | { + expected_keys = {"spyre_" + k for k in FAKE_METRIC_KEYS} | { "spyre_total_prefill_chunks", "spyre_num_requests_missing_blocks", } for key in expected_keys: assert key in data, f"expected key {key!r} missing from result JSON" + # The alignment-only stamp must not leak into the written columns. + assert "spyre_" + _SPYRE_START_TIME_KEY not in data @pytest.mark.cpu @@ -126,7 +148,7 @@ def test_inject_values_correct(tmp_path): data = json.loads(result_file.read_text()) # Per-request lists map directly: "spyre_" + key → [m[key] for m in FAKE_METRICS] - for key in FAKE_METRICS[0]: + for key in FAKE_METRIC_KEYS: if key == "num_chunked_prefills": continue expected = [m[key] for m in FAKE_METRICS] @@ -142,6 +164,63 @@ def test_inject_values_correct(tmp_path): assert data["backend"] == "spyre-chat" +@pytest.mark.cpu +def test_inject_reorders_by_start_time(tmp_path): + """The core feature: metrics are collected in completion order but must be + written in the result's submission order, keyed by start_time — so a reversed + collection order still lands aligned with the rows.""" + # start_times in submission order match FAKE_METRICS[0], then [1]. + result_file = _write_fake_result(tmp_path) + # Hand the metrics in the opposite (completion) order; the stamp must drive + # placement, not list position. + _inject_spyre_metrics_into_result_file(_make_args(tmp_path), list(reversed(FAKE_METRICS))) + data = json.loads(result_file.read_text()) + + # queued_time_s is distinct per request, so it pins the ordering unambiguously. + assert data["spyre_queued_time_s"] == [ + FAKE_METRICS[0]["queued_time_s"], + FAKE_METRICS[1]["queued_time_s"], + ] + + +@pytest.mark.cpu +def test_inject_pads_unmatched_rows_with_sentinels(tmp_path): + """A result row with no matching collected metric (e.g. a failed request that + never produced Spyre metrics) gets a sentinel, keeping every spyre_* array the + same length as start_times/ttfts.""" + # Three submission-ordered rows; only the first and third have metrics. + result_file = _write_fake_result( + tmp_path, start_times=[FAKE_START_TIMES[0], 999.999, FAKE_START_TIMES[1]] + ) + _inject_spyre_metrics_into_result_file(_make_args(tmp_path), FAKE_METRICS) + data = json.loads(result_file.read_text()) + + assert len(data["spyre_queued_time_s"]) == 3 + assert data["spyre_queued_time_s"] == [ + FAKE_METRICS[0]["queued_time_s"], + None, # unmatched row → sentinel + FAKE_METRICS[1]["queued_time_s"], + ] + # List-valued columns get their [] sentinel, not None. + assert data["spyre_tkvs"][1] == [] + + +@pytest.mark.cpu +def test_inject_skips_when_no_start_times(tmp_path, caplog_sendnn_inference): + """Without start_times there is no key to realign the completion-ordered, + metrics-only list against vllm's submission-ordered arrays, so injection must + refuse to write per-request columns rather than emit misaligned data.""" + result_file = _write_fake_result(tmp_path, start_times=None) + original = result_file.read_text() + with caplog_sendnn_inference.at_level("WARNING"): + _inject_spyre_metrics_into_result_file(_make_args(tmp_path), FAKE_METRICS) + # File untouched, and no spyre_* columns written. + data = json.loads(result_file.read_text()) + assert result_file.read_text() == original + assert not any(k.startswith("spyre_") for k in data) + assert "no 'start_times'" in caplog_sendnn_inference.text + + @pytest.mark.cpu def test_inject_noop_when_save_result_false(tmp_path): result_file = _write_fake_result(tmp_path) From 55acec1798dc8598783e82200a035b191c9716d2 Mon Sep 17 00:00:00 2001 From: Yannick Schnider Date: Fri, 14 Aug 2026 16:11:30 +0200 Subject: [PATCH 105/106] tolerate present-but-null metric values in _print_spyre_section Signed-off-by: Yannick Schnider --- .../benchmarks/spyre_bench_serve.py | 38 +++++++++++++------ tests/benchmarks/test_bench_metrics.py | 14 +++++++ 2 files changed, 40 insertions(+), 12 deletions(-) diff --git a/sendnn_inference/benchmarks/spyre_bench_serve.py b/sendnn_inference/benchmarks/spyre_bench_serve.py index 9195f795d..d34a8724b 100644 --- a/sendnn_inference/benchmarks/spyre_bench_serve.py +++ b/sendnn_inference/benchmarks/spyre_bench_serve.py @@ -332,14 +332,20 @@ def _print_spyre_section( if not metrics_list: return - queue_times_ms = [m["queued_time_s"] * 1000 for m in metrics_list if "queued_time_s" in m] + # Values may be present-but-None (sentinel rows); skip nulls for scalars and + # coerce None to [] for lists so a single null never aborts the whole section. + queue_times_ms = [ + m["queued_time_s"] * 1000 for m in metrics_list if m.get("queued_time_s") is not None + ] num_chunks_list = [ - m["num_chunked_prefills"] for m in metrics_list if "num_chunked_prefills" in m + m["num_chunked_prefills"] for m in metrics_list if m.get("num_chunked_prefills") is not None ] chunk_lats_ms = [ - lat * 1000 for m in metrics_list for lat in m.get("chunk_prefill_latencies_s", []) + lat * 1000 for m in metrics_list for lat in (m.get("chunk_prefill_latencies_s") or []) + ] + decode_lats_ms = [ + lat * 1000 for m in metrics_list for lat in (m.get("decode_latencies_s") or []) ] - decode_lats_ms = [lat * 1000 for m in metrics_list for lat in m.get("decode_latencies_s", [])] total_prefill_chunks = sum(num_chunks_list) total_missing_blocks = sum(1 for m in metrics_list if m.get("was_missing_blocks", False)) @@ -364,19 +370,27 @@ def _section(header: str, values: list[float], label: str) -> None: ) cache_hit_pcts = [ - m["prefix_cache_hit_pct"] * 100 for m in metrics_list if "prefix_cache_hit_pct" in m + m["prefix_cache_hit_pct"] * 100 + for m in metrics_list + if m.get("prefix_cache_hit_pct") is not None ] prefill_elapsed_ms = [ - m["prefill_elapsed_s"] * 1000 for m in metrics_list if "prefill_elapsed_s" in m + m["prefill_elapsed_s"] * 1000 + for m in metrics_list + if m.get("prefill_elapsed_s") is not None + ] + prefill_busy_ms = [ + m["prefill_busy_s"] * 1000 for m in metrics_list if m.get("prefill_busy_s") is not None + ] + prefill_idle_ms = [ + m["prefill_idle_s"] * 1000 for m in metrics_list if m.get("prefill_idle_s") is not None ] - prefill_busy_ms = [m["prefill_busy_s"] * 1000 for m in metrics_list if "prefill_busy_s" in m] - prefill_idle_ms = [m["prefill_idle_s"] * 1000 for m in metrics_list if "prefill_idle_s" in m] - left_padding_blocks = [v for m in metrics_list for v in m.get("left_padding_blocks", [])] - pause_lats_ms = [lat * 1000 for m in metrics_list for lat in m.get("pause_latencies_s", [])] - pause_counts = [float(len(m.get("pause_latencies_s", []))) for m in metrics_list] - total_pause_ms = [float(sum(m.get("pause_latencies_s", []))) * 1000 for m in metrics_list] + left_padding_blocks = [v for m in metrics_list for v in (m.get("left_padding_blocks") or [])] + pause_lats_ms = [lat * 1000 for m in metrics_list for lat in (m.get("pause_latencies_s") or [])] + pause_counts = [float(len(m.get("pause_latencies_s") or [])) for m in metrics_list] + total_pause_ms = [float(sum(m.get("pause_latencies_s") or [])) * 1000 for m in metrics_list] _section("Queue Wait Time", queue_times_ms, "Queue Wait Time (ms)") _section("Chunked Prefill Count", num_chunks_list, "Num Chunked Prefills") diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index c93cbf738..4ae82e17f 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -473,6 +473,20 @@ def test_print_missing_keys_show_zeros(capsys): assert "0.00" in out +@pytest.mark.cpu +def test_print_tolerates_null_values(capsys): + # Present-but-None values (scalars and lists) must not abort the section. + metrics = [ + {k: None for k in FAKE_METRIC_KEYS}, + FAKE_METRICS[0], + ] + _print_spyre_section(metrics, SELECTED_PERCENTILES) + out = capsys.readouterr().out + # Non-null row still contributes its data; section completes to the end marker. + assert "Queue Wait Time" in out + assert out.rstrip().endswith("=" * 50) + + # --------------------------------------------------------------------------- # Test 3B — get_and_clear_chunk_stats (pure unit, no engine) # --------------------------------------------------------------------------- From 81b8c47801f7eb364506139ebf8474e88cc6cbb9 Mon Sep 17 00:00:00 2001 From: Yannick Schnider Date: Mon, 17 Aug 2026 10:59:47 +0200 Subject: [PATCH 106/106] Update tests/benchmarks/test_bench_metrics.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Sophie du Couédic Signed-off-by: Yannick Schnider --- tests/benchmarks/test_bench_metrics.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/benchmarks/test_bench_metrics.py b/tests/benchmarks/test_bench_metrics.py index 4ae82e17f..88dca4841 100644 --- a/tests/benchmarks/test_bench_metrics.py +++ b/tests/benchmarks/test_bench_metrics.py @@ -30,10 +30,7 @@ # Shared test data # --------------------------------------------------------------------------- -# Per-request start_time stamps carried on each collected metric dict. vllm's -# result arrays (ttfts/itls/start_times) are submission-ordered, so injection -# realigns metrics against them by this key. The values are distinct so the -# reorder is observable. +# Per-request start_time stamps carried on each collected metric dict FAKE_START_TIMES = [111.111, 222.222] # Synthetic values. FAKE_METRICS[i] carries FAKE_START_TIMES[i] under