code: def forward(
      ctx, qkv, cu_seqlens, p_dropout, max_s, is_training, zero_tensors
  ):
      batch_size = cu_seqlens.numel() - 1
      if batch_size < 4:
          max_s = 512
          context, S_dmask = mha.fwd_nl(
              qkv,
              cu_seqlens,
              p_dropout,
              max_s,
              is_training,
              True,
              zero_tensors,
              None,
          )
      else:
          context, S_dmask = mha.fwd(
              qkv,
              cu_seqlens,
              p_dropout,
              max_s,
              is_training,
              False,
              zero_tensors,
              None,
          )
      ctx.save_for_backward(qkv, S_dmask)
      ctx.cu_seqlens = cu_seqlens
      ctx.p_dropout = p_dropout
      ctx.max_s = max_s
      ctx.zero_tensors = zero_tensors
      return context
  
file_path: apex/apex/contrib/fmha/fmha.py
imports: import torch
  import torch.nn.functional
  import fmhalib
label_desc: Use image node_id 1 for calling the FMHAFun obj's underlying member method code with example usage: obj.forward(ctx, qkv, cu_seqlens, p_dropout, max_s, is_training, zero_tensors) and returns: context