working except tests.

benchmarks/benchmark_latency.py

@@ -10,13 +10,6 @@
 
 from vllm import LLM, SamplingParams
 
-SAMPLE_PROMPTS = [
-    "The president of the United States is",
-    "Hello, my name is",
-    "The capital of France is",
-    "The future of AI is",
-]
-
 
 def main(args: argparse.Namespace):
     print(args)
@@ -35,7 +28,6 @@ def main(args: argparse.Namespace):
         device=args.device,
         block_size=args.block_size,
         max_chunked_prefill_len=args.max_chunked_prefill_len,
-        max_num_prompt_seqs=args.max_num_prompt_seqs,
         ray_workers_use_nsight=args.ray_workers_use_nsight,
     )
 
@@ -68,25 +60,16 @@ def run_to_completion(profile_dir: Optional[str] = None):
             print(p.key_averages())
         else:
             start_time = time.perf_counter()
-            if args.use_sample:
-                batch = (SAMPLE_PROMPTS *
-                         (args.batch_size // len(SAMPLE_PROMPTS) +
-                          1))[:args.batch_size]
-                outputs = llm.generate(prompts=batch,
-                                       sampling_params=sampling_params,
-                                       use_tqdm=False)
-            else:
-                outputs = llm.generate(prompt_token_ids=dummy_prompt_token_ids,
-                                       sampling_params=sampling_params,
-                                       use_tqdm=False)
+            outputs = llm.generate(prompt_token_ids=dummy_prompt_token_ids,
+                                   sampling_params=sampling_params,
+                                   use_tqdm=False)
             end_time = time.perf_counter()
             if args.verbose:
                 for output in outputs:
                     prompt = output.prompt
                     generated_text = output.outputs[0].text
-                    print(
-                        f"Prompt: {prompt!r}, Generated text: {generated_text!r}"
-                    )
+                    print(f"Prompt: {prompt!r}, Generated text: "
+                          f"{generated_text!r}")
             latency = end_time - start_time
             return latency
 
@@ -182,7 +165,6 @@ def run_to_completion(profile_dir: Optional[str] = None):
                         action='store_true',
                         help='print generated text')
     parser.add_argument('--max-chunked-prefill-len', type=int, default=-1)
-    parser.add_argument('--max-num-prompt-seqs', type=int, default=1000)
     parser.add_argument(
         "--ray-workers-use-nsight",
         action='store_true',

benchmarks/kernels/benchmark_paged_attention.py

@@ -7,7 +7,6 @@
 
 from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, create_kv_caches_with_random
 from vllm._C import ops
-from vllm.model_executor.layers.attention import flash_attn_with_kvcache_paged
 
 NUM_BLOCKS = 1024
 PARTITION_SIZE = 512

tests/conftest.py

@@ -168,7 +168,6 @@ def __init__(
         tensor_parallel_size: int = 1,
         block_size: int = 16,
         max_chunked_prefill_len: int = -1,
-        max_num_prompt_seqs: int = 1000,
         max_num_batched_tokens: int = 4096,
         **kwargs,
     ) -> None:
@@ -182,7 +181,6 @@ def __init__(
             tensor_parallel_size=tensor_parallel_size,
             block_size=block_size,
             max_chunked_prefill_len=max_chunked_prefill_len,
-            max_num_prompt_seqs=max_num_prompt_seqs,
             max_num_batched_tokens=max_num_batched_tokens,
             **kwargs,
         )

tests/core/test_scheduler.py

@@ -134,34 +134,28 @@ def test_scheduler_schedule_chunked_prefill():
     num_seq_group = 2
     max_model_len = 16
     max_chunked_prefill_len = 2
-    max_num_prompt_seqs = 1
     scheduler_config = SchedulerConfig(
         64,
         num_seq_group,
         max_model_len,
-        flash_style=True,
-        max_chunked_prefill_len=max_chunked_prefill_len,
-        max_num_prompt_seqs=max_num_prompt_seqs)
-    cache_config = CacheConfig(block_size, 1.0, 1)
+        max_chunked_prefill_len=max_chunked_prefill_len)
+    cache_config = CacheConfig(block_size, 1.0, 1, "auto")
     cache_config.num_cpu_blocks = 8
     cache_config.num_gpu_blocks = 8
     scheduler = Scheduler(scheduler_config, cache_config, None)
 
     # Add seq groups to scheduler.
     seq_groups: List[SequenceGroup] = []
     for i in range(num_seq_group):
-        _, seq_group = create_dummy_prompt(str(i),
-                                           prompt_length=block_size,
-                                           num_processed_token_ids=block_size -
-                                           1)
+        _, seq_group = create_dummy_prompt(str(i), prompt_length=block_size)
         scheduler.add_seq_group(seq_group)
         seq_groups.append(seq_group)
 
     # Schedule chunk prefill. Only the first seq_group should be scheduled.
     seq_group_meta, out = scheduler.schedule()
     assert set(out.scheduled_seq_groups) == set(seq_groups[:1])
-    seq_groups[0].get_num_unprefilled() == 2
-    seq_groups[1].get_num_unprefilled() == 4
+    assert seq_groups[0].get_num_unprefilled() == 2
+    assert seq_groups[1].get_num_unprefilled() == 4
     assert out.num_batched_tokens == 2
     assert (not out.blocks_to_copy and not out.blocks_to_swap_in
             and not out.blocks_to_swap_out)
@@ -170,11 +164,12 @@ def test_scheduler_schedule_chunked_prefill():
     assert seq_group_meta[0].is_chunked_prefill
     assert seq_group_meta[0].is_prompt
 
-    # Schedule chunk prefill. Still Only the first seq_group should be scheduled.
+    # Schedule chunk prefill. Still Only the first seq_group should be
+    # scheduled.
     seq_group_meta, out = scheduler.schedule()
     assert set(out.scheduled_seq_groups) == set(seq_groups[:1])
-    seq_groups[0].get_num_unprefilled() == 0
-    seq_groups[1].get_num_unprefilled() == 4
+    assert seq_groups[0].get_num_unprefilled() == 0
+    assert seq_groups[1].get_num_unprefilled() == 4
     assert out.num_batched_tokens == 2
     assert (not out.blocks_to_copy and not out.blocks_to_swap_in
             and not out.blocks_to_swap_out)
@@ -187,8 +182,8 @@ def test_scheduler_schedule_chunked_prefill():
     # for chunk prefill, and the first seq_group should be select for decoding.
     seq_group_meta, out = scheduler.schedule()
     assert set(out.scheduled_seq_groups) == set(seq_groups)
-    seq_groups[0].get_num_unprefilled() == 0
-    seq_groups[1].get_num_unprefilled() == 2
+    assert seq_groups[0].get_num_unprefilled() == 0
+    assert seq_groups[1].get_num_unprefilled() == 2
     assert out.num_batched_tokens == 3
     assert (not out.blocks_to_copy and not out.blocks_to_swap_in
             and not out.blocks_to_swap_out)

tests/samplers/test_sampler.py

@@ -285,49 +285,6 @@ def test_sampling(model_runner: ModelRunner):
     del model_runner
 
 
-@pytest.mark.parametrize("seed", RANDOM_SEEDS)
-@pytest.mark.parametrize("device", CUDA_DEVICES)
-def test_sampler_logits_processors(seed: int, device: str):
-    set_random_seed(seed)
-    torch.set_default_device(device)
-    batch_size = random.randint(1, 256)
-    input_tensor, _, sampler, model_runner = _prepare_test(batch_size)
-
-    # This sample logits processor gives infinite score to the i-th token,
-    # where i is the length of the input sequence.
-    # We therefore expect the output token sequence to be [0, 1, 2, ...]
-    def pick_ith(token_ids, logits):
-        logits[len(token_ids)] = float("inf")
-        return logits
-
-    seq_group_metadata_list = []
-    prompt_lens = []
-    for i in range(batch_size):
-        seq_group_metadata_list.append(
-            SequenceGroupMetadata(
-                request_id=f"test_{i}",
-                is_prompt=True,
-                is_chunked_prefill=False,
-                seq_data={0: SequenceData([1, 2, 3])},
-                sampling_params=SamplingParams(temperature=0,
-                                               logits_processors=[pick_ith]),
-                block_tables={0: [1]},
-            ))
-        prompt_lens.append(seq_group_metadata_list[-1].seq_data[0].get_len())
-
-    sampling_metadata = model_runner._prepare_sample(seq_group_metadata_list,
-                                                     prompt_lens,
-                                                     subquery_lens=prompt_lens)
-    sampler_output = sampler(embedding=None,
-                             hidden_states=input_tensor,
-                             sampling_metadata=sampling_metadata)
-    for _, sequence_output in enumerate(sampler_output):
-        for idx, nth_output in enumerate(sequence_output.samples):
-            assert nth_output.output_token == idx
-
-    del model_runner
-
-
 @pytest.mark.parametrize("seed", RANDOM_SEEDS)
 @pytest.mark.parametrize("device", CUDA_DEVICES)
 def test_sampler_top_k_top_p(seed: int, device: str):

vllm/config.py

@@ -539,8 +539,6 @@ class SchedulerConfig:
             requests. Longer requests will be chunked into multiple chunks.
             -1 means no chunking (disabled). This features is only supported
             for flash style attention.
-        max_num_prompt_seqs: The maximum number of prompt sequences that can be
-            processed in a single iteration.
     """
 
     def __init__(
@@ -549,7 +547,6 @@ def __init__(
         max_num_seqs: int,
         max_model_len: int,
         max_chunked_prefill_len: int = -1,
-        max_num_prompt_seqs: int = 1024,
     ) -> None:
         if max_num_batched_tokens is not None:
             self.max_num_batched_tokens = max_num_batched_tokens
@@ -561,7 +558,6 @@ def __init__(
         self.max_model_len = max_model_len
         self.chunked_prefill_enabled = max_chunked_prefill_len != -1
         self.max_chunked_prefill_len = max_chunked_prefill_len
-        self.max_num_prompt_seqs = max_num_prompt_seqs
         self._verify_args()
 
     def _verify_args(self) -> None:

vllm/core/scheduler.py

@@ -371,7 +371,6 @@ def schedule(self) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs]:
                 block_tables[seq_id] = self.block_manager.get_block_table(seq)
                 self.block_manager.access_all_blocks_in_seq(seq, now)
 
-            # SANG-TODO Update chunked prefill related info.
             seq_group_metadata = SequenceGroupMetadata(
                 request_id=seq_group.request_id,
                 is_prompt=scheduler_outputs.prompt_run,

vllm/engine/arg_utils.py

@@ -52,7 +52,6 @@ class EngineArgs:
     device: str = 'auto'
     ray_workers_use_nsight: bool = False
     max_chunked_prefill_len: int = -1
-    max_num_prompt_seqs: int = 256
 
     def __post_init__(self):
         if self.tokenizer is None:
@@ -357,8 +356,7 @@ def create_engine_configs(
             self.max_num_batched_tokens,
             self.max_num_seqs,
             model_config.max_model_len,
-            max_chunked_prefill_len=self.max_chunked_prefill_len,
-            max_num_prompt_seqs=self.max_num_prompt_seqs)
+            max_chunked_prefill_len=self.max_chunked_prefill_len)
         lora_config = LoRAConfig(
             max_lora_rank=self.max_lora_rank,
             max_loras=self.max_loras,

vllm/engine/llm_engine.py

@@ -615,8 +615,6 @@ def step(self) -> List[RequestOutput]:
             >>>         break
         """
         seq_group_metadata_list, scheduler_outputs = self.scheduler.schedule()
-        # print("SANG-TODO step seq_group_metadata_list length: ",
-        #   len(seq_group_metadata_list))
         if not scheduler_outputs.is_empty():
             output = self.model_executor.execute_model(
                 seq_group_metadata_list, scheduler_outputs.blocks_to_swap_in,

vllm/entrypoints/llm.py

@@ -146,7 +146,6 @@ def generate(
             A list of `RequestOutput` objects containing the generated
             completions in the same order as the input prompts.
         """
-        # print("SANG-TODO generate: ", prompts, prompt_token_ids)
         if prompts is None and prompt_token_ids is None:
             raise ValueError("Either prompts or prompt_token_ids must be "
                              "provided.")

vllm/utils.py

@@ -322,7 +322,6 @@ def create_kv_caches_with_random(
         key_caches.append(key_cache)
 
     value_cache_shape = (num_blocks, num_heads, head_size, block_size)
-
     value_caches = []
     for _ in range(num_layers):
         value_cache = torch.empty(size=value_cache_shape,