class_docstr: C++ implementation to populate the attention atoms for the blocked attention
  kernel.
docstr: Populates the attention atoms for the blocked attention kernel.
  
  Args:
      atoms (torch.Tensor): Pre-allocated int32 tensor of shape [max_atoms, 8]
      ragged_batch (torch.Tensor): Wrapper for the ragged batch.
      q_block_size (int): The block size for the queries (as determined by the
          attention implementation)
      kv_block_size (int): The block size for the keys/values (as determined by the
          attention implementation)
  
  Returns:
docstr_tok: ["Populates","the","attention","atoms","for","the","blocked","attention","kernel",".","Args",":","atoms","(","torch.Tensor",")",":","Pre-allocated","int32","tensor","of","shape","[","max_atoms",",","8","]","ragged_batch","(","torch.Tensor",")",":","Wrapper","for","the","ragged","batch",".","q_block_size","(","int",")",":","The","block","size","for","the","queries","(","as","determined","by","the","attention","implementation",")","kv_block_size","(","int",")",":","The","block","size","for","the","keys\/values","(","as","determined","by","the","attention","implementation",")","Returns",":"]
code: def __call__(
      self,
      atoms: torch.Tensor,
      ragged_batch: RaggedBatchWrapper,
      q_block_size: int,
      kv_block_size: int,
  ) -> Tuple[torch.Tensor, int]:
      """
      Populates the attention atoms for the blocked attention kernel.
  
      Args:
          atoms (torch.Tensor): Pre-allocated int32 tensor of shape [max_atoms, 8]
          ragged_batch (torch.Tensor): Wrapper for the ragged batch.
          q_block_size (int): The block size for the queries (as determined by the
              attention implementation)
          kv_block_size (int): The block size for the keys/values (as determined by the
              attention implementation)
  
      Returns:
  
      """
      if atoms.device != torch.device("cpu"):
          raise RuntimeError("AtomBuilder must be called on tensors")
  
      n_atoms = self.kernel(
          atoms,
          ragged_batch.batch_metadata_buffer(on_device=False),
          ragged_batch.inflight_seq_descriptors(on_device=False),
          ragged_batch.kv_ptrs(on_device=False),
          q_block_size,
          kv_block_size,
      )
      return atoms, n_atoms
  
file_path: deepspeed/deepspeed/inference/v2/kernels/ragged_ops/atom_builder/atom_builder.py
imports: from typing import Tuple
  import torch
  from ...None import DSKernelBase
  from deepspeed.ops.op_builder import RaggedOpsBuilder
  from ....ragged import RaggedBatchWrapper
label_desc: Use image node_id 2 for calling the AtomBuilder obj's underlying member method code with example usage: obj.__call__(atoms, ragged_batch, q_block_size, kv_block_size) and returns: atoms,  n_atoms