Fix per comments

Signed-off-by: Chendi.Xue <[email protected]>

Author: xuechendi

tests/full_tests/spec_decode.py

@@ -20,13 +20,14 @@
 def time_generation(llm: LLM,
                     prompts: list[str],
                     sampling_params: SamplingParams,
-                    num_spec_tokens=5):
+                    num_spec_tokens=5,
+                    num_warmups=1):
     # Generate texts from the prompts. The output is a list of RequestOutput
     # objects that contain the prompt, generated text, and other information.
     # Warmup first
     logging.info("Warming up the model...")
-    llm.generate(prompts, sampling_params)
-    llm.generate(prompts, sampling_params)
+    for _ in range(num_warmups):
+        llm.generate(prompts, sampling_params)
     logging.info("Starting generation...")
     start = time.time()
     outputs = llm.generate(prompts, sampling_params)
@@ -103,7 +104,7 @@ def test_ngram(is_enable, args, prompts, sampling_params, task_key,
         )
 
     result_dict = time_generation(llm, prompts, sampling_params,
-                                  args.num_spec_tokens)
+                                  args.num_spec_tokens, args.num_warmups)
 
     result_queue.put((task_key, result_dict))
 
@@ -128,7 +129,7 @@ def test_eagle_model(is_enable, args, prompts, sampling_params, task_key,
         )
 
     result_dict = time_generation(llm, prompts, sampling_params,
-                                  args.num_spec_tokens)
+                                  args.num_spec_tokens, args.num_warmups)
     result_queue.put((task_key, result_dict))
 
 
@@ -152,7 +153,7 @@ def test_medusa_model(is_enable, args, prompts, sampling_params, task_key,
         )
 
     result_dict = time_generation(llm, prompts, sampling_params,
-                                  args.num_spec_tokens)
+                                  args.num_spec_tokens, args.num_warmups)
     result_queue.put((task_key, result_dict))
 
 
@@ -175,7 +176,7 @@ def test_mtp_model(is_enable, args, prompts, sampling_params, task_key,
         )
 
     result_dict = time_generation(llm, prompts, sampling_params,
-                                  args.num_spec_tokens)
+                                  args.num_spec_tokens, args.num_warmups)
     result_queue.put((task_key, result_dict))
 
 
@@ -199,6 +200,10 @@ def test_mtp_model(is_enable, args, prompts, sampling_params, task_key,
     parser.add_argument("--enforce_eager",
                         action="store_true",
                         help="Enforce eager execution for Eagle model.")
+    parser.add_argument("--num_warmups",
+                        type=int,
+                        default=1,
+                        help="Number of warmup runs before timing.")
 
     # 'ngram', 'eagle', 'eagle3', 'medusa', 'mlp_speculator',
     # 'draft_model' or 'deepseek_mtp

vllm_gaudi/v1/sample/hpu_rejection_sampler.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+
 from vllm.v1.sample import rejection_sampler
 import torch
 from typing import Optional
@@ -68,7 +70,9 @@ def rejection_greedy_sample_pytorch(
         # This loop is a direct translation of the Triton kernel's core logic.
         rejected = False
         for pos in range(num_draft_tokens):
-            if not rejected:
+            if rejected:
+                break
+            else:
                 draft_token = draft_token_ids[start_idx + pos]
                 target_token = target_argmax[start_idx + pos]
 
@@ -79,11 +83,6 @@ def rejection_greedy_sample_pytorch(
                 # all subsequent tokens.
                 if draft_token != target_token:
                     rejected = True
-            else:
-                # This `break` is a Pythonic optimization. The original Triton
-                # kernel continues the loop but the `if not rejected` check
-                # prevents further operations. Breaking is more efficient here.
-                break
 
         # If the entire draft sequence was accepted without any rejection,
         # append the bonus token.
@@ -148,8 +147,7 @@ def rejection_sample(
         bonus_token_ids,
         is_greedy,
     )
-    if sampling_metadata.all_greedy:
-        return output_token_ids
+    return output_token_ids
 
 
 rejection_sampler.rejection_sample = rejection_sample

vllm_gaudi/v1/worker/hpu_model_runner.py

@@ -1629,14 +1629,18 @@ def _prepare_decode_inputs(self,
         # self.input_batch.num_computed_tokens_cpu[req_indices]
         positions = torch.zeros((padded_batch_size, num_tokens),
                                 dtype=torch.int32)
-        # per request using universal self.positions_cpu then pad
-        position_split_tensors = torch.split(
-            self.positions_cpu[:total_num_scheduled_tokens],
-            num_tokens_per_req)
-        positions[:num_decodes] = \
-            pad_sequence(list(position_split_tensors),
-                            batch_first=True,
-                            padding_value=0)[:num_decodes]
+        if num_tokens == 1:
+            positions[:num_decodes] = self.positions_cpu[:num_decodes].reshape(
+                -1, 1)
+        else:
+            # per request using universal self.positions_cpu then pad
+            position_split_tensors = torch.split(
+                self.positions_cpu[:total_num_scheduled_tokens],
+                num_tokens_per_req)
+            positions[:num_decodes] = \
+                pad_sequence(list(position_split_tensors),
+                                batch_first=True,
+                                padding_value=0)[:num_decodes]
 
         padded_index = torch.zeros((padded_batch_size, num_tokens),
                                    dtype=torch.int64)
@@ -1680,13 +1684,17 @@ def _prepare_decode_inputs(self,
         # self.input_batch.token_ids_cpu[:total_num_scheduled_tokens]
         token_ids = torch.zeros((padded_batch_size, num_tokens),
                                 dtype=torch.int32)
-        token_ids_split_tensors = torch.split(
-            self.input_ids_cpu[:total_num_scheduled_tokens],
-            num_tokens_per_req)
-        token_ids[:num_decodes] = \
-            pad_sequence(list(token_ids_split_tensors),
-                            batch_first=True,
-                            padding_value=0)[:num_decodes]
+        if num_tokens == 1:
+            token_ids[:num_decodes] = self.input_ids_cpu[:num_decodes].reshape(
+                -1, 1)
+        else:
+            token_ids_split_tensors = torch.split(
+                self.input_ids_cpu[:total_num_scheduled_tokens],
+                num_tokens_per_req)
+            token_ids[:num_decodes] = \
+                pad_sequence(list(token_ids_split_tensors),
+                                batch_first=True,
+                                padding_value=0)[:num_decodes]
 
         ###################################
         # SLOT_MAPPING [batch, 1]