docstr: qkv: (batch, seqlen, 3, nheads, headdim)
  bias: optional, shape broadcastible to (batch, nheads, seqlen, seqlen).
      For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen).
      ALiBi mask for non-causal would have shape (1, nheads, seqlen, seqlen)
docstr_tok: ["qkv",":","(","batch",",","seqlen",",","3",",","nheads",",","headdim",")","bias",":","optional",",","shape","broadcastible","to","(","batch",",","nheads",",","seqlen",",","seqlen",")",".","For","example",",","ALiBi","mask","for","causal","would","have","shape","(","1",",","nheads",",","1",",","seqlen",")",".","ALiBi","mask","for","non-causal","would","have","shape","(","1",",","nheads",",","seqlen",",","seqlen",")"]
code: def forward(ctx, qkv, bias=None, causal=False, softmax_scale=None):
      """
      qkv: (batch, seqlen, 3, nheads, headdim)
      bias: optional, shape broadcastible to (batch, nheads, seqlen, seqlen).
          For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen).
          ALiBi mask for non-causal would have shape (1, nheads, seqlen, seqlen)
      """
      # Make sure that the last dimension is contiguous
      if qkv.stride(-1) != 1:
          qkv = qkv.contiguous()
      o, lse, ctx.softmax_scale = _flash_attn_forward(
          qkv[:, :, 0],
          qkv[:, :, 1],
          qkv[:, :, 2],
          bias=bias,
          causal=causal,
          softmax_scale=softmax_scale,
      )
      ctx.save_for_backward(qkv, 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 FlashAttnQKVPackedFunc obj's underlying member method code with example usage: obj.forward(ctx, qkv, bias, causal, softmax_scale) and returns: o