class_docstr: A module which performs QKV attention.
  Attention mask is accepted as input.
code: def count_flops(model, _x, y):
      return count_flops_attn(model, _x, y)
  
imports: import math
  import numpy
  import torch
  import torch.nn
  from torch.cuda.amp import custom_bwd, custom_fwd
label_desc: Use image node_id 3 for calling the QKVStableMaskedAttention obj's underlying member method code with example usage: obj.count_flops(model, _x, y) and returns: count_flops_attn