code: def __init__(
      self,
      B: int,
      Mq: int,
      Mkv: int,
      Hq: int,
      Hkv: int,
      K: int,
      bw: bool,
  ) -> None:
      dtype = torch.float16
      self.sub_label = (
          f"B={B} Mq={Mq} Mkv={Mkv} Hq={Hq} Hkv={Hkv} K={K}"
      )
      self.label = "attn_decoding"
      self.shapes = (B, Mq, Mkv, Hq, Hkv, K)
  
      assert Hkv <= Hq
      assert Hq % Hkv == 0
      self.q = torch.randn(
          [B, Mq, Hkv, Hq // Hkv, K],
          device="cuda",
          dtype=dtype,
          requires_grad=bw,
      )
      self.k = torch.randn(
          [B, Mkv, Hkv, 1, K],
          device="cuda",
          dtype=dtype,
          requires_grad=bw,
      ).expand(-1, -1, -1, Hq // Hkv, -1)
      self.v = torch.randn(
          [B, Mkv, Hkv, 1, K],
          device="cuda",
          dtype=dtype,
          requires_grad=bw,
      ).expand(-1, -1, -1, Hq // Hkv, -1)
  
      if Hq == Hkv:
          self.q = self.q[:, :, :, 0]
          self.k = self.k[:, :, :, 0]
          self.v = self.v[:, :, :, 0]
      if Hkv == 1:
          self.q = self.q[:, :, 0]
          self.k = self.k[:, :, 0]
          self.v = self.v[:, :, 0]
  
file_path: xformers/xformers/benchmarks/benchmark_attn_decoding.py
imports: from typing import Any
  import torch
  from torch.utils import benchmark
  from utils import benchmark_main_helper2
  import xformers.ops
label_desc: Use image node_id 1 to create a new AttentionDecodingFlashDecoding object  with example: obj = AttentionDecodingFlashDecoding(B, Mq, Mkv, Hq, Hkv, K, bw)