docstr: If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel
  with sequence parallelism: we do an all_gather of x before doing the matmul.
  If sequence_parallel=False, then the input is already gathered.
  
  checkpoint_lvl:
  0: no recomputation in the bwd
  1: recompute gelu_out / relu_out in the bwd
  2: recompute pre_act and gelu_out / relu_out in the bwd
docstr_tok: ["If","process_group","is","not","None","and","sequence_parallel=True",",","we","'re","doing","Tensor","Parallel","with","sequence","parallelism",":","we","do","an","all_gather","of","x","before","doing","the","matmul",".","If","sequence_parallel=False",",","then","the","input","is","already","gathered",".","checkpoint_lvl",":","0",":","no","recomputation","in","the","bwd","1",":","recompute","gelu_out","\/","relu_out","in","the","bwd","2",":","recompute","pre_act","and","gelu_out","\/","relu_out","in","the","bwd"]
code: def forward(
      ctx,
      x,
      weight1,
      bias1,
      weight2,
      bias2,
      activation="gelu_approx",
      save_pre_act=True,
      return_residual=False,
      checkpoint_lvl=0,
      heuristic=0,
      process_group=None,
      sequence_parallel=True,
  ):
      """
      If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel
      with sequence parallelism: we do an all_gather of x before doing the matmul.
      If sequence_parallel=False, then the input is already gathered.
  
      checkpoint_lvl:
      0: no recomputation in the bwd
      1: recompute gelu_out / relu_out in the bwd
      2: recompute pre_act and gelu_out / relu_out in the bwd
      """
      assert -1 <= heuristic <= 4
      assert activation in ["gelu_approx", "relu", "sqrelu"]
      if activation == "sqrelu":
          assert heuristic == -1
      if not save_pre_act:
          checkpoint_lvl = 2
      assert checkpoint_lvl in [0, 1, 2]
      ctx.return_residual = return_residual
      ctx.process_group = process_group
      ctx.sequence_parallel = sequence_parallel
      ctx.checkpoint_lvl = checkpoint_lvl
      ctx.activation = activation
      ctx.heuristic = heuristic
  
      if torch.is_autocast_enabled():
          x = x.to(dtype=torch.get_autocast_gpu_dtype())
      x = x.contiguous()
      if process_group is not None and sequence_parallel:
          # We want to kick off the all_gather early, before weight dtype conversion
          total_x, handle_x = all_gather_raw(
              x, process_group, async_op=True
          )
      else:
          total_x = x
  
      if torch.is_autocast_enabled():
          dtype = torch.get_autocast_gpu_dtype()
          weight1, weight2 = [
              a.to(dtype=dtype) for a in [weight1, weight2]
          ]
          bias1 = bias1.to(dtype=dtype) if bias1 is not None else None
          bias2 = bias2.to(dtype=dtype) if bias2 is not None else None
      weight1 = weight1.contiguous()
      bias1 = bias1.contiguous() if bias1 is not None else None
      weight2 = weight2.contiguous()
      bias2 = bias2.contiguous() if bias2 is not None else None
      if process_group is not None and sequence_parallel:
          handle_x.wait()
      batch_shape, n = total_x.shape[:-1], total_x.shape[-1]
      batch_dim = batch_shape.numel()
      # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174
      if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32:
          raise RuntimeError(
              "fused_dense only supports matrix dims <= 2M"
          )
      if heuristic == -1:
          pre_act = F.linear(total_x, weight1, bias1)
          activation_fn = (
              partial(F.gelu, approximate="tanh")
              if activation == "gelu_approx"
              else (sqrelu_fwd if activation == "sqrelu" else F.relu)
          )
          with torch.jit.fuser("fuser2"):
              output1 = activation_fn(pre_act)
          # This is before adding bias1
          # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1)
          # with torch.jit.fuser('fuser2'):
          #     output1 = bias_gelu(pre_act, bias1)
      else:
          is_gelu = activation == "gelu_approx"
          output1, *rest = fused_dense_cuda.linear_act_forward(
              total_x.reshape(batch_dim, n),
              weight1,
              bias1,
              is_gelu,
              save_pre_act,
              heuristic,
          )
          if save_pre_act:
              pre_act = rest[0]
      output2 = F.linear(output1, weight2, bias2)
      if checkpoint_lvl == 0 or (
          checkpoint_lvl == 1 and activation == "relu"
      ):
          # For RELU the pre_act is very small (just a bit-mask) so we just save it
          ctx.save_for_backward(x, weight1, weight2, pre_act, output1)
      elif checkpoint_lvl == 1:
          ctx.save_for_backward(x, weight1, weight2, pre_act)
      elif checkpoint_lvl == 2:
          ctx.save_for_backward(x, weight1, weight2, bias1)
      output2 = output2.reshape(*batch_shape, output2.shape[-1])
      return output2 if not return_residual else (output2, x)
  
imports: from functools import partial
  from typing import Optional
  import fused_dense_lib
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
  from torch.cuda.amp import custom_bwd, custom_fwd
  from torch.distributed import ProcessGroup
  from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd
  from flash_attn.utils.distributed import all_gather_raw, all_reduce, all_reduce_raw, reduce_scatter, reduce_scatter_raw
label_desc: Use image node_id 1 for calling the FusedMLPFunc obj's underlying member method code with example usage: obj.forward(ctx, x, weight1, bias1, weight2, bias2, activation, save_pre_act, return_residual, checkpoint_lvl, heuristic, process_group, sequence_parallel) and returns: unknown