docstr: q: (batch_size, seqlen_q, nheads, headdim)
  k, v: (batch_size, seqlen_k, nheads, headdim)
  bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).
      For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).
      ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)
docstr_tok: ["q",":","(","batch_size",",","seqlen_q",",","nheads",",","headdim",")","k",",","v",":","(","batch_size",",","seqlen_k",",","nheads",",","headdim",")","bias",":","optional",",","shape","broadcastible","to","(","batch",",","nheads",",","seqlen_q",",","seqlen_k",")",".","For","example",",","ALiBi","mask","for","causal","would","have","shape","(","1",",","nheads",",","1",",","seqlen_k",")",".","ALiBi","mask","for","non-causal","would","have","shape","(","1",",","nheads",",","seqlen_q",",","seqlen_k",")"]
code: def forward(
      ctx, q, k, v, bias=None, causal=False, softmax_scale=None
  ):
      """
      q: (batch_size, seqlen_q, nheads, headdim)
      k, v: (batch_size, seqlen_k, nheads, headdim)
      bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).
          For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).
          ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)
      """
      # Make sure that the last dimension is contiguous
      q, k, v = [
          x if x.stride(-1) == 1 else x.contiguous() for x in [q, k, v]
      ]
      o, lse, ctx.softmax_scale = _flash_attn_forward(
          q, k, v, bias=bias, causal=causal, softmax_scale=softmax_scale
      )
      ctx.save_for_backward(q, k, v, o, lse, bias)
      ctx.causal = causal
      return o
  
imports: import math
  import torch
  import triton
  import triton.language
label_desc: Use image node_id 1 for calling the FlashAttnFunc obj's underlying member method code with example usage: obj.forward(ctx, q, k, v, bias, causal, softmax_scale) and returns: o