Repository: vllm-project/vllm-spyre
Pull request #26: Use execute_model for warmup
Author: login: rafvasq
Created: 2025-03-13T15:26:46Z | Merged: 2025-03-21T15:52:35Z

PR to use `execute_model` in the model warmup `v1.worker.SpyreWorker.compile_or_warm_up_model` instead of a separate dummy forward pass. 

Closes #12 

Linked issue #12: [V1] Use the "real" model forward pass for warmup
Currently we implement a separate dummy forward pass for the model warmup in `v1.worker.SpyreWorker.compile_or_warm_up_model`

If we can instead figure out how to use `execute_model` during the warmup, we can both reduce the code required here and ensure that the torch compilation that runs at warmup is building the correct graphs. 🤞 this is a simple refactor, but I have not looked into it to know about any potential blockers.
Linked issue comments: user: login: joerunde
created_at: 2025-03-07T22:31:01Z
body: @rafvasq is taking a look at this, we just don't have the ability to assign him yet
author_association: COLLABORATOR

Diff:
diff --git a/vllm_spyre/v1/worker/spyre_model_runner.py b/vllm_spyre/v1/worker/spyre_model_runner.py
index 45d63bccd..76ad21fd8 100644
--- a/vllm_spyre/v1/worker/spyre_model_runner.py
+++ b/vllm_spyre/v1/worker/spyre_model_runner.py
@@ -459,27 +459,6 @@ def pad_input_ids(
 
         return input_ids, position_ids, mask
 
-    def _raw_model_forward(
-        self,
-        input_ids: torch.Tensor,
-        mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        past_key_value_states: Optional[List[Tuple[torch.Tensor,
-                                                   torch.Tensor]]] = None,
-        use_cache: bool = False,
-        only_last_token: bool = False,
-        attn_algorithm: Optional[str] = None
-    ) -> Tuple[torch.Tensor, Optional[List[Tuple[torch.Tensor,
-                                                 torch.Tensor]]]]:
-
-        return self.model.model(input_ids,
-                                mask=mask,
-                                position_ids=position_ids,
-                                past_key_value_states=past_key_value_states,
-                                use_cache=use_cache,
-                                only_last_token=only_last_token,
-                                attn_algorithm=attn_algorithm)
-
     def get_kv_cache_spec(self) -> KVCacheSpec:
         """
         This method should generate the KVCache spec by parsing the kv cache
diff --git a/vllm_spyre/v1/worker/spyre_worker.py b/vllm_spyre/v1/worker/spyre_worker.py
index 78c06a1bb..6fca9cd37 100644
--- a/vllm_spyre/v1/worker/spyre_worker.py
+++ b/vllm_spyre/v1/worker/spyre_worker.py
@@ -14,7 +14,9 @@
 from vllm.logger import init_logger
 from vllm.model_executor import set_random_seed
 from vllm.platforms import current_platform
-from vllm.v1.core.scheduler import SchedulerOutput
+from vllm.sampling_params import SamplingParams
+from vllm.v1.core.scheduler import (CachedRequestData, NewRequestData,
+                                    SchedulerOutput)
 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.worker.worker_base import WorkerBase as WorkerBaseV1
@@ -233,12 +235,7 @@ def load_model(self):
 
     def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
                                  special_token_ids, batch_size):
-        # TODO See if we can use `self.execute_model` instead for the warmup
-        # It's slightly risky to implement different forward pass logic here,
-        # which can go out of sync with the real forward pass and cause problems
-        # for torch.compile
 
-        # warmup the model
         warmup_start_t = time.time()
         # NOTE(ngl): empty tensor causes spyre to hang, so using
         # randint without 0 and the eos and bos token
@@ -254,6 +251,7 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
         valid_token_ids_tensor = torch.tensor(valid_token_ids,
                                               dtype=torch.long,
                                               device=torch.device("cpu"))
+
         # Sample from the valid token ids
         warmup_tokens_tensor = valid_token_ids_tensor[torch.randint(
             0, len(valid_token_ids_tensor), (batch_size, prompt_len))]
@@ -266,82 +264,93 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
             # impl when padding too much
             extra_kwargs["attn_algorithm"] = "math"
 
-        print(f"[SpyreWorker] warmup for prompt length "
-              f"{prompt_len} and max output tokens {num_decode_tokens}.")
+        # Set up dummy requests for prefill steps
+        dummy_requests = [
+            NewRequestData(
+                req_id="warmup",
+                prompt_token_ids=warmup_tokens_tensor[i].tolist(),
+                prompt="test",
+                mm_inputs=[],
+                mm_hashes=[],
+                mm_positions=[],
+                sampling_params=SamplingParams(max_tokens=num_decode_tokens),
+                block_ids=[0],
+                num_computed_tokens=0,
+                lora_request=None,
+            ) for i in range(batch_size)
+        ]
+
+        # Set up dummy cached_requests for decode steps
+        cached_requests = [
+            CachedRequestData(
+                req_id=req.req_id,
+                resumed_from_preemption=False,
+                new_token_ids=[
+                    valid_token_ids_tensor[torch.randint(
+                        0, len(valid_token_ids_tensor), (1, )).item()]
+                ],  # placeholder token
+                new_block_ids=req.block_ids,
+                num_computed_tokens=req.num_computed_tokens,
+            ) for req in dummy_requests
+        ]
 
-        # 1. trace
-        print("[SpyreWorker] warmup 1/2...")
-        # TODO: torch_sendnn.CleanGraph() should be necessary?
-        # warmup 1st forward pass
-        self._warmup_model_forward_pass(warmup_tokens_tensor,
-                                        valid_token_ids_tensor, prompt_len,
-                                        num_decode_tokens, batch_size,
-                                        extra_kwargs)
+        # Set up scheduler_output for execute_model
+        scheduler_output = SchedulerOutput(
+            scheduled_new_reqs=dummy_requests,
+            scheduled_cached_reqs=[],
+            num_scheduled_tokens={i: prompt_len
+                                  for i in range(batch_size)},
+            total_num_scheduled_tokens=sum(prompt_len
+                                           for _ in range(batch_size)),
+            scheduled_spec_decode_tokens={},
+            scheduled_encoder_inputs={},
+            num_common_prefix_blocks=0,
+            finished_req_ids=set(),
+            free_encoder_input_ids=[],
+        )
+
+        # First full forward pass
+        logger.info("Warmup forward pass 1/2...")
+        self._warmup_model_forward_pass(scheduler_output, dummy_requests,
+                                        cached_requests, num_decode_tokens)
 
-        # 2. compile
-        print("[SpyreWorker] warmup 2/2...")
         if envs_spyre.VLLM_SPYRE_DYNAMO_BACKEND == "sendnn_decoder":
             from torch_sendnn import torch_sendnn
             ul_start_time = time.time()
             torch_sendnn.update_lazyhandle()
             ul_stop_time = time.time()
-            ul_total_t = ul_stop_time - ul_start_time
-            print(f"update_lazyhandle() done (duration: {ul_total_t}s)")
+            logger.info("update_lazyhandle() done (duration: %.3fs)",
+                        ul_stop_time - ul_start_time)
 
-        # warmup 2nd forward pass
-        self._warmup_model_forward_pass(warmup_tokens_tensor,
-                                        valid_token_ids_tensor, prompt_len,
-                                        num_decode_tokens, batch_size,
-                                        extra_kwargs)
+        # Second full forward pass
+        logger.info("Warmup forward pass 2/2...")
+        self._warmup_model_forward_pass(scheduler_output, dummy_requests,
+                                        cached_requests, num_decode_tokens)
 
         warmup_end_t = time.time()
         warmup_total_t = warmup_end_t - warmup_start_t
-        print("[SpyreWorker] ... warmup finished.")
-        print(f"\twarmup took {warmup_total_t}s (for prompt length"
-              f"{prompt_len} and max output tokens {num_decode_tokens})")
-
-    def _warmup_model_forward_pass(self, warmup_tokens_tensor,
-                                   valid_token_ids_tensor, prompt_len,
-                                   num_decode_tokens, batch_size,
-                                   extra_kwargs):
-        # padding warmup tokens to obtain the
-        # corresponding position ids and mask
-        warmup_tokens_pad, self.model_runner._position_ids, \
-        self.model_runner._mask = self.model_runner.pad_input_ids(
-            warmup_tokens_tensor, min_pad_length=prompt_len)
-
-        logits, past_key_value_states = self.model_runner._raw_model_forward(
-            warmup_tokens_pad,
-            position_ids=self.model_runner._position_ids,
-            mask=self.model_runner._mask,
-            past_key_value_states=None,
-            use_cache=True,
-            only_last_token=True,
-            **extra_kwargs)
-        # decoding
-        for i in range(num_decode_tokens - 1):
-            # sampling next input token from vocab without bos and eos tokens
-            decode_tokens = valid_token_ids_tensor[torch.randint(
-                0, len(valid_token_ids_tensor), (batch_size, 1))]
-
-            # update mask and position_ids
-            self.model_runner._update_mask()
-            self.model_runner._update_position_ids()
-
-            if past_key_value_states is not None:
-                for layer in past_key_value_states:
-                    for tensor in layer:
-                        torch._dynamo.mark_dynamic(tensor, 2)
-
-            logits, past_key_value_states = self.model_runner.\
-                _raw_model_forward(
-                decode_tokens,
-                position_ids=self.model_runner._position_ids,
-                mask=self.model_runner._mask,
-                past_key_value_states=past_key_value_states,
-                use_cache=True,
-                only_last_token=True,
-                **extra_kwargs)
+        logger.info("Warmup finished.")
+        logger.info(
+            "Warmup took %.3fs (for prompt length %d and max output tokens %d)",
+            warmup_total_t, prompt_len, num_decode_tokens)
+
+    def _warmup_model_forward_pass(
+        self,
+        scheduler_output: SchedulerOutput,
+        requests: List[NewRequestData],
+        cached_requests: List[CachedRequestData],
+        num_decode_tokens,
+    ):
+        """Handle a complete forward pass"""
+        scheduler_output.scheduled_new_reqs = requests
+        scheduler_output.scheduled_cached_reqs = []
+        self.execute_model(scheduler_output)  # Prefill
+
+        # Switch to cached requests to trigger decoding steps
+        scheduler_output.scheduled_new_reqs = []
+        scheduler_output.scheduled_cached_reqs = cached_requests
+        for _ in range(num_decode_tokens - 1):
+            self.execute_model(scheduler_output)
 
     @property
     def do_metadata_broadcast(self) -> bool:


Reviews:
user: login: yannicks1
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-03-15T11:15:52Z
user: login: joerunde
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-17T22:27:28Z
user: login: joerunde
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-17T22:34:04Z
user: login: rafvasq
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-19T18:43:55Z
user: login: joerunde
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-19T19:10:14Z
user: login: joerunde
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-19T19:10:35Z
user: login: rafvasq
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-19T19:18:54Z
user: login: joerunde
body: LPGTM!

I'll let @yannicks1 take another look through since he has a better understanding of the requirements
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-19T19:56:24Z
user: login: yannicks1
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-20T10:10:15Z
user: login: yannicks1
body: Other than the two minor comments this looks good to me! Thanks for contributing!

PS: I am mildly confused why the spyre-tests did not succeed. 
`tests/test_spyre_embeddings.py` fails with V0. But as your code changes only touch V1, the reason has got to be somewhere else... 
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-20T10:20:02Z
user: login: rafvasq
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-20T16:22:38Z
user: login: yannicks1
body: LGTM! Thanks for contributing. 
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-21T15:52:09Z

Review comments:
diff_hunk: @@ -68,8 +69,51 @@ def compile_or_warm_up_model(self) -> None:
             print(f"[SpyreWorker] Warming up for prompt length {prompt_len}, "
                   f"decoding {num_decode_tokens} tokens with batch "
                   f"size {batch_size}")
-            self._warmup_spyre_fixed_size(prompt_len, num_decode_tokens,
path: vllm_spyre/v1/worker/spyre_worker.py
user: login: yannicks1
body: personally, I like the idea of a helper function here to make things more readable.
created_at: 2025-03-15T10:58:01Z
author_association: COLLABORATOR
diff_hunk: @@ -68,8 +69,51 @@ def compile_or_warm_up_model(self) -> None:
             print(f"[SpyreWorker] Warming up for prompt length {prompt_len}, "
                   f"decoding {num_decode_tokens} tokens with batch "
                   f"size {batch_size}")
-            self._warmup_spyre_fixed_size(prompt_len, num_decode_tokens,
-                                          self.restricted_tokens, batch_size)
+
+            num_scheduled_tokens: dict = {}
+            total_num_scheduled_tokens: int = 0
+            dummy_requests: list = []
+            for i in range(batch_size):
+                dummy_requests.append(
+                    NewRequestData(
+                        req_id=f"warmup-{i}",
+                        prompt_token_ids=[1] * prompt_len,
path: vllm_spyre/v1/worker/spyre_worker.py
user: login: yannicks1
body: We previously sampled these tokens from a list of valid tokens: 
```
valid_token_ids = [ i for i in range(1, vocab_size) if i not in set(special_token_ids)]
```
where `special_token_ids` contains BOS, EOS and pad token ids. Not sure whether this was needed, or what happens if in your case any of the above special token ids are 1. 
created_at: 2025-03-15T11:02:23Z
author_association: COLLABORATOR
diff_hunk: @@ -68,8 +69,51 @@ def compile_or_warm_up_model(self) -> None:
             print(f"[SpyreWorker] Warming up for prompt length {prompt_len}, "
                   f"decoding {num_decode_tokens} tokens with batch "
                   f"size {batch_size}")
-            self._warmup_spyre_fixed_size(prompt_len, num_decode_tokens,
-                                          self.restricted_tokens, batch_size)
+
+            num_scheduled_tokens: dict = {}
+            total_num_scheduled_tokens: int = 0
+            dummy_requests: list = []
+            for i in range(batch_size):
+                dummy_requests.append(
+                    NewRequestData(
+                        req_id=f"warmup-{i}",
+                        prompt_token_ids=[1] * prompt_len,
+                        prompt="test",
+                        mm_inputs=[],
+                        mm_hashes=[],
+                        mm_positions=[],
+                        sampling_params=SamplingParams(
+                            max_tokens=num_decode_tokens),
+                        block_ids=[0],
+                        num_computed_tokens=0,
+                        lora_request=None,
+                    ))
+                num_scheduled_tokens[i] = prompt_len
+                total_num_scheduled_tokens += num_scheduled_tokens[i]
+
+            scheduler_output = SchedulerOutput(
+                scheduled_new_reqs=dummy_requests,
+                scheduled_cached_reqs=[],
+                num_scheduled_tokens=num_scheduled_tokens,
+                total_num_scheduled_tokens=total_num_scheduled_tokens,
+                scheduled_spec_decode_tokens={},
+                scheduled_encoder_inputs={},
+                num_common_prefix_blocks=0,
+                finished_req_ids=set(),
+                free_encoder_input_ids=[],
+            )
+
+            # Use execute_model for warm up
+            self.execute_model(scheduler_output)
path: vllm_spyre/v1/worker/spyre_worker.py
user: login: yannicks1
body: As far as I understand this executes the model only for one step. In your case it does the prefill and generates the 1st token. We need to warm up not only for prefill, but also for `(num_decode_tokens - 1)` decode steps (since prefill already produced 1 token).  
created_at: 2025-03-15T11:08:28Z
author_association: COLLABORATOR
diff_hunk: @@ -68,8 +69,51 @@ def compile_or_warm_up_model(self) -> None:
             print(f"[SpyreWorker] Warming up for prompt length {prompt_len}, "
                   f"decoding {num_decode_tokens} tokens with batch "
                   f"size {batch_size}")
-            self._warmup_spyre_fixed_size(prompt_len, num_decode_tokens,
-                                          self.restricted_tokens, batch_size)
+
+            num_scheduled_tokens: dict = {}
+            total_num_scheduled_tokens: int = 0
+            dummy_requests: list = []
+            for i in range(batch_size):
+                dummy_requests.append(
+                    NewRequestData(
+                        req_id=f"warmup-{i}",
+                        prompt_token_ids=[1] * prompt_len,
+                        prompt="test",
+                        mm_inputs=[],
+                        mm_hashes=[],
+                        mm_positions=[],
+                        sampling_params=SamplingParams(
+                            max_tokens=num_decode_tokens),
+                        block_ids=[0],
+                        num_computed_tokens=0,
+                        lora_request=None,
+                    ))
+                num_scheduled_tokens[i] = prompt_len
+                total_num_scheduled_tokens += num_scheduled_tokens[i]
+
+            scheduler_output = SchedulerOutput(
+                scheduled_new_reqs=dummy_requests,
+                scheduled_cached_reqs=[],
+                num_scheduled_tokens=num_scheduled_tokens,
+                total_num_scheduled_tokens=total_num_scheduled_tokens,
+                scheduled_spec_decode_tokens={},
+                scheduled_encoder_inputs={},
+                num_common_prefix_blocks=0,
+                finished_req_ids=set(),
+                free_encoder_input_ids=[],
+            )
+
+            # Use execute_model for warm up
+            self.execute_model(scheduler_output)
+
+            if envs_spyre.VLLM_SPYRE_DYNAMO_BACKEND == "sendnn_decoder":
+                from torch_sendnn import torch_sendnn
+                ul_start_time = time.time()
+                torch_sendnn.update_lazyhandle()
+                ul_stop_time = time.time()
+                ul_total_t = ul_stop_time - ul_start_time
+                print(f"update_lazyhandle() done (duration: {ul_total_t}s)")
+
path: vllm_spyre/v1/worker/spyre_worker.py
user: login: yannicks1
body: After `torch_sendnn.update_lazyhandle()` there is a second _**complete**_ warmup needed. 

To sum up: 

1. complete forward pass: prefill plus `(num_decode_tokens - 1)` decode steps
2. torch_sendnn.update_lazyhandle()
3. complete forward pass: prefill plus `(num_decode_tokens - 1)` decode steps

See comment below. 

created_at: 2025-03-15T11:12:59Z
author_association: COLLABORATOR
diff_hunk: @@ -209,118 +253,6 @@ def load_model(self):
         load_model_total_t = load_model_end_t - load_model_start_t
         print(f"\tload model took {load_model_total_t}s")
 
-    def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
-                                 special_token_ids, batch_size):
-        # TODO See if we can use `self.execute_model` instead for the warmup
-        # It's slightly risky to implement different forward pass logic here,
-        # which can go out of sync with the real forward pass and cause problems
-        # for torch.compile
-
-        # warmup the model
-        warmup_start_t = time.time()
-        # NOTE(ngl): empty tensor causes spyre to hang, so using
-        # randint without 0 and the eos and bos token
-
-        # Create a list of valid values between 1 (inclusive) and vocab
-        # size (exclusive) by excluding the eos and bos token ids
-        # (in special_token_ids)
-        vocab_size = self.model_runner.vocab_size
-        valid_token_ids = [
-            i for i in range(1, vocab_size) if i not in set(special_token_ids)
-        ]
-        # Convert to tensor for sampling
-        valid_token_ids_tensor = torch.tensor(valid_token_ids,
-                                              dtype=torch.long,
-                                              device=torch.device("cpu"))
-        # Sample from the valid token ids
-        warmup_tokens_tensor = valid_token_ids_tensor[torch.randint(
-            0, len(valid_token_ids_tensor), (batch_size, prompt_len))]
-
-        extra_kwargs = {}
-        if envs_spyre.VLLM_SPYRE_DYNAMO_BACKEND not in [
-                "sendnn", "sendnn_decoder"
-        ]:
-            # Bug in 2.3.1 fixed in 2.4.1 for SDPA flash cpu
-            # impl when padding too much
-            extra_kwargs["attn_algorithm"] = "math"
-
-        print(f"[SpyreWorker] warmup for prompt length "
-              f"{prompt_len} and max output tokens {num_decode_tokens}.")
-
-        # 1. trace
-        print("[SpyreWorker] warmup 1/2...")
-        # TODO: torch_sendnn.CleanGraph() should be necessary?
-        # warmup 1st forward pass
-        self._warmup_model_forward_pass(warmup_tokens_tensor,
-                                        valid_token_ids_tensor, prompt_len,
-                                        num_decode_tokens, batch_size,
-                                        extra_kwargs)
-
-        # 2. compile
-        print("[SpyreWorker] warmup 2/2...")
-        if envs_spyre.VLLM_SPYRE_DYNAMO_BACKEND == "sendnn_decoder":
-            from torch_sendnn import torch_sendnn
-            ul_start_time = time.time()
-            torch_sendnn.update_lazyhandle()
-            ul_stop_time = time.time()
-            ul_total_t = ul_stop_time - ul_start_time
-            print(f"update_lazyhandle() done (duration: {ul_total_t}s)")
-
-        # warmup 2nd forward pass
-        self._warmup_model_forward_pass(warmup_tokens_tensor,
-                                        valid_token_ids_tensor, prompt_len,
-                                        num_decode_tokens, batch_size,
-                                        extra_kwargs)
path: vllm_spyre/v1/worker/spyre_worker.py
user: login: yannicks1
body: this has to happen for each warmup shape (combination of `prompt_len`, `num_decode_tokens`, `batch_size`)
created_at: 2025-03-15T11:14:55Z
author_association: COLLABORATOR
diff_hunk: @@ -68,8 +69,51 @@ def compile_or_warm_up_model(self) -> None:
             print(f"[SpyreWorker] Warming up for prompt length {prompt_len}, "
                   f"decoding {num_decode_tokens} tokens with batch "
                   f"size {batch_size}")
-            self._warmup_spyre_fixed_size(prompt_len, num_decode_tokens,
path: vllm_spyre/v1/worker/spyre_worker.py
user: login: joerunde
body: +1, specifically avoiding double-nested for loops is nice to do
created_at: 2025-03-17T22:27:28Z
author_association: COLLABORATOR
diff_hunk: @@ -68,8 +69,51 @@ def compile_or_warm_up_model(self) -> None:
             print(f"[SpyreWorker] Warming up for prompt length {prompt_len}, "
                   f"decoding {num_decode_tokens} tokens with batch "
                   f"size {batch_size}")
-            self._warmup_spyre_fixed_size(prompt_len, num_decode_tokens,
-                                          self.restricted_tokens, batch_size)
+
+            num_scheduled_tokens: dict = {}
+            total_num_scheduled_tokens: int = 0
+            dummy_requests: list = []
+            for i in range(batch_size):
+                dummy_requests.append(
+                    NewRequestData(
+                        req_id=f"warmup-{i}",
+                        prompt_token_ids=[1] * prompt_len,
path: vllm_spyre/v1/worker/spyre_worker.py
user: login: joerunde
body: Hah, for the very first dummy model I checked this is true: https://huggingface.co/JackFram/llama-160m/blob/main/config.json#L6

I think vllm uses a bunch of repeated token id 0 for profiling, since the input ids tensor is just initialized with `torch.zeros` and for text-only models it's not updated for profiling. 

The general idea with setting a repeated token ID was to have the model continue the sequence, so it doesn't end up hitting an eos token early and stopping. But if we control the loop here that keeps invoking the model, maybe that doesn't matter.
created_at: 2025-03-17T22:34:04Z
author_association: COLLABORATOR
diff_hunk: @@ -68,8 +69,51 @@ def compile_or_warm_up_model(self) -> None:
             print(f"[SpyreWorker] Warming up for prompt length {prompt_len}, "
                   f"decoding {num_decode_tokens} tokens with batch "
                   f"size {batch_size}")
-            self._warmup_spyre_fixed_size(prompt_len, num_decode_tokens,
-                                          self.restricted_tokens, batch_size)
+
+            num_scheduled_tokens: dict = {}
+            total_num_scheduled_tokens: int = 0
+            dummy_requests: list = []
+            for i in range(batch_size):
+                dummy_requests.append(
+                    NewRequestData(
+                        req_id=f"warmup-{i}",
+                        prompt_token_ids=[1] * prompt_len,
+                        prompt="test",
+                        mm_inputs=[],
+                        mm_hashes=[],
+                        mm_positions=[],
+                        sampling_params=SamplingParams(
+                            max_tokens=num_decode_tokens),
+                        block_ids=[0],
+                        num_computed_tokens=0,
+                        lora_request=None,
+                    ))
+                num_scheduled_tokens[i] = prompt_len
+                total_num_scheduled_tokens += num_scheduled_tokens[i]
+
+            scheduler_output = SchedulerOutput(
+                scheduled_new_reqs=dummy_requests,
+                scheduled_cached_reqs=[],
+                num_scheduled_tokens=num_scheduled_tokens,
+                total_num_scheduled_tokens=total_num_scheduled_tokens,
+                scheduled_spec_decode_tokens={},
+                scheduled_encoder_inputs={},
+                num_common_prefix_blocks=0,
+                finished_req_ids=set(),
+                free_encoder_input_ids=[],
+            )
+
+            # Use execute_model for warm up
+            self.execute_model(scheduler_output)
+
+            if envs_spyre.VLLM_SPYRE_DYNAMO_BACKEND == "sendnn_decoder":
+                from torch_sendnn import torch_sendnn
+                ul_start_time = time.time()
+                torch_sendnn.update_lazyhandle()
+                ul_stop_time = time.time()
+                ul_total_t = ul_stop_time - ul_start_time
+                print(f"update_lazyhandle() done (duration: {ul_total_t}s)")
+
path: vllm_spyre/v1/worker/spyre_worker.py
user: login: rafvasq
body: Thanks for clarifying this, learning as I go but it makes sense to me now.

I took another stab at it, still trying to use `execute_model` to avoid doing anything manually except the dummy data setup. 
created_at: 2025-03-19T18:43:54Z
author_association: COLLABORATOR
diff_hunk: @@ -253,95 +250,99 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
         # Convert to tensor for sampling
         valid_token_ids_tensor = torch.tensor(valid_token_ids,
                                               dtype=torch.long,
-                                              device=torch.device("cpu"))
+                                              device="cpu")
+
         # Sample from the valid token ids
         warmup_tokens_tensor = valid_token_ids_tensor[torch.randint(
             0, len(valid_token_ids_tensor), (batch_size, prompt_len))]
 
-        extra_kwargs = {}
-        if envs_spyre.VLLM_SPYRE_DYNAMO_BACKEND not in [
-                "sendnn", "sendnn_decoder"
-        ]:
-            # Bug in 2.3.1 fixed in 2.4.1 for SDPA flash cpu
-            # impl when padding too much
-            extra_kwargs["attn_algorithm"] = "math"
-
-        print(f"[SpyreWorker] warmup for prompt length "
-              f"{prompt_len} and max output tokens {num_decode_tokens}.")
-
-        # 1. trace
-        print("[SpyreWorker] warmup 1/2...")
-        # TODO: torch_sendnn.CleanGraph() should be necessary?
-        # warmup 1st forward pass
-        self._warmup_model_forward_pass(warmup_tokens_tensor,
-                                        valid_token_ids_tensor, prompt_len,
-                                        num_decode_tokens, batch_size,
-                                        extra_kwargs)
-
-        # 2. compile
-        print("[SpyreWorker] warmup 2/2...")
+        # Create requests to be used for prefill steps
+        dummy_requests = [
+            NewRequestData(
+                req_id="warmup",
+                prompt_token_ids=warmup_tokens_tensor[i].tolist(),
+                prompt="test",
+                mm_inputs=[],
+                mm_hashes=[],
+                mm_positions=[],
+                sampling_params=SamplingParams(max_tokens=num_decode_tokens),
+                block_ids=[0],
+                num_computed_tokens=0,
+                lora_request=None,
+            ) for i in range(batch_size)
+        ]
+
+        # Set up dummy cached_requests to be used for decode steps
+        cached_requests = [
+            CachedRequestData(
+                req_id=req.req_id,
+                resumed_from_preemption=False,
+                new_token_ids=[
+                    valid_token_ids_tensor[torch.randint(
+                        0, len(valid_token_ids_tensor), (1, )).item()]
+                ],  # placeholder token
+                new_block_ids=req.block_ids,
+                num_computed_tokens=req.num_computed_tokens,
+            ) for req in dummy_requests
+        ]
+
+        # To be used for execute_model, start with scheduled_new_reqs
+        # for prefill
+        scheduler_output = SchedulerOutput(
+            scheduled_new_reqs=dummy_requests,
+            scheduled_cached_reqs=[],
+            num_scheduled_tokens={i: prompt_len
+                                  for i in range(batch_size)},
+            total_num_scheduled_tokens=sum(prompt_len
+                                           for _ in range(batch_size)),
+            scheduled_spec_decode_tokens={},
+            scheduled_encoder_inputs={},
+            num_common_prefix_blocks=0,
+            finished_req_ids=set(),
+            free_encoder_input_ids=[],
+        )
+
+        # First full forward pass
+        logger.info("Warmup 1/2: Prefill...")
+        self.execute_model(scheduler_output)  # Prefill step
+
+        # Switch to cached requests to trigger decoding steps
+        scheduler_output.scheduled_new_reqs = []
+        scheduler_output.scheduled_cached_reqs = cached_requests
+
+        logger.info("Warmup 1/2: Decoding...")
+        for _ in range(num_decode_tokens - 1):
+            self.execute_model(scheduler_output)
+
+        # update_lazyhandle
         if envs_spyre.VLLM_SPYRE_DYNAMO_BACKEND == "sendnn_decoder":
             from torch_sendnn import torch_sendnn
             ul_start_time = time.time()
             torch_sendnn.update_lazyhandle()
             ul_stop_time = time.time()
-            ul_total_t = ul_stop_time - ul_start_time
-            print(f"update_lazyhandle() done (duration: {ul_total_t}s)")
+            logger.info("update_lazyhandle() done (duration: %.3fs",
path: vllm_spyre/v1/worker/spyre_worker.py
user: login: joerunde
body: ```suggestion
            logger.info("update_lazyhandle() done (duration: %.3fs)",
```
created_at: 2025-03-19T19:10:14Z
author_association: COLLABORATOR
diff_hunk: @@ -253,95 +250,99 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
         # Convert to tensor for sampling
         valid_token_ids_tensor = torch.tensor(valid_token_ids,
                                               dtype=torch.long,
-                                              device=torch.device("cpu"))
+                                              device="cpu")
+
         # Sample from the valid token ids
         warmup_tokens_tensor = valid_token_ids_tensor[torch.randint(
             0, len(valid_token_ids_tensor), (batch_size, prompt_len))]
 
-        extra_kwargs = {}
-        if envs_spyre.VLLM_SPYRE_DYNAMO_BACKEND not in [
-                "sendnn", "sendnn_decoder"
-        ]:
-            # Bug in 2.3.1 fixed in 2.4.1 for SDPA flash cpu
-            # impl when padding too much
-            extra_kwargs["attn_algorithm"] = "math"
-
-        print(f"[SpyreWorker] warmup for prompt length "
-              f"{prompt_len} and max output tokens {num_decode_tokens}.")
-
-        # 1. trace
-        print("[SpyreWorker] warmup 1/2...")
-        # TODO: torch_sendnn.CleanGraph() should be necessary?
-        # warmup 1st forward pass
-        self._warmup_model_forward_pass(warmup_tokens_tensor,
-                                        valid_token_ids_tensor, prompt_len,
-                                        num_decode_tokens, batch_size,
-                                        extra_kwargs)
-
-        # 2. compile
-        print("[SpyreWorker] warmup 2/2...")
+        # Create requests to be used for prefill steps
+        dummy_requests = [
+            NewRequestData(
+                req_id="warmup",
+                prompt_token_ids=warmup_tokens_tensor[i].tolist(),
+                prompt="test",
+                mm_inputs=[],
+                mm_hashes=[],
+                mm_positions=[],
+                sampling_params=SamplingParams(max_tokens=num_decode_tokens),
+                block_ids=[0],
+                num_computed_tokens=0,
+                lora_request=None,
+            ) for i in range(batch_size)
+        ]
+
+        # Set up dummy cached_requests to be used for decode steps
+        cached_requests = [
+            CachedRequestData(
+                req_id=req.req_id,
+                resumed_from_preemption=False,
+                new_token_ids=[
+                    valid_token_ids_tensor[torch.randint(
+                        0, len(valid_token_ids_tensor), (1, )).item()]
+                ],  # placeholder token
+                new_block_ids=req.block_ids,
+                num_computed_tokens=req.num_computed_tokens,
+            ) for req in dummy_requests
+        ]
+
+        # To be used for execute_model, start with scheduled_new_reqs
+        # for prefill
+        scheduler_output = SchedulerOutput(
+            scheduled_new_reqs=dummy_requests,
+            scheduled_cached_reqs=[],
+            num_scheduled_tokens={i: prompt_len
+                                  for i in range(batch_size)},
+            total_num_scheduled_tokens=sum(prompt_len
+                                           for _ in range(batch_size)),
+            scheduled_spec_decode_tokens={},
+            scheduled_encoder_inputs={},
+            num_common_prefix_blocks=0,
+            finished_req_ids=set(),
+            free_encoder_input_ids=[],
+        )
+
+        # First full forward pass
+        logger.info("Warmup 1/2: Prefill...")
+        self.execute_model(scheduler_output)  # Prefill step
+
+        # Switch to cached requests to trigger decoding steps
+        scheduler_output.scheduled_new_reqs = []
+        scheduler_output.scheduled_cached_reqs = cached_requests
+
+        logger.info("Warmup 1/2: Decoding...")
+        for _ in range(num_decode_tokens - 1):
+            self.execute_model(scheduler_output)
+
+        # update_lazyhandle
         if envs_spyre.VLLM_SPYRE_DYNAMO_BACKEND == "sendnn_decoder":
             from torch_sendnn import torch_sendnn
             ul_start_time = time.time()
             torch_sendnn.update_lazyhandle()
             ul_stop_time = time.time()
-            ul_total_t = ul_stop_time - ul_start_time
-            print(f"update_lazyhandle() done (duration: {ul_total_t}s)")
+            logger.info("update_lazyhandle() done (duration: %.3fs",
+                        ul_stop_time - ul_start_time)
+
+        # Second full forward pass
+        logger.info("Warmup 2/2: Prefill step...")
+        scheduler_output.scheduled_new_reqs = dummy_requests
+        scheduler_output.scheduled_cached_reqs = []
+        self.execute_model(scheduler_output)
 
-        # warmup 2nd forward pass
-        self._warmup_model_forward_pass(warmup_tokens_tensor,
-                                        valid_token_ids_tensor, prompt_len,
-                                        num_decode_tokens, batch_size,
-                                        extra_kwargs)
+        # Switch to cached requests to trigger decoding steps
+        scheduler_output.scheduled_new_reqs = []
+        scheduler_output.scheduled_cached_reqs = cached_requests
+
+        logger.info("[Warmup 2/2: Decoding steps...")
path: vllm_spyre/v1/worker/spyre_worker.py
user: login: joerunde
body: ```suggestion
        logger.info("Warmup 2/2: Decoding steps...")
```
created_at: 2025-03-19T19:10:35Z
author_association: COLLABORATOR
diff_hunk: @@ -253,95 +250,99 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
         # Convert to tensor for sampling
         valid_token_ids_tensor = torch.tensor(valid_token_ids,
                                               dtype=torch.long,
-                                              device=torch.device("cpu"))
+                                              device="cpu")
+
         # Sample from the valid token ids
         warmup_tokens_tensor = valid_token_ids_tensor[torch.randint(
             0, len(valid_token_ids_tensor), (batch_size, prompt_len))]
 
-        extra_kwargs = {}
-        if envs_spyre.VLLM_SPYRE_DYNAMO_BACKEND not in [
-                "sendnn", "sendnn_decoder"
-        ]:
-            # Bug in 2.3.1 fixed in 2.4.1 for SDPA flash cpu
-            # impl when padding too much
-            extra_kwargs["attn_algorithm"] = "math"
-
-        print(f"[SpyreWorker] warmup for prompt length "
-              f"{prompt_len} and max output tokens {num_decode_tokens}.")
-
-        # 1. trace
-        print("[SpyreWorker] warmup 1/2...")
-        # TODO: torch_sendnn.CleanGraph() should be necessary?
-        # warmup 1st forward pass
-        self._warmup_model_forward_pass(warmup_tokens_tensor,
-                                        valid_token_ids_tensor, prompt_len,
-                                        num_decode_tokens, batch_size,
-                                        extra_kwargs)
-
-        # 2. compile
-        print("[SpyreWorker] warmup 2/2...")
+        # Create requests to be used for prefill steps
+        dummy_requests = [
+            NewRequestData(
+                req_id="warmup",
+                prompt_token_ids=warmup_tokens_tensor[i].tolist(),
+                prompt="test",
+                mm_inputs=[],
+                mm_hashes=[],
+                mm_positions=[],
+                sampling_params=SamplingParams(max_tokens=num_decode_tokens),
+                block_ids=[0],
+                num_computed_tokens=0,
+                lora_request=None,
+            ) for i in range(batch_size)
+        ]
+
+        # Set up dummy cached_requests to be used for decode steps
+        cached_requests = [
+            CachedRequestData(
+                req_id=req.req_id,
+                resumed_from_preemption=False,
+                new_token_ids=[
+                    valid_token_ids_tensor[torch.randint(
+                        0, len(valid_token_ids_tensor), (1, )).item()]
+                ],  # placeholder token
+                new_block_ids=req.block_ids,
+                num_computed_tokens=req.num_computed_tokens,
+            ) for req in dummy_requests
+        ]
+
+        # To be used for execute_model, start with scheduled_new_reqs
+        # for prefill
+        scheduler_output = SchedulerOutput(
+            scheduled_new_reqs=dummy_requests,
+            scheduled_cached_reqs=[],
+            num_scheduled_tokens={i: prompt_len
+                                  for i in range(batch_size)},
+            total_num_scheduled_tokens=sum(prompt_len
+                                           for _ in range(batch_size)),
+            scheduled_spec_decode_tokens={},
+            scheduled_encoder_inputs={},
+            num_common_prefix_blocks=0,
+            finished_req_ids=set(),
+            free_encoder_input_ids=[],
+        )
+
+        # First full forward pass
+        logger.info("Warmup 1/2: Prefill...")
+        self.execute_model(scheduler_output)  # Prefill step
+
+        # Switch to cached requests to trigger decoding steps
+        scheduler_output.scheduled_new_reqs = []
+        scheduler_output.scheduled_cached_reqs = cached_requests
+
+        logger.info("Warmup 1/2: Decoding...")
+        for _ in range(num_decode_tokens - 1):
+            self.execute_model(scheduler_output)
+
+        # update_lazyhandle
         if envs_spyre.VLLM_SPYRE_DYNAMO_BACKEND == "sendnn_decoder":
             from torch_sendnn import torch_sendnn
             ul_start_time = time.time()
             torch_sendnn.update_lazyhandle()
             ul_stop_time = time.time()
-            ul_total_t = ul_stop_time - ul_start_time
-            print(f"update_lazyhandle() done (duration: {ul_total_t}s)")
+            logger.info("update_lazyhandle() done (duration: %.3fs",
path: vllm_spyre/v1/worker/spyre_worker.py
user: login: rafvasq
body: Thanks for catching these!
created_at: 2025-03-19T19:18:54Z
author_association: COLLABORATOR
diff_hunk: @@ -68,8 +69,51 @@ def compile_or_warm_up_model(self) -> None:
             print(f"[SpyreWorker] Warming up for prompt length {prompt_len}, "
                   f"decoding {num_decode_tokens} tokens with batch "
                   f"size {batch_size}")
-            self._warmup_spyre_fixed_size(prompt_len, num_decode_tokens,
-                                          self.restricted_tokens, batch_size)
+
+            num_scheduled_tokens: dict = {}
+            total_num_scheduled_tokens: int = 0
+            dummy_requests: list = []
+            for i in range(batch_size):
+                dummy_requests.append(
+                    NewRequestData(
+                        req_id=f"warmup-{i}",
+                        prompt_token_ids=[1] * prompt_len,
path: vllm_spyre/v1/worker/spyre_worker.py
user: login: yannicks1
body: yeah, I also suspect it does not matter as we force the decode steps in the loop, but rather save than sorry:) 
created_at: 2025-03-20T10:10:15Z
author_association: COLLABORATOR
diff_hunk: @@ -253,95 +250,99 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
         # Convert to tensor for sampling
         valid_token_ids_tensor = torch.tensor(valid_token_ids,
                                               dtype=torch.long,
-                                              device=torch.device("cpu"))
+                                              device="cpu")
+
         # Sample from the valid token ids
         warmup_tokens_tensor = valid_token_ids_tensor[torch.randint(
             0, len(valid_token_ids_tensor), (batch_size, prompt_len))]
 
-        extra_kwargs = {}
-        if envs_spyre.VLLM_SPYRE_DYNAMO_BACKEND not in [
-                "sendnn", "sendnn_decoder"
-        ]:
-            # Bug in 2.3.1 fixed in 2.4.1 for SDPA flash cpu
-            # impl when padding too much
-            extra_kwargs["attn_algorithm"] = "math"
-
-        print(f"[SpyreWorker] warmup for prompt length "
-              f"{prompt_len} and max output tokens {num_decode_tokens}.")
-
-        # 1. trace
-        print("[SpyreWorker] warmup 1/2...")
-        # TODO: torch_sendnn.CleanGraph() should be necessary?
-        # warmup 1st forward pass
-        self._warmup_model_forward_pass(warmup_tokens_tensor,
-                                        valid_token_ids_tensor, prompt_len,
-                                        num_decode_tokens, batch_size,
-                                        extra_kwargs)
-
-        # 2. compile
-        print("[SpyreWorker] warmup 2/2...")
+        # Create requests to be used for prefill steps
+        dummy_requests = [
+            NewRequestData(
+                req_id="warmup",
+                prompt_token_ids=warmup_tokens_tensor[i].tolist(),
+                prompt="test",
+                mm_inputs=[],
+                mm_hashes=[],
+                mm_positions=[],
+                sampling_params=SamplingParams(max_tokens=num_decode_tokens),
+                block_ids=[0],
+                num_computed_tokens=0,
+                lora_request=None,
+            ) for i in range(batch_size)
+        ]
+
+        # Set up dummy cached_requests to be used for decode steps
+        cached_requests = [
+            CachedRequestData(
+                req_id=req.req_id,
+                resumed_from_preemption=False,
+                new_token_ids=[
+                    valid_token_ids_tensor[torch.randint(
+                        0, len(valid_token_ids_tensor), (1, )).item()]
+                ],  # placeholder token
+                new_block_ids=req.block_ids,
+                num_computed_tokens=req.num_computed_tokens,
+            ) for req in dummy_requests
+        ]
+
+        # To be used for execute_model, start with scheduled_new_reqs
+        # for prefill
+        scheduler_output = SchedulerOutput(
+            scheduled_new_reqs=dummy_requests,
+            scheduled_cached_reqs=[],
+            num_scheduled_tokens={i: prompt_len
+                                  for i in range(batch_size)},
+            total_num_scheduled_tokens=sum(prompt_len
+                                           for _ in range(batch_size)),
+            scheduled_spec_decode_tokens={},
+            scheduled_encoder_inputs={},
+            num_common_prefix_blocks=0,
+            finished_req_ids=set(),
+            free_encoder_input_ids=[],
+        )
+
+        # First full forward pass
+        logger.info("Warmup 1/2: Prefill...")
+        self.execute_model(scheduler_output)  # Prefill step
+
+        # Switch to cached requests to trigger decoding steps
+        scheduler_output.scheduled_new_reqs = []
+        scheduler_output.scheduled_cached_reqs = cached_requests
+
+        logger.info("Warmup 1/2: Decoding...")
+        for _ in range(num_decode_tokens - 1):
+            self.execute_model(scheduler_output)
+
path: vllm_spyre/v1/worker/spyre_worker.py
user: login: yannicks1
body: Personally, I am for the use of helper functions wherever they help reducing code duplication. I am aware its just 10 lines here, but they could be eliminated by reusing the `_warmup_model_forward_pass` we introduced in the original implementation. 
created_at: 2025-03-20T10:14:35Z
author_association: COLLABORATOR
diff_hunk: @@ -253,95 +250,99 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
         # Convert to tensor for sampling
         valid_token_ids_tensor = torch.tensor(valid_token_ids,
                                               dtype=torch.long,
-                                              device=torch.device("cpu"))
+                                              device="cpu")
+
         # Sample from the valid token ids
         warmup_tokens_tensor = valid_token_ids_tensor[torch.randint(
             0, len(valid_token_ids_tensor), (batch_size, prompt_len))]
 
-        extra_kwargs = {}
-        if envs_spyre.VLLM_SPYRE_DYNAMO_BACKEND not in [
-                "sendnn", "sendnn_decoder"
-        ]:
-            # Bug in 2.3.1 fixed in 2.4.1 for SDPA flash cpu
-            # impl when padding too much
-            extra_kwargs["attn_algorithm"] = "math"
-
-        print(f"[SpyreWorker] warmup for prompt length "
-              f"{prompt_len} and max output tokens {num_decode_tokens}.")
-
-        # 1. trace
-        print("[SpyreWorker] warmup 1/2...")
-        # TODO: torch_sendnn.CleanGraph() should be necessary?
-        # warmup 1st forward pass
-        self._warmup_model_forward_pass(warmup_tokens_tensor,
-                                        valid_token_ids_tensor, prompt_len,
-                                        num_decode_tokens, batch_size,
-                                        extra_kwargs)
-
-        # 2. compile
-        print("[SpyreWorker] warmup 2/2...")
+        # Create requests to be used for prefill steps
+        dummy_requests = [
+            NewRequestData(
+                req_id="warmup",
+                prompt_token_ids=warmup_tokens_tensor[i].tolist(),
+                prompt="test",
+                mm_inputs=[],
+                mm_hashes=[],
+                mm_positions=[],
+                sampling_params=SamplingParams(max_tokens=num_decode_tokens),
+                block_ids=[0],
+                num_computed_tokens=0,
+                lora_request=None,
+            ) for i in range(batch_size)
+        ]
+
+        # Set up dummy cached_requests to be used for decode steps
+        cached_requests = [
+            CachedRequestData(
+                req_id=req.req_id,
+                resumed_from_preemption=False,
+                new_token_ids=[
+                    valid_token_ids_tensor[torch.randint(
+                        0, len(valid_token_ids_tensor), (1, )).item()]
+                ],  # placeholder token
+                new_block_ids=req.block_ids,
+                num_computed_tokens=req.num_computed_tokens,
+            ) for req in dummy_requests
+        ]
+
+        # To be used for execute_model, start with scheduled_new_reqs
+        # for prefill
+        scheduler_output = SchedulerOutput(
+            scheduled_new_reqs=dummy_requests,
+            scheduled_cached_reqs=[],
+            num_scheduled_tokens={i: prompt_len
+                                  for i in range(batch_size)},
+            total_num_scheduled_tokens=sum(prompt_len
+                                           for _ in range(batch_size)),
+            scheduled_spec_decode_tokens={},
+            scheduled_encoder_inputs={},
+            num_common_prefix_blocks=0,
+            finished_req_ids=set(),
+            free_encoder_input_ids=[],
+        )
+
+        # First full forward pass
+        logger.info("Warmup 1/2: Prefill...")
+        self.execute_model(scheduler_output)  # Prefill step
+
+        # Switch to cached requests to trigger decoding steps
+        scheduler_output.scheduled_new_reqs = []
+        scheduler_output.scheduled_cached_reqs = cached_requests
+
+        logger.info("Warmup 1/2: Decoding...")
+        for _ in range(num_decode_tokens - 1):
+            self.execute_model(scheduler_output)
+
+        # update_lazyhandle
         if envs_spyre.VLLM_SPYRE_DYNAMO_BACKEND == "sendnn_decoder":
             from torch_sendnn import torch_sendnn
             ul_start_time = time.time()
             torch_sendnn.update_lazyhandle()
             ul_stop_time = time.time()
-            ul_total_t = ul_stop_time - ul_start_time
-            print(f"update_lazyhandle() done (duration: {ul_total_t}s)")
+            logger.info("update_lazyhandle() done (duration: %.3fs)",
+                        ul_stop_time - ul_start_time)
+
+        # Second full forward pass
+        logger.info("Warmup 2/2: Prefill step...")
path: vllm_spyre/v1/worker/spyre_worker.py
user: login: yannicks1
body: please be consistent with the logging: 
```
logger.info("Warmup 1/2: Prefill...")
logger.info("Warmup 1/2: Decoding...")
```
vs. 
```
logger.info("Warmup 2/2: Prefill step...")
logger.info("Warmup 2/2: Decoding steps...")
```
created_at: 2025-03-20T10:16:35Z
author_association: COLLABORATOR
diff_hunk: @@ -253,95 +250,99 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
         # Convert to tensor for sampling
         valid_token_ids_tensor = torch.tensor(valid_token_ids,
                                               dtype=torch.long,
-                                              device=torch.device("cpu"))
+                                              device="cpu")
+
         # Sample from the valid token ids
         warmup_tokens_tensor = valid_token_ids_tensor[torch.randint(
             0, len(valid_token_ids_tensor), (batch_size, prompt_len))]
 
-        extra_kwargs = {}
-        if envs_spyre.VLLM_SPYRE_DYNAMO_BACKEND not in [
-                "sendnn", "sendnn_decoder"
-        ]:
-            # Bug in 2.3.1 fixed in 2.4.1 for SDPA flash cpu
-            # impl when padding too much
-            extra_kwargs["attn_algorithm"] = "math"
-
-        print(f"[SpyreWorker] warmup for prompt length "
-              f"{prompt_len} and max output tokens {num_decode_tokens}.")
-
-        # 1. trace
-        print("[SpyreWorker] warmup 1/2...")
-        # TODO: torch_sendnn.CleanGraph() should be necessary?
-        # warmup 1st forward pass
-        self._warmup_model_forward_pass(warmup_tokens_tensor,
-                                        valid_token_ids_tensor, prompt_len,
-                                        num_decode_tokens, batch_size,
-                                        extra_kwargs)
-
-        # 2. compile
-        print("[SpyreWorker] warmup 2/2...")
+        # Create requests to be used for prefill steps
+        dummy_requests = [
+            NewRequestData(
+                req_id="warmup",
+                prompt_token_ids=warmup_tokens_tensor[i].tolist(),
+                prompt="test",
+                mm_inputs=[],
+                mm_hashes=[],
+                mm_positions=[],
+                sampling_params=SamplingParams(max_tokens=num_decode_tokens),
+                block_ids=[0],
+                num_computed_tokens=0,
+                lora_request=None,
+            ) for i in range(batch_size)
+        ]
+
+        # Set up dummy cached_requests to be used for decode steps
+        cached_requests = [
+            CachedRequestData(
+                req_id=req.req_id,
+                resumed_from_preemption=False,
+                new_token_ids=[
+                    valid_token_ids_tensor[torch.randint(
+                        0, len(valid_token_ids_tensor), (1, )).item()]
+                ],  # placeholder token
+                new_block_ids=req.block_ids,
+                num_computed_tokens=req.num_computed_tokens,
+            ) for req in dummy_requests
+        ]
+
+        # To be used for execute_model, start with scheduled_new_reqs
+        # for prefill
+        scheduler_output = SchedulerOutput(
+            scheduled_new_reqs=dummy_requests,
+            scheduled_cached_reqs=[],
+            num_scheduled_tokens={i: prompt_len
+                                  for i in range(batch_size)},
+            total_num_scheduled_tokens=sum(prompt_len
+                                           for _ in range(batch_size)),
+            scheduled_spec_decode_tokens={},
+            scheduled_encoder_inputs={},
+            num_common_prefix_blocks=0,
+            finished_req_ids=set(),
+            free_encoder_input_ids=[],
+        )
+
+        # First full forward pass
+        logger.info("Warmup 1/2: Prefill...")
+        self.execute_model(scheduler_output)  # Prefill step
+
+        # Switch to cached requests to trigger decoding steps
+        scheduler_output.scheduled_new_reqs = []
+        scheduler_output.scheduled_cached_reqs = cached_requests
+
+        logger.info("Warmup 1/2: Decoding...")
+        for _ in range(num_decode_tokens - 1):
+            self.execute_model(scheduler_output)
+
path: vllm_spyre/v1/worker/spyre_worker.py
user: login: rafvasq
body: Sounds good, I re-introduced `_warmup_model_forward_pass` to handle the duplicate pass code
created_at: 2025-03-20T16:22:38Z
author_association: COLLABORATOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-03-13T15:26:59Z
body: 👋 Hi! Thank you for contributing to vLLM support on Spyre.
 Just a reminder: Make sure that your code passes all the linting checks, otherwise your PR won't be able to be merged. To do so, first install the linting requirements, then run `format.sh` and commit the changes:
```
pip install -r requirements-lint.txt
bash format.sh
```
Now you are good to go 🚀
author_association: NONE
user: login: joerunde
created_at: 2025-03-14T18:45:33Z
body: Can `_raw_model_forward` on the model runner be removed too?
author_association: COLLABORATOR
user: login: wallashss
created_at: 2025-03-17T15:59:33Z
body: Unfortunately, I could not yet tried your changes because of problems in my dev environment. But, besides the feedback from other reviewers, I feel that the final code should have more comments, specially the review comments pointed out by @yannicks1 for better understanding later. 
author_association: COLLABORATOR
user: login: joerunde
created_at: 2025-03-18T15:44:47Z
body: @rafvasq Can you also update all the prints in this code to logging statements?

See https://github.com/vllm-project/vllm-spyre/pull/31 for context. I didn't update the prints in this warmup code because I knew you were rewriting it
author_association: COLLABORATOR