docstr: checkpoint_lvl:
  0: no recomputation in the bwd
  1: recompute gelu_out in the bwd
  2: recompute act_input and gelu_out in the bwd
docstr_tok: ["checkpoint_lvl",":","0",":","no","recomputation","in","the","bwd","1",":","recompute","gelu_out","in","the","bwd","2",":","recompute","act_input","and","gelu_out","in","the","bwd"]
code: def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0):
      """checkpoint_lvl:
      0: no recomputation in the bwd
      1: recompute gelu_out in the bwd
      2: recompute act_input and gelu_out in the bwd
      """
      if torch.is_autocast_enabled():
          dtype = torch.get_autocast_gpu_dtype()
          x, weight1, bias1, weight2, bias2 = [
              a.to(dtype=dtype)
              for a in [x, weight1, bias1, weight2, bias2]
          ]
      is_bf16 = x.dtype == torch.bfloat16
      assert checkpoint_lvl in [0, 1, 2]
      x = x.contiguous()
      weight1 = weight1.contiguous()
      bias1 = bias1.contiguous()
      weight2 = weight2.contiguous()
      bias2 = bias2.contiguous()
      batch_shape, n = x.shape[:-1], x.shape[-1]
      batch_dim = batch_shape.numel()
      if is_bf16:
          act_input = fused_dense_cuda.linear_bias_forward(
              x.reshape(batch_dim, n), weight1, bias1
          )
          output1 = sqrelu_fwd(act_input)
      else:
          save_act_input = checkpoint_lvl != 2
          result = triton_linear_act(
              x.reshape(batch_dim, n),
              weight1,
              bias1,
              activation="squared_relu",
              save_act_input=save_act_input,
          )
          if save_act_input:
              output1, act_input = result
          else:
              output1 = result
      output2 = fused_dense_cuda.linear_bias_forward(
          output1, weight2, bias2
      )
      ctx.checkpoint_lvl = checkpoint_lvl
      if checkpoint_lvl == 0:
          ctx.save_for_backward(
              x, weight1, bias1, weight2, act_input, output1
          )
      elif checkpoint_lvl == 1:
          ctx.save_for_backward(x, weight1, bias1, weight2, act_input)
      elif checkpoint_lvl == 2:
          ctx.save_for_backward(x, weight1, bias1, weight2)
      return output2.reshape(*batch_shape, output2.shape[-1])
  
file_path: flash-attention/flash_attn/ops/triton/mlp.py
imports: import fused_dense_lib
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.cuda.amp import custom_bwd, custom_fwd
  from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd
  from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act
label_desc: Use image node_id 1 for calling the FusedDenseSqreluDenseFunc obj's underlying member method code with example usage: obj.forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl) and returns: output2