class_docstr: This class is used to record the infomation of the first stage of pipeline, the first stage should have attributes `input_ids` and `attention_mask`
  and `new_tokens`, and the `new_tokens` is the tokens generated by the first stage. Also due to the schdule of pipeline, the operation to update the
  information and the condition to determine the state is different from other stages.
  
  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`.
class_docstr_tok: ["This","class","is","used","to","record","the","infomation","of","the","first","stage","of","pipeline",",","the","first","stage","should","have","attributes","`","input_ids","`","and","`","attention_mask","`","and","`","new_tokens","`",",","and","the","`","new_tokens","`","is","the","tokens","generated","by","the","first","stage",".","Also","due","to","the","schdule","of","pipeline",",","the","operation","to","update","the","information","and","the","condition","to","determine","the","state","is","different","from","other","stages",".","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","`","."]
code: def _update_attnmask(self):
      self.attn_mask = torch.cat(
          (
              self.attn_mask,
              torch.ones(
                  (self.attn_mask.shape[0], 1),
                  dtype=torch.int64,
                  device="cuda",
              ),
          ),
          dim=-1,
      )
  
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 4 for calling the HeadMicroBatchDescription obj's underlying member method code with example usage: obj._update_attnmask() without return types