Add TransformerEngine

recipes/TransformerEngine/recipe.yaml

@@ -0,0 +1,71 @@
+schema_version: 1
+
+context:
+  name: TransformerEngine
+  version: 2.1
+
+package:
+  name: ${{ name|lower }}
+  version: ${{ version }}
+
+source:
+  url: https://github.com/NVIDIA/TransformerEngine/archive/v${{ version }}.tar.gz
+  sha256: 144338fda47712b647adcb71ed40bc9ee69c626c10dac5c408487208fb4bdc98
+
+build:
+  number: 0
+  skip:
+    - not linux
+  script:
+    env:
+      NVTE_FRAMEWORK: pytorch
+    content:
+      - mkdir -p ${SRC_DIR}/3rdparty/cudnn-frontend/
+      - ln -s ${PREFIX}/lib/python${PY_VER}/site-packages/include ${SRC_DIR}/3rdparty/cudnn-frontend/include
+      - ${{ PYTHON }} -m pip install . -vvv --no-deps --no-build-isolation
+
+requirements:
+  build:
+    - ${{ compiler('c') }}
+    - ${{ compiler('cxx') }}
+    - ${{ compiler('cuda') }}
+    - ${{ stdlib('c') }}
+    - cmake >=3.21
+    - ninja
+  host:
+    - cuda-cudart-dev
+    - cuda-crt-dev_linux-64
+    - cudnn
+    - pip
+    - pybind11
+    - python
+    - python-cudnn-frontend
+    - pytorch
+  run:
+    # - flash-attn >=2.1.1 # https://github.com/NVIDIA/TransformerEngine/pull/1226
+    - packaging
+    - pydantic
+    - python
+    - pytorch
+
+tests:
+  - python:
+      imports:
+        - transformer_engine
+      pip_check: true
+  - requirements:
+      run:
+        - pip
+    script: commands
+
+about:
+  summary: Transformer acceleration library
+  license: Apache-2.0
+  license_file: LICENSE
+  homepage: https://github.com/NVIDIA/TransformerEngine
+  repository: https://github.com/NVIDIA/TransformerEngine
+  documentation: https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/index.html
+
+extra:
+  recipe-maintainers:
+    - weiji14