[Enhancement] Update abs function for half_t and bfloat_t to use cutlass implementation (#1023)

* [Enhancement] Update abs function for half_t and bfloat_t to use cutlass implementation

* [Lint]: [pre-commit.ci] auto fixes [...]

* optimize amd ci

---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>

.github/workflows/rocm-ci.yml

@@ -31,7 +31,7 @@ jobs:
     - name: Ensure venv (local & persistent)
       run: |
         set -e
-        REQS_HASH=$(sha256sum requirements-test.txt 2>/dev/null | awk '{print $1}' || echo "no_requirements")
+        REQS_HASH=$(sha256sum requirements-rocm.txt 2>/dev/null | awk '{print $1}' || echo "no_requirements")
         MARKER="${{ runner.tool_cache }}/.venv_marker_${{ env.PYTHON_VERSION }}_${REQS_HASH:0:8}"
 
         if [[ -f "$MARKER" ]] && [[ -f "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate" ]]; then
@@ -43,8 +43,8 @@ jobs:
           # shellcheck source=/dev/null
           source "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate"
           python -m pip install --upgrade pip --no-user
-          [[ -f requirements-test.txt ]] && \
-            PIP_NO_BUILD_ISOLATION=1 pip install -r requirements-test.txt --no-user
+          [[ -f requirements-rocm.txt ]] && \
+            PIP_NO_BUILD_ISOLATION=1 pip install -r requirements-rocm.txt --no-user
           touch "$MARKER"
         fi
 
@@ -84,26 +84,21 @@ jobs:
 
     - name: Ensure venv (local & persistent)
       run: |
-        echo "Running on AMD GPU"
         set -e
-        REQS_HASH=$(sha256sum requirements-rocm.txt | cut -d ' ' -f 1)
+        REQS_HASH=$(sha256sum requirements-rocm.txt 2>/dev/null | awk '{print $1}' || echo "no_requirements")
         MARKER="${{ runner.tool_cache }}/.venv_marker_${{ env.PYTHON_VERSION }}_${REQS_HASH:0:8}"
 
-        echo "Installing requirements"
         if [[ -f "$MARKER" ]] && [[ -f "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate" ]]; then
           echo "venv exists and hash matches – reuse it"
         else
           echo "venv stale or missing – recreating"
           rm -rf "${{ runner.tool_cache }}/${{ env.VENV_DIR }}" "$MARKER"
           python -m venv "${{ runner.tool_cache }}/${{ env.VENV_DIR }}"
+          # shellcheck source=/dev/null
           source "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate"
           python -m pip install --upgrade pip --no-user
-          if [[ -f requirements-rocm.txt ]]; then
-            pip install --pre torch torchvision torchaudio --index-url ${{ env.PYTORCH_INDEX_URL }}
-            PIP_NO_BUILD_ISOLATION=1 pip install -r requirements-rocm.txt
-          fi
-
-          USE_ROCM=True pip install . --no-user
+          [[ -f requirements-rocm.txt ]] && \
+            PIP_NO_BUILD_ISOLATION=1 pip install -r requirements-rocm.txt --no-user
           touch "$MARKER"
         fi
 

src/tl_templates/cuda/common.h

@@ -54,11 +54,11 @@ using int4_t = int4;
     }                                                                          \
   } while (0)
 
-// abs function for bfloat_t and half_t since there is no implicit conversion
-// method
-TL_PATCH TL_DEVICE half_t __habs(const half_t x) {
-  return half_t(__habs(x.to_half()));
-}
+// using cutlass abs function for half_t
+TL_PATCH TL_DEVICE half_t __habs(const half_t x) { return abs(x); }
+
+// using cutlass abs function for bfloat_t
+TL_PATCH TL_DEVICE bfloat16_t __habs(const bfloat16_t x) { return abs(x); }
 
 // hrsqrt function for half_t
 TL_PATCH TL_DEVICE half_t hrsqrt(const half_t x) {