code: def concatenate(
      cls, batches: List["FlashCausalLMBatch"]
  ) -> "FlashCausalLMBatch":
      # Batch attributes
      requests = []
      requests_idx_mapping = {}
  
      blocks = 0
      total_batch_size = 0
      total_slots = 0
      max_blocks = 0
      max_length = 0
      max_seqlen = 0
      for b in batches:
          total_batch_size += len(b)
          total_slots += len(b.slots)
          blocks += b.blocks
          speculative_length = (
              b.speculative_ids.shape[1]
              if b.speculative_ids is not None
              else 0
          )
          max_blocks = max(max_blocks, b.max_blocks)
          max_seqlen = max(max_seqlen, b.max_seqlen)
          max_length = max(
              max_length,
              max(
                  input_length
                  + stopping_criteria.max_new_tokens
                  + speculative_length
                  - stopping_criteria.current_tokens
                  for input_length, stopping_criteria in zip(
                      b.input_lengths, b.stopping_criterias
                  )
              ),
          )
  
      input_ids = batches[0].input_ids.new_empty(total_batch_size)
      position_ids = batches[0].position_ids.new_empty(total_batch_size)
      slots = batches[0].slots.new_empty(total_slots)
      slot_indices = batches[0].slot_indices.new_empty(total_batch_size)
      input_lengths_tensor = batches[0].input_lengths_tensor.new_empty(
          total_batch_size
      )
      block_tables_tensor = batches[0].block_tables_tensor.new_zeros(
          (total_batch_size, max_blocks)
      )
      all_input_ids_tensor = batches[0].all_input_ids_tensor.new_zeros(
          (total_batch_size, max_length)
      )
      top_n_tokens_tensor = batches[0].top_n_tokens_tensor.new_zeros(
          total_batch_size,
      )
  
      start_slots = []
      block_tables = []
      all_input_ids = []
  
      input_lengths = []
      prefix_offsets = []
      read_offsets = []
  
      next_token_chooser_parameters = []
      stopping_criterias = []
      top_n_tokens = []
  
      # Cumulative length
      cumulative_batch_size = 0
      cumulative_slots = 0
  
      for i, batch in enumerate(batches):
          requests.extend(batch.requests)
  
          if i == 0:
              requests_idx_mapping = batch.requests_idx_mapping
          else:
              # We need to offset the mapping for each batch by the cumulative batch size
              for k, v in batch.requests_idx_mapping.items():
                  requests_idx_mapping[k] = v + cumulative_batch_size
  
          start_index = cumulative_batch_size
          end_index = cumulative_batch_size + len(batch)
          slots_start_index = cumulative_slots
          slots_end_index = cumulative_slots + len(batch.slots)
  
          # Copy tensors (GPU)
          input_ids[start_index:end_index] = batch.input_ids
          position_ids[start_index:end_index] = batch.position_ids
          slot_indices[start_index:end_index] = (
              batch.slot_indices + cumulative_slots
          )
          input_lengths_tensor[
              start_index:end_index
          ] = batch.input_lengths_tensor
          top_n_tokens_tensor[
              start_index:end_index
          ] = batch.top_n_tokens_tensor
          slots[slots_start_index:slots_end_index] = batch.slots
  
          all_input_ids_tensor[
              start_index:end_index,
              : batch.all_input_ids_tensor.shape[1],
          ] = batch.all_input_ids_tensor[:, :max_length]
  
          block_tables_tensor[
              start_index:end_index,
              : batch.block_tables_tensor.shape[1],
          ] = batch.block_tables_tensor[:, :max_blocks]
  
          start_slots.append(batch.start_slots + cumulative_slots)
  
          block_tables.extend(batch.block_tables)
          all_input_ids.extend(batch.all_input_ids)
  
          input_lengths.extend(batch.input_lengths)
          prefix_offsets.extend(batch.prefix_offsets)
          read_offsets.extend(batch.read_offsets)
  
          next_token_chooser_parameters.extend(
              [r.parameters for r in batch.requests]
          )
          stopping_criterias.extend(batch.stopping_criterias)
  
          top_n_tokens.extend(batch.top_n_tokens)
  
          # Update
          cumulative_batch_size += len(batch)
          cumulative_slots += len(batch.slots)
  
      start_slots = torch.concat(start_slots)
  
      next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
          next_token_chooser_parameters,
          dtype=batches[0].next_token_chooser.dtype,
          device=batches[0].next_token_chooser.device,
      )
  
      speculative_ids = (
          torch.cat([b.speculative_ids for b in batches], dim=0)
          if batches[0].speculative_ids is not None
          else None
      )
  
      # Needed to avoid dropping blocks when the batches will go out of scope
      for b in batches:
          b.block_tables = None
          del b
  
      return cls(
          batch_id=batches[0].batch_id,
          requests=requests,
          requests_idx_mapping=requests_idx_mapping,
          input_ids=input_ids,
          position_ids=position_ids,
          cu_seqlen_prefill=None,
          start_slots=start_slots,
          slot_indices=slot_indices,
          needed_blocks_slots=None,
          block_tables=block_tables,
          block_tables_tensor=block_tables_tensor,
          slots=slots,
          max_seqlen=max_seqlen,
          prefill_head_indices=None,
          prefill_next_token_indices=None,
          prefill_cu_outlens=None,
          input_lengths=input_lengths,
          input_lengths_tensor=input_lengths_tensor,
          prefix_offsets=prefix_offsets,
          read_offsets=read_offsets,
          all_input_ids=all_input_ids,
          all_input_ids_tensor=all_input_ids_tensor,
          next_token_chooser=next_token_chooser,
          stopping_criterias=stopping_criterias,
          top_n_tokens=top_n_tokens,
          top_n_tokens_tensor=top_n_tokens_tensor,
          blocks=blocks,
          max_blocks=max_blocks,
          speculative_ids=speculative_ids,
      )
  
imports: import math
  import time
  import itertools
  import torch
  import torch.distributed
  import numpy
  from dataclasses import dataclass
  from opentelemetry import trace
  from transformers import PreTrainedTokenizerBase
  from typing import Optional, Tuple, List, Type, Dict
  from text_generation_server.models import Model
  from text_generation_server.utils.tokens import batch_top_tokens
  from text_generation_server.utils.speculate import get_speculate
  from text_generation_server.models.types import Batch, Tokens, Generation, GeneratedText
  from text_generation_server.models.cache_manager import get_cache_manager, set_cache_manager, BLOCK_SIZE
  from text_generation_server.pb import generate_pb2
  from text_generation_server.utils import StoppingCriteria, HeterogeneousNextTokenChooser
  from text_generation_server.utils.dist import MEMORY_FRACTION
label_desc: Use image node_id 4 for calling the FlashCausalLMBatch obj's underlying member method code with example usage: obj.concatenate(cls, batches) and returns: cls