class_docstr: Implementation of a model for inference composable with ragged batching.
class_docstr_tok: ["Implementation","of","a","model","for","inference","composable","with","ragged","batching","."]
docstr: This will be called before each forward with the intent of building forward-specific metadata
  about a batch. The intent here is to build data structures like attention atoms without necessarily
  needing to implement graphable kernels to do so.
  
  Abstract so as to force model implementations to opt out of doing anything here explicitly.
docstr_tok: ["This","will","be","called","before","each","forward","with","the","intent","of","building","forward-specific","metadata","about","a","batch",".","The","intent","here","is","to","build","data","structures","like","attention","atoms","without","necessarily","needing","to","implement","graphable","kernels","to","do","so",".","Abstract","so","as","to","force","model","implementations","to","opt","out","of","doing","anything","here","explicitly","."]
code: def prepare_batch(self, wrapped_batch: RaggedBatchWrapper) -> None:
      """
      This will be called before each forward with the intent of building forward-specific metadata
      about a batch. The intent here is to build data structures like attention atoms without necessarily
      needing to implement graphable kernels to do so.
  
      Abstract so as to force model implementations to opt out of doing anything here explicitly.
      """
      raise NotImplementedError()
  
imports: from abc import ABC, abstractmethod
  from typing import Iterable, Optional, Tuple, Type
  import torch
  import deepspeed.comm
  from ..ragged import DSStateManager, RaggedBatchWrapper
  from ..ragged.manager_configs import KVCacheConfig
  from ..ragged import DSSequenceDescriptor
  from ..model_implementations.layer_container_base import LayerContainer
  from ..config_v2 import RaggedInferenceEngineConfig
  from .flat_model_helpers import ModelMetadata
label_desc: Use image node_id 16 for calling the DSInferenceModelBase obj's underlying member method code with example usage: obj.prepare_batch(wrapped_batch) without return types