class_docstr: xFormers' MHA kernel based on CUTLASS.
  Supports a large number of settings (including without TensorCores, f32 ...)
  and GPUs as old as P100 (Sm60)
code: def apply_bmhk(
      cls, inp: Inputs, needs_gradient: bool
  ) -> Tuple[torch.Tensor, Optional[Context]]:
      if type(inp.attn_bias) not in FwOp.SUPPORTED_ATTN_BIAS_TYPES:
          raise NotImplementedError("Unsupported attn_bias type")
      seqstart_k, seqstart_q, max_seqlen_q, _ = _get_seqlen_info(inp)
      out, lse, rng_seed, rng_offset = cls.OPERATOR(
          query=inp.query,
          key=inp.key,
          value=inp.value,
          attn_bias=_get_tensor_bias(inp.attn_bias),
          seqstart_q=seqstart_q,
          seqstart_k=seqstart_k,
          max_seqlen_q=max_seqlen_q,
          dropout_p=inp.p,
          compute_logsumexp=needs_gradient,
          custom_mask_type=_custom_mask_type(inp.attn_bias),
          scale=inp.scale,
          seqlen_k=inp.attn_bias.k_seqinfo.seqlen
          if isinstance(
              inp.attn_bias, BlockDiagonalCausalWithOffsetPaddedKeysMask
          )
          else None,
          window_size=inp.attn_bias._window_size
          if isinstance(
              inp.attn_bias,
              (
                  BlockDiagonalCausalLocalAttentionMask,
                  BlockDiagonalCausalLocalAttentionFromBottomRightMask,
                  LowerTriangularFromBottomRightLocalAttentionMask,
              ),
          )
          else None,
      )
      ctx: Optional[Context] = None
      if needs_gradient:
          ctx = Context(
              out=out,
              lse=lse,
              # cutlass forward is only compatible with cutlass backward if
              # dropout is used (because of the way RNG states are passed and the
              # way random numbers are generated during backward)
              op_bw=BwOp if inp.p != 0 else None,
          )
          if inp.p != 0:
              ctx.rng_state = torch.tensor(
                  [rng_seed, rng_offset],
                  dtype=torch.int64,
                  device="cpu",
              )
      return out, ctx
  
imports: from dataclasses import replace
  from enum import Enum
  from functools import partial
  from typing import Any, List, Optional, Set, Tuple, Union
  import torch
  from ..common import get_xformers_operator, register_operator
  from .None import attn_bias
  from .attn_bias import AttentionBias, BlockDiagonalCausalLocalAttentionFromBottomRightMask, BlockDiagonalCausalLocalAttentionMask, BlockDiagonalCausalMask, BlockDiagonalCausalWithOffsetPaddedKeysMask, BlockDiagonalMask, LowerTriangularFromBottomRightLocalAttentionMask, LowerTriangularFromBottomRightMask, LowerTriangularMask, LowerTriangularMaskWithTensorBias
  from .common import AttentionBwOpBase, AttentionFwOpBase, Context, Gradients, Inputs, _attn_bias_apply, check_lastdim_alignment_stride1
label_desc: Use image node_id 2 for calling the FwOp obj's underlying member method code with example usage: obj.apply_bmhk(cls, inp, needs_gradient) and returns: out,  ctx