docstr: Extend the KV-cache for the sequence.
  
  Arguments:
      new_ids (Union[List[torch.IntTensor], torch.IntTensor]): For each allocation group, the IDs
          to add to the KV-cache. If there is only one allocation group, a single tensor can be
          provided. Otherwise, a list of tensors should be provided. The tensors do not need
          to have the same shape.
docstr_tok: ["Extend","the","KV-cache","for","the","sequence",".","Arguments",":","new_ids","(","Union","[","List","[","torch.IntTensor","]",",","torch.IntTensor","]",")",":","For","each","allocation","group",",","the","IDs","to","add","to","the","KV-cache",".","If","there","is","only","one","allocation","group",",","a","single","tensor","can","be","provided",".","Otherwise",",","a","list","of","tensors","should","be","provided",".","The","tensors","do","not","need","to","have","the","same","shape","."]
code: def extend_kv_cache(
      self,
      new_ids: Union[List[torch.IntTensor], torch.IntTensor],
      cache_group: int = 0,
  ) -> None:
      """
      Extend the KV-cache for the sequence.
  
      Arguments:
          new_ids (Union[List[torch.IntTensor], torch.IntTensor]): For each allocation group, the IDs
              to add to the KV-cache. If there is only one allocation group, a single tensor can be
              provided. Otherwise, a list of tensors should be provided. The tensors do not need
              to have the same shape.
      """
      if isinstance(new_ids, torch.Tensor):
          new_ids = [new_ids]
  
      if len(new_ids) != self._num_allocation_groups[cache_group]:
          raise ValueError(
              f"Only {len(new_ids)} allocation groups provided, expected {self._num_allocation_groups[cache_group]}"
          )
  
      for group_id, new_group_ids in enumerate(new_ids):
          new_blocks = new_group_ids.numel()
  
          if new_blocks == 0:
              # If we have multiple groups, it's possible to have an empty group.
              continue
  
          shadow_alloc_group = self._kv_cache_ids_shadow[cache_group][
              group_id
          ]
          alloc_group = self._kv_cache_ids[cache_group][group_id]
          cur_blocks = self._blocks_per_allocation_group[cache_group][
              group_id
          ]
  
          shadow_alloc_group[
              cur_blocks : cur_blocks + new_blocks
          ].copy_(new_group_ids)
          alloc_group[cur_blocks : cur_blocks + new_blocks].copy_(
              shadow_alloc_group[cur_blocks : cur_blocks + new_blocks],
              non_blocking=True,
          )
  
          self._blocks_per_allocation_group[cache_group][
              group_id
          ] += new_blocks
  
imports: from typing import List, Tuple, Union
  import torch
label_desc: Use image node_id 12 for calling the DSSequenceDescriptor obj's underlying member method code with example usage: obj.extend_kv_cache(new_ids, cache_group) without return types