class_docstr: Operator that computes memory-efficient attention using         `Flash-Attention <https://github.com/HazyResearch/flash-attention>`_         implementation.
      
class_docstr_tok: ["Operator","that","computes","memory-efficient","attention","using","`","Flash-Attention","<","https",":","\/\/github.com\/HazyResearch\/flash-attention",">","`","_","implementation","."]
code: def apply(
      cls, inp: Inputs, needs_gradient: bool
  ) -> Tuple[torch.Tensor, Optional[Context]]:
      return_softmax = False
      out_shape = [
          *inp.query.shape[:-1],
          inp.value.shape[-1],
      ]
      # no cumulative seqlen
      (
          inp,
          cu_seqlens_q,
          max_seqlen_q,
          cu_seqlens_k,
          max_seqlen_k,
          seqused_k,
      ) = _convert_input_format(inp, supports_mqa=True)
      if inp.query.numel() > 0 and inp.key.numel() > 0:
          win_left, win_right = _window_size(inp.attn_bias)
          out, softmax_lse, rng_state = cls.OPERATOR(
              inp.query,
              inp.key,
              inp.value,
              cu_seqlens_q,
              cu_seqlens_k,
              seqused_k,
              max_seqlen_q,
              max_seqlen_k,
              inp.p,
              inp.scale_float,
              _is_causal(inp.attn_bias),
              window_left=win_left,
              window_right=win_right,
              return_softmax=return_softmax,
          )
          out = out.reshape(out_shape)
      else:
          out = torch.zeros(
              out_shape, device=inp.query.device, dtype=inp.query.dtype
          )
          rng_state = None
          softmax_lse = torch.empty(
              [
                  inp.query.shape[0],
                  inp.query.shape[2],
                  inp.query.shape[1],
              ],
              device=inp.query.device,
              dtype=torch.float32,
          )
      ctx = Context(out=out, lse=softmax_lse)
      if inp.p != 0.0:
          ctx.op_bw = BwOp
          ctx.rng_state = rng_state
      return (out, ctx)
  
imports: import os
  from dataclasses import replace
  from itertools import zip_longest
  from typing import Any, List, Optional, Set, Tuple, Union
  import torch
  from ..common import _get_storage_base, get_operator, register_operator
  from .attn_bias import AttentionBias, BlockDiagonalCausalFromBottomRightMask, BlockDiagonalCausalLocalAttentionFromBottomRightMask, BlockDiagonalCausalLocalAttentionMask, BlockDiagonalCausalMask, BlockDiagonalCausalWithOffsetPaddedKeysMask, BlockDiagonalMask, LocalAttentionFromBottomRightMask, LowerTriangularFromBottomRightLocalAttentionMask, LowerTriangularFromBottomRightMask, LowerTriangularMask
  from .common import AttentionBwOpBase, AttentionFwOpBase, Context, Gradients, Inputs, 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(cls, inp, needs_gradient) and returns: out,  ctx