docstr: If process_group is not None, we're doing Tensor Parallel with sequence parallelism:
  we do an all_gather of x before doing the matmul, gelu, then matmul.
  Finally we do a reduce_scatter of the output.
  
  checkpoint_lvl (increasing lvl means slower but more memory saving):
      0: no recomputation in the bwd
      1: recompute gelu_out in the bwd
      2: recompute pre_act and gelu_out in the bwd
  heuristic:
      -1: don't fuse gemm + gelu (separate kernel)
      0..4: use this heuristic for the algo section in the fused gemm + gelu
      'auto': heuristic will be picked automatically:
          For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf.
          For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16.
          For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation
          is slower than the unfused version.
  return_residual: whether to return the input x along with the output. This is for
      performance reason: for post-norm architecture, returning the input allows us
      to fuse the backward of nn.Linear with the residual connection.
docstr_tok: ["If","process_group","is","not","None",",","we","'re","doing","Tensor","Parallel","with","sequence","parallelism",":","we","do","an","all_gather","of","x","before","doing","the","matmul",",","gelu",",","then","matmul",".","Finally","we","do","a","reduce_scatter","of","the","output",".","checkpoint_lvl","(","increasing","lvl","means","slower","but","more","memory","saving",")",":","0",":","no","recomputation","in","the","bwd","1",":","recompute","gelu_out","in","the","bwd","2",":","recompute","pre_act","and","gelu_out","in","the","bwd","heuristic",":","-1",":","do","n't","fuse","gemm","+","gelu","(","separate","kernel",")","0","..","4",":","use","this","heuristic","for","the","algo","section","in","the","fused","gemm","+","gelu","'auto","'",":","heuristic","will","be","picked","automatically",":","For","CUDA",">","=","11.8",",","we","set","heuristic=0","for","both","fp16","and","bf16","for","best","perf",".","For","CUDA","<","=","11.7",",","we","set","heuristic=1","for","fp16","and","heuristic=-1","for","bf16",".","For","H100",",","we","set","heuristic=-1","for","both","fp16","and","bf16","as","the","fused","cuBlasLt","implementation","is","slower","than","the","unfused","version",".","return_residual",":","whether","to","return","the","input","x","along","with","the","output",".","This","is","for","performance","reason",":","for","post-norm","architecture",",","returning","the","input","allows","us","to","fuse","the","backward","of","nn.Linear","with","the","residual","connection","."]
code: def __init__(
      self,
      in_features,
      hidden_features=None,
      out_features=None,
      bias1=True,
      bias2=True,
      activation="gelu_approx",
      return_residual=False,
      checkpoint_lvl=0,
      heuristic="auto",
      device=None,
      dtype=None,
  ):
      """
      If process_group is not None, we're doing Tensor Parallel with sequence parallelism:
      we do an all_gather of x before doing the matmul, gelu, then matmul.
      Finally we do a reduce_scatter of the output.
  
      checkpoint_lvl (increasing lvl means slower but more memory saving):
          0: no recomputation in the bwd
          1: recompute gelu_out in the bwd
          2: recompute pre_act and gelu_out in the bwd
      heuristic:
          -1: don't fuse gemm + gelu (separate kernel)
          0..4: use this heuristic for the algo section in the fused gemm + gelu
          'auto': heuristic will be picked automatically:
              For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf.
              For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16.
              For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation
              is slower than the unfused version.
      return_residual: whether to return the input x along with the output. This is for
          performance reason: for post-norm architecture, returning the input allows us
          to fuse the backward of nn.Linear with the residual connection.
      """
      assert checkpoint_lvl in [0, 1, 2]
      assert activation in ["gelu_approx", "relu", "sqrelu"]
      factory_kwargs = {"device": device, "dtype": dtype}
      super().__init__()
      out_features = out_features or in_features
      hidden_features = hidden_features or in_features * 4
      self.activation = activation
      self.return_residual = return_residual
      self.checkpoint_lvl = checkpoint_lvl
      self.heuristic = heuristic if activation != "sqrelu" else -1
      self.fc1 = nn.Linear(
          in_features, hidden_features, bias=bias1, **factory_kwargs
      )
      self.fc2 = nn.Linear(
          hidden_features, out_features, bias=bias2, **factory_kwargs
      )
  
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 to create a new FusedMLP object from inherited base classes: nn with example: obj = FusedMLP(in_features, hidden_features, out_features, bias1, bias2, activation, return_residual, checkpoint_lvl, heuristic, device, dtype)