class_docstr: MicroBatchManager is a class that manages the micro batch.
  
  Args:
      stage (int): stage id of current stage.
      micro_batch_size (int): the micro batch size.
      micro_batch_buffer_size (int): the buffer size for micro batch. Normally, it should be the same as the number of pipeline stages.
class_docstr_tok: ["MicroBatchManager","is","a","class","that","manages","the","micro","batch",".","Args",":","stage","(","int",")",":","stage","id","of","current","stage",".","micro_batch_size","(","int",")",":","the","micro","batch","size",".","micro_batch_buffer_size","(","int",")",":","the","buffer","size","for","micro","batch",".","Normally",",","it","should","be","the","same","as","the","number","of","pipeline","stages","."]
docstr: Update the state if microbatch manager, 2 conditions.
  1. For first stage in PREFILL, receive inputs and outputs, `_add_descrption` will save its inputs.
  2. For other conditon, only receive the output of previous stage, and update the descrption.
  
  Args:
      inputs_dict (Dict[str, torch.Tensor]): the inputs of current stage. The key should have `input_ids` and `attention_mask`.
      output_dict (Dict[str, torch.Tensor]): the outputs of previous stage. The key should have `hidden_states` and `past_key_values`.
      new_token (torch.Tensor): the new token generated by current stage.
docstr_tok: ["Update","the","state","if","microbatch","manager",",","2","conditions",".","1",".","For","first","stage","in","PREFILL",",","receive","inputs","and","outputs",",","`","_add_descrption","`","will","save","its","inputs",".","2",".","For","other","conditon",",","only","receive","the","output","of","previous","stage",",","and","update","the","descrption",".","Args",":","inputs_dict","(","Dict","[","str",",","torch.Tensor","]",")",":","the","inputs","of","current","stage",".","The","key","should","have","`","input_ids","`","and","`","attention_mask","`",".","output_dict","(","Dict","[","str",",","torch.Tensor","]",")",":","the","outputs","of","previous","stage",".","The","key","should","have","`","hidden_states","`","and","`","past_key_values","`",".","new_token","(","torch.Tensor",")",":","the","new","token","generated","by","current","stage","."]
code: def step(self, new_token: torch.Tensor = None):
      """
      Update the state if microbatch manager, 2 conditions.
      1. For first stage in PREFILL, receive inputs and outputs, `_add_descrption` will save its inputs.
      2. For other conditon, only receive the output of previous stage, and update the descrption.
  
      Args:
          inputs_dict (Dict[str, torch.Tensor]): the inputs of current stage. The key should have `input_ids` and `attention_mask`.
          output_dict (Dict[str, torch.Tensor]): the outputs of previous stage. The key should have `hidden_states` and `past_key_values`.
          new_token (torch.Tensor): the new token generated by current stage.
      """
      # Add descrption first if the descrption is None
      self.cur_descrption.update(new_token)
      return self.cur_state
  
imports: from enum import Enum
  from typing import Dict
  import torch
  from ..tensor_parallel.batch_infer_state import BatchInferState
  from ..tensor_parallel.kvcache_manager import MemoryManager
label_desc: Use image node_id 3 for calling the MicroBatchManager obj's underlying member method code with example usage: obj.step(new_token) and returns: self