class_docstr: Operator that computes memory-efficient attention using         `Tri Dao's <https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/flash_attn_triton.py>`_         implementation, based on
  `Phil Tillet's code <https://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py>`_
code: def apply(
      cls, inp: Inputs, needs_gradient: bool
  ) -> Tuple[torch.Tensor, Optional[Context]]:
      inp = _prepare_inputs(inp)
  
      out, lse, softmax_scale = triton_flash_forward(
          q=inp.query,
          k=inp.key,
          v=inp.value,
          bias=inp.attn_bias
          if isinstance(inp.attn_bias, torch.Tensor)
          else None,
          softmax_scale=inp.scale_float,
          causal=isinstance(inp.attn_bias, LowerTriangularMask),
      )
      return out, Context(lse=lse, out=out)
  
file_path: xformers/xformers/ops/fmha/triton.py
imports: from dataclasses import replace
  from typing import TYPE_CHECKING, Any, List, Optional, Set, Tuple
  import torch
  from ...None import _is_triton_available
  from ..common import register_operator
  from .attn_bias import 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,  Context