class_docstr: Engine class for tensor parallel inference.
  
  Args:
      model (Module): original model, e.g. huggingface CausalLM
      shard_config (ShardConfig): The config for sharding original model
      max_batch_size (int): maximum batch size
      max_input_len (int): maximum input length of sequence
      max_output_len (int): maximum output length of output tokens
      dtype (torch.dtype): datatype used to init KV cache space
      device (str): device the KV cache of engine to be initialized on
  
  Examples:
      >>> # define model and shard config for your inference
      >>> model = ...
      >>> generate_kwargs = ...
      >>> shard_config = ShardConfig(enable_tensor_parallelism=True, extra_kwargs={"inference_only": True})
      >>> infer_engine = TPInferEngine(model, shard_config, MAX_BATCH_SIZE, MAX_INPUT_LEN, MAX_OUTPUT_LEN)
      >>> outputs = infer_engine.generate(input_ids, **generate_kwargs)
class_docstr_tok: ["Engine","class","for","tensor","parallel","inference",".","Args",":","model","(","Module",")",":","original","model",",","e.g",".","huggingface","CausalLM","shard_config","(","ShardConfig",")",":","The","config","for","sharding","original","model","max_batch_size","(","int",")",":","maximum","batch","size","max_input_len","(","int",")",":","maximum","input","length","of","sequence","max_output_len","(","int",")",":","maximum","output","length","of","output","tokens","dtype","(","torch.dtype",")",":","datatype","used","to","init","KV","cache","space","device","(","str",")",":","device","the","KV","cache","of","engine","to","be","initialized","on","Examples",":",">",">",">","#","define","model","and","shard","config","for","your","inference",">",">",">","model","=","...",">",">",">","generate_kwargs","=","...",">",">",">","shard_config","=","ShardConfig","(","enable_tensor_parallelism=True",",","extra_kwargs=","{","``","inference_only","''",":","True","}",")",">",">",">","infer_engine","=","TPInferEngine","(","model",",","shard_config",",","MAX_BATCH_SIZE",",","MAX_INPUT_LEN",",","MAX_OUTPUT_LEN",")",">",">",">","outputs","=","infer_engine.generate","(","input_ids",",","*","*","generate_kwargs",")"]
docstr: Forward is used in Dynamic Batching Manager
docstr_tok: ["Forward","is","used","in","Dynamic","Batching","Manager"]
code: def forward(self, batch_id, is_prefill):
      """
      Forward is used in Dynamic Batching Manager
      """
      batch = self.cache.pop(batch_id)
      if is_prefill:
          input_ = torch.tensor(batch.all_input_ids).cuda()
      else:
          input_ = batch.input_ids.reshape(len(batch), 1)
  
      batch_args = {
          "batch_size": len(batch),
          "max_len_in_batch": batch.nopad_max_len_in_batch,
          "block_loc": batch.nopad_b_loc,
          "start_loc": batch.nopad_b_start_loc,
          "seq_len": batch.nopad_b_seq_len,
          "cache_manager": batch.cache_manager,
          "is_context_stage": is_prefill,
      }
  
      infer_state = BatchInferState(**batch_args)
      model = self.model
      if isinstance(model, LlamaForCausalLM):
          model = self.model.model
      elif isinstance(model, BloomForCausalLM):
          model = self.model.transformer
  
      setattr(model, "infer_state", infer_state)
      output = self.model.forward(input_ids=input_)
      logits = output.logits
      # bsz, seq_len, vocab_size
      prob_out = torch.softmax(
          logits[
              :,
              -1,
          ],
          dim=-1,
      ).squeeze(1)
      # prob_out: bsz, vocab_size
      predict_ids = torch.argmax(prob_out, dim=-1, keepdim=True)
      prob_out = torch.log(prob_out).detach().cpu().numpy()
      predict_ids = predict_ids.detach().cpu().numpy()
      # [ batch_size, 1 ]
  
      output_dict = {}
      new_input_ids = []
      for i, (
          r,
          all_input_ids,
          next_token_id,
          next_token_logprob,
      ) in enumerate(
          zip(
              batch.requests, batch.all_input_ids, predict_ids, prob_out
          )
      ):
          next_token_id = int(next_token_id)
          next_token_logprob = next_token_logprob[next_token_id]
          # all_input_ids_tensor = torch.tensor(all_input_ids, dtype=torch.long, device="cuda")
          all_input_ids.append(next_token_id)
          # all_input_ids_tensor = None
          new_input_ids.append(next_token_id)
          batch.all_input_ids[i] = all_input_ids
          batch.input_lengths[i] += 1
          batch.out_token_id_counts[i][next_token_id] += 1
          metadata = {
              "id": int(next_token_id),
              "logprob": float(next_token_logprob),
          }
          output_dict[r["request_id"]] = (int(next_token_id), metadata)
  
      batch.input_ids = torch.tensor(
          new_input_ids, dtype=torch.long
      ).cuda()
      batch.nopad_total_token_num += len(batch)
      batch.nopad_max_len_in_batch += 1  # NOTE: we may repalce this
      self.cache[batch.batch_id] = batch
      return output_dict
  
imports: from typing import Any, Callable, List, Optional, Union
  import torch
  import torch.nn
  from transformers import BloomForCausalLM, LlamaForCausalLM
  from transformers.generation import GenerationConfig
  from transformers.generation.stopping_criteria import StoppingCriteriaList
  from transformers.tokenization_utils_base import BatchEncoding
  from colossalai.shardformer import ShardConfig, ShardFormer
  from colossalai.shardformer.policies.auto_policy import get_autopolicy
  from .batch_infer_state import BatchInferState
  from .kvcache_manager import MemoryManager
label_desc: Use image node_id 13 for calling the TPInferEngine obj's underlying member method code with example usage: obj.forward(batch_id, is_prefill) and returns: output_dict