### Setting up kernels

- **Custom CUDA layernorm kernels** modified from [FastFold](https://github.com/hpcaitech/FastFold) and [Oneflow](https://github.com/Oneflow-Inc/oneflow) accelerate about 30%-50% during different training stages. To use this feature, run the following command:
  ```bash
  export LAYERNORM_TYPE=fast_layernorm
  ```
  If the environment variable `LAYERNORM_TYPE` is set to `fast_layernorm`, the model will employ the layernorm we have developed; otherwise, the naive PyTorch layernorm will be adopted. The kernels will be compiled when `fast_layernorm` is called for the first time.
- **[DeepSpeed DS4Sci_EvoformerAttention kernel](https://www.deepspeed.ai/tutorials/ds4sci_evoformerattention/)** is a memory-efficient attention kernel developed as part of a collaboration between OpenFold and the DeepSpeed4Science initiative. To use this feature, run the following command:
  ```bash
  export USE_DEEPSPEED_EVO_ATTTENTION=true
  ```
  DS4Sci_EvoformerAttention is implemented based on [CUTLASS](https://github.com/NVIDIA/cutlass). If you use this feature, You need to clone the CUTLASS repository and specify the path to it in the environment variable CUTLASS_PATH. The [Dockerfile](Dockerfile) has already include this setting:
  ```bash
  RUN git clone -b v3.5.1 https://github.com/NVIDIA/cutlass.git  /opt/cutlass
  ENV CUTLASS_PATH=/opt/cutlass
  ```
  If you set up `Protenix` by `pip`, you can set environment variable `CUTLASS_PATH` as follows:

  ```bash
  git clone -b v3.5.1 https://github.com/NVIDIA/cutlass.git  /path/to/cutlass
  export CUTLASS_PATH=/path/to/cutlass
  ```

  The kernels will be compiled when DS4Sci_EvoformerAttention is called for the first time.