refactor: update variable names and logging for better code clarity (#105)

* refactor: update variable name, type annotations and improve readability

* refactor: convert info logging to debug level

Author: pei0033

vllm_rbln/attention/backends/flash_attention.py

@@ -453,15 +453,15 @@ def build(
             attn_masks=attn_masks,
             kv_caches=None,
         )
-        logger.info("RBLNAttentionMetadata = %s", attn_metadata)
-        logger.info("\tslot_mapping size = %s", slot_mapping.size())
-        logger.info("\tblock_tables size = %s", block_tables.size())
+        logger.debug("RBLNAttentionMetadata = %s", attn_metadata)
+        logger.debug("\tslot_mapping size = %s", slot_mapping.size())
+        logger.debug("\tblock_tables size = %s", block_tables.size())
         if not envs.RBLN_FLASH_CAUSAL_ATTN and attn_masks is not None:
-            logger.info("\tattn_masks size = %s", attn_masks.size())
-            logger.info("\tattn_masks = %s", attn_masks[:, :, :, :, :32])
+            logger.debug("\tattn_masks size = %s", attn_masks.size())
+            logger.debug("\tattn_masks = %s", attn_masks[:, :, :, :, :32])
         else:
             assert attn_masks is None
-        logger.info("\tseq_lens_tensor size= %s", seq_lens_tensor.size())
+        logger.debug("\tseq_lens_tensor size= %s", seq_lens_tensor.size())
         return attn_metadata
 
 

vllm_rbln/v1/attention/backends/flash_attention.py

@@ -514,12 +514,12 @@ def build(
             kv_caches=None,
         )
 
-        logger.info("RBLNAttentionMetadata = %s", attn_metadata)
-        logger.info("\tslot_mapping size = %s", slot_mapping.size())
-        logger.info("\tblock_tables size = %s", block_tables_tensor.size())
-        logger.info("\tattn_masks size = %s", attn_masks.size())
-        logger.info("\tattn_masks = %s", attn_masks[:, :, :, :, :32])
-        logger.info("\tseq_lens_tensor size= %s", seq_lens_tensor.size())
+        logger.debug("RBLNAttentionMetadata = %s", attn_metadata)
+        logger.debug("\tslot_mapping size = %s", slot_mapping.size())
+        logger.debug("\tblock_tables size = %s", block_tables_tensor.size())
+        logger.debug("\tattn_masks size = %s", attn_masks.size())
+        logger.debug("\tattn_masks = %s", attn_masks[:, :, :, :, :32])
+        logger.debug("\tseq_lens_tensor size= %s", seq_lens_tensor.size())
         return attn_metadata
 
     def use_cascade_attention(self, *args, **kwargs) -> bool:

vllm_rbln/v1/worker/rbln_model_runner.py

@@ -653,19 +653,19 @@ def _prepare_inputs(
         # Hot-Swap lora model
         # if self.lora_config:
         #     self.set_active_loras(self.input_batch, num_scheduled_tokens)
-        logger.info("num_reqs: %s", num_reqs)
-        logger.info("token_indices: %s", token_indices)
-        logger.info("input_batch: %s", vars(self.input_batch))
-        logger.info(
+        logger.debug("num_reqs: %s", num_reqs)
+        logger.debug("token_indices: %s", token_indices)
+        logger.debug("input_batch: %s", vars(self.input_batch))
+        logger.debug(
             "input_ids: %s",
             self.input_ids[:scheduler_output.total_num_scheduled_tokens],
         )
-        logger.info(
+        logger.debug(
             "positions: %s",
             self.positions[:scheduler_output.total_num_scheduled_tokens],
         )
-        logger.info("attn_metadata: %s", next(iter(attn_metadata.items())))
-        logger.info("logits_indices: %s", logits_indices)
+        logger.debug("attn_metadata: %s", next(iter(attn_metadata.items())))
+        logger.debug("logits_indices: %s", logits_indices)
         return attn_metadata, logits_indices, spec_decode_metadata
 
     def _compile_model(self, model):

vllm_rbln/worker/model_runner.py

@@ -267,12 +267,12 @@ def _prepare_prompt(
                                                dtype=torch.long,
                                                device=self.device)
 
-        logger.info("[RBLN] model input builder, prepare_prompt")
-        logger.info("\tpadded input_tokens = %s", input_tokens)
-        logger.info("\tpadded input_positions = %s", input_positions)
-        logger.info("\tinput_block_ids = %s", input_block_ids)
-        logger.info("\tseq_lens = %s", data.seq_lens)
-        logger.info("\tquery_lens = %s", data.query_lens)
+        logger.debug("[RBLN] model input builder, prepare_prompt")
+        logger.debug("\tpadded input_tokens = %s", input_tokens)
+        logger.debug("\tpadded input_positions = %s", input_positions)
+        logger.debug("\tinput_block_ids = %s", input_block_ids)
+        logger.debug("\tseq_lens = %s", data.seq_lens)
+        logger.debug("\tquery_lens = %s", data.query_lens)
         return (input_tokens, input_positions, input_block_ids)
 
     def _prepare_decode(
@@ -340,12 +340,12 @@ def _prepare_decode(
                                                dtype=torch.long,
                                                device=self.device)
 
-        logger.info("[RBLN] model input builder, prepare_decode")
-        logger.info("\tpadded input_tokens = %s", data.input_tokens)
-        logger.info("\tpadded input_positions = %s", data.input_positions)
-        logger.info("\tinput_block_ids = %s", input_block_ids)
-        logger.info("\tseq_lens = %s", data.seq_lens)
-        logger.info("\tquery_lens = %s", data.query_lens)
+        logger.debug("[RBLN] model input builder, prepare_decode")
+        logger.debug("\tpadded input_tokens = %s", data.input_tokens)
+        logger.debug("\tpadded input_positions = %s", data.input_positions)
+        logger.debug("\tinput_block_ids = %s", input_block_ids)
+        logger.debug("\tseq_lens = %s", data.seq_lens)
+        logger.debug("\tquery_lens = %s", data.query_lens)
 
         assert input_tokens.shape[0] == self.max_num_seqs
         assert input_positions.shape[0] == self.max_num_seqs
@@ -520,10 +520,10 @@ def model_wrapper(
                     model_output = model_output[:, selected_token_indices]
                 logits = self.compute_logits_model.compute_logits(
                     model_output, None)
+                return logits
             else:
                 # non last rank create intermediate tensors, bypass it
-                logits = model_output
-            return logits
+                return model_output
 
         if self.model_config.enforce_eager or not envs.RBLN_COMPILE_MODEL:
             self.model_executable = model_wrapper
@@ -583,9 +583,9 @@ def prepare_model_input(
 
         is_prompt = seq_group_metadata_list[
             0].is_prompt if seq_group_metadata_list else None
-        logger.info("[RBLN] num_requests = %d", len(seq_group_metadata_list))
-        logger.info("[RBLN] input_ids = %s", model_input.input_tokens)
-        logger.info("[RBLN] positions = %s", model_input.input_positions)
+        logger.debug("[RBLN] num_requests = %d", len(seq_group_metadata_list))
+        logger.debug("[RBLN] input_ids = %s", model_input.input_tokens)
+        logger.debug("[RBLN] positions = %s", model_input.input_positions)
         return dataclasses.replace(model_input,
                                    sampling_metadata=sampling_metadata,
                                    virtual_engine=virtual_engine,
@@ -594,12 +594,12 @@ def prepare_model_input(
     @torch.inference_mode()
     def execute_model(
         self,
-        model_input: ModelInputForRebel,
+        model_input: ModelInputForRebelWithSamplingMetadata,
         kv_caches: Optional[List[torch.Tensor]] = None,
         intermediate_tensors: Optional[IntermediateTensors] = None,
         num_steps: int = 1,
         previous_hidden_states: Optional[torch.Tensor] = None,
-    ) -> Optional[SamplerOutput]:
+    ) -> Optional[Union[List[SamplerOutput], IntermediateTensors]]:
         assert kv_caches is not None
         if num_steps > 1:
             raise ValueError(
@@ -613,6 +613,7 @@ def execute_model(
         assert model_input.attn_metadata is not None
         token_indices = None
         if get_pp_group().is_last_rank:
+            assert model_input.sampling_metadata is not None
             num_prefills = model_input.attn_metadata.num_prefills
             selected_token_indices = \
                 model_input.sampling_metadata.selected_token_indices
@@ -633,30 +634,29 @@ def execute_model(
             if model_input.attn_metadata is not None:
                 model_input.attn_metadata.kv_caches = kv_caches
 
-            hidden_states = self.model_executable(
+            logits_or_intermediate_states = self.model_executable(
                 input_ids=model_input.input_tokens,
                 positions=model_input.input_positions,
                 intermediate_tensors=intermediate_tensors,
                 selected_token_indices=token_indices,
                 **execute_model_kwargs,
             )
 
-            if get_pp_group().is_last_rank:
-                # Gather logits for TP
-                logits_processor = self.compute_logits_model.logits_processor
-                hidden_states = logits_processor._gather_logits(hidden_states)
-                hidden_states = hidden_states.view(-1, hidden_states.size(-1))
+        if get_pp_group().is_last_rank:
+            # Gather logits for TP
+            logits_processor = self.compute_logits_model.logits_processor
+            logits = logits_processor._gather_logits(
+                logits_or_intermediate_states)
+            logits = logits.view(-1, logits.size(-1))
 
-        if not get_pp_group().is_last_rank:
-            intermediate_states = hidden_states
+        else:
+            intermediate_states = logits_or_intermediate_states
             assert isinstance(intermediate_states, IntermediateTensors)
             return intermediate_states
 
         # Compute the logits. -> moved to model executable
-        if num_prefills > 0 and len_token_indices != 0:
-            logits = hidden_states
-        else:
-            logits = hidden_states[selected_token_indices]
+        if not (num_prefills > 0 and len_token_indices != 0):
+            logits = logits[selected_token_indices]
 
         # Only perform sampling in the driver worker.
         if not self.is_driver_worker: