[2/N] Chunked prefill data update

vllm/entrypoints/llm.py

@@ -152,7 +152,7 @@ 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)
+        # 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/model_executor/layers/attention.py

@@ -141,7 +141,7 @@ def forward(
             shape = [batch_size, seq_len, num_heads * head_size]
         """
         batch_size, seq_len, hidden_size = query.shape
-        print("SANG-TODO query size: ", query.size())
+        # print("SANG-TODO query size: ", query.size())
         # Reshape the query, key, and value tensors.
         query = query.view(-1, self.num_heads, self.head_size)
         key = key.view(-1, self.num_kv_heads, self.head_size)
@@ -153,12 +153,12 @@ def forward(
         # profiling run.
         if key_cache is not None and value_cache is not None:
             if input_metadata.flash_style:
-                print("SANG-TODO reshape cache flash.")
+                # print("SANG-TODO reshape cache flash.")
                 cache_ops.reshape_and_cache_flash(
                     key, value, key_cache, value_cache,
                     input_metadata.slot_mapping.flatten())
             else:
-                print("SANG-TODO reshape cache.")
+                # print("SANG-TODO reshape cache.")
                 cache_ops.reshape_and_cache(
                     key,
                     value,
@@ -173,33 +173,11 @@ def forward(
             if (key_cache is None or value_cache is None
                     # or input_metadata.block_tables.numel() == 0):
                     or not input_metadata.prefix_enabled):
-                print("SANG-TODO flash attn is used.")
-                print(
-                    "SANG-TODO query size: ",
-                    query.view(batch_size, seq_len, self.num_heads,
-                               self.head_size).size())
-                # if key_cache is not None and value_cache is not None:
-                #     output2 = flash_attn_with_kvcache_paged(
-                #         query.view(batch_size, seq_len, self.num_heads,
-                #                     self.head_size),
-                #         key_cache,
-                #         value_cache,
-                #         self.scale,
-                #         input_metadata.block_tables,
-                #         input_metadata.context_lens + seq_len,
-                #         self.alibi_slopes,
-                #     )
-                #     from flash_attn import flash_attn_func
-                #     breakpoint()
-                #     output3 = flash_attn_func(
-                #         q=query.view(batch_size, seq_len, self.num_heads,
-                #                     self.head_size),
-                #         k=key.view(batch_size, seq_len, self.num_kv_heads, self.head_size),
-                #         v=value.view(batch_size, seq_len, self.num_kv_heads, self.head_size),
-                #         softmax_scale=self.scale,
-                #         causal=True,
-                #         alibi_slopes=self.alibi_slopes,
-                #     )
+                # print("SANG-TODO flash attn is used.")
+                # print(
+                #     "SANG-TODO query size: ",
+                #     query.view(batch_size, seq_len, self.num_heads,
+                #                self.head_size).size())
                 if self.num_kv_heads != self.num_heads:
                     # As of Nov 2023, xformers only supports MHA. For MQA/GQA,
                     # project the key and value tensors to the desired number of

vllm/worker/model_runner.py

@@ -138,8 +138,8 @@ def _prepare_prompt(
         context_lens: List[int] = []
         subquery_lens: List[int] = []
         prefix_block_tables: List[List[int]] = []
-        print("SANG-TODO # of requests (seq_group_metadata_list): ",
-              len(seq_group_metadata_list))
+        # print("SANG-TODO # of requests (seq_group_metadata_list): ",
+        #       len(seq_group_metadata_list))
         for seq_group_metadata in seq_group_metadata_list:
             assert seq_group_metadata.is_prompt
             seq_ids = list(seq_group_metadata.seq_data.keys())
@@ -152,7 +152,7 @@ def _prepare_prompt(
             prompt_lens.append(prompt_len)
             prefix_len = 0
             prefix = seq_group_metadata.prefix
-            print("SANG-TODO prefix, ", prefix)
+            # print("SANG-TODO prefix, ", prefix)
             if prefix is not None and prefix.computed:
                 prefix_len = prefix.get_length()
                 prompt_tokens = prompt_tokens[prefix_len:]
@@ -500,12 +500,12 @@ def prepare_input_tensors(
             # SANG-TODO set num prompt tokens and generations?
             # Prepare input tensors.
             if is_prompt:
-                print("SANG-TODO execute model prompt.")
+                # print("SANG-TODO execute model prompt.")
                 (input_tokens, input_positions, input_metadata, prompt_lens,
                  subquery_lens, lora_index_mapping, lora_prompt_mapping,
                  lora_requests) = self._prepare_prompt(seq_group_metadata_list)
             else:
-                print("SANG-TODO execute model decode.")
+                # print("SANG-TODO execute model decode.")
                 (input_tokens, input_positions, input_metadata,
                  lora_index_mapping, lora_prompt_mapping,
                  lora_requests) = self._prepare_decode(seq_group_metadata_list)

vllm/worker/worker.py

@@ -117,7 +117,7 @@ def profile_num_available_blocks(
             gpu_memory_utilization: The fraction of the total GPU memory to use.
             cpu_swap_space: The size of the CPU swap space in bytes.
         """
-        print("SANG-TODO profile_num_available_blocks")
+        # print("SANG-TODO profile_num_available_blocks")
         # Profile the memory usage of the model and get the maximum number of
         # cache blocks that can be allocated with the remaining free memory.
         torch.cuda.empty_cache()
@@ -154,7 +154,7 @@ def profile_num_available_blocks(
                                  MAX_INT_32 // cache_block_size)
             num_cpu_blocks = min(num_cpu_blocks,
                                  MAX_INT_32 // cache_block_size)
-        print("SANG-TODO profile_num_available_blocks done")
+        # print("SANG-TODO profile_num_available_blocks done")
 
         return num_gpu_blocks, num_cpu_blocks
 
@@ -207,7 +207,7 @@ def execute_model(
         blocks_to_swap_out: Optional[Dict[int, int]] = None,
         blocks_to_copy: Optional[Dict[int, List[int]]] = None,
     ) -> Optional[SamplerOutput]:
-        print("SANG-TODO execute model.")
+        # print("SANG-TODO execute model.")
         if self.is_driver_worker:
             assert seq_group_metadata_list is not None
             num_seq_groups = len(seq_group_metadata_list)