class_docstr: Manage tensors stored in a compressed format.
class_docstr_tok: ["Manage","tensors","stored","in","a","compressed","format","."]
code: def init_attention_compute_workspace(self, config, task, policy):
      if self.base_device.device_type != DeviceType.CPU:
          return  # Only CPU requires this fp32 workspace
  
      b = policy.gpu_batch_size
      n_head = config.n_head
      head_dim = config.input_dim // n_head
      max_seq_len = task.prompt_len + task.gen_len - 1
      shape = (max_seq_len, b * n_head, head_dim)
  
      group_size, group_dim = (
          policy.comp_cache_config.group_size,
          policy.comp_cache_config.group_dim,
      )
      num_groups = (shape[group_dim] + group_size - 1) // group_size
      new_shape = (
          shape[:group_dim]
          + (num_groups, group_size)
          + shape[group_dim + 1 :]
      )
  
      self.data_decompress_workspace = [
          torch.empty(
              *new_shape,
              dtype=torch.float32,
              device=self.base_device.dev,
          ),
          torch.empty(
              *new_shape,
              dtype=torch.float32,
              device=self.base_device.dev,
          ),
      ]
  
imports: import dataclasses
  import torch
  import numpy
  from flexgen.pytorch_backend import TorchTensor, TorchDevice, DeviceType, general_copy, fix_recursive_import
  from flexgen.utils import np_dtype_to_torch_dtype
label_desc: Use image node_id 4 for calling the TorchCompressedDevice obj's underlying member method code with example usage: obj.init_attention_compute_workspace(config, task, policy) without return types