initial commit

scripts/build.sh

@@ -0,0 +1,49 @@
+#!/bin/bash
+# shellcheck disable=SC2329
+
+set -e
+
+set_max_jobs() {
+  export MAX_JOBS="${MAX_JOBS:-16}"
+  echo "Setting MAX_JOBS to $MAX_JOBS"
+}
+
+PY_BUILDTAG=$(python -c "import sys, sysconfig; print(f'{sysconfig.get_platform()}-{sys.implementation.cache_tag}')")
+PLATFORM_TAG="${DEEPINFER_ACCELERATOR:-nvidia}-$(uname -m)"
+BUILD_BASE="build/$PLATFORM_TAG"
+SO_DST="deepinfer/_so/$PLATFORM_TAG"
+
+_build() {
+  set_max_jobs
+  python setup.py build --build-base "$BUILD_BASE"
+  (
+    set +e
+    set -x
+    shopt -s failglob
+    mkdir -p "$SO_DST"
+    for so in "${BUILD_BASE}/lib.${PY_BUILDTAG}/deepinfer/"*.so; do
+      ln -fs "$(realpath --relative-to="$SO_DST" "$so")" "$SO_DST/"
+    done
+  )
+  [ -n "$(ls -A "$SO_DST")" ] || { echo "No .so was linked into $SO_DST" >&2; exit 1; }
+  python -m compileall deepinfer -j "$MAX_JOBS" -q
+}
+
+_clean() {
+  set +e
+  set -x
+  rm -rf "$BUILD_BASE" "$SO_DST"
+  rm -f deepinfer/*.so
+  find . -type d -name "__pycache__" -exec rm -rf {} +
+}
+
+case "$1" in
+  build | clean)
+    "_$1"
+    exit 0
+    ;;
+  *)
+    echo "Unknown command: '$1'"
+    exit 1
+    ;;
+esac

scripts/build_vendor.sh

@@ -0,0 +1,56 @@
+#!/bin/bash
+set -e
+
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+VENDOR_DIR="${REPO_ROOT}/vendor"
+
+case "${DEEPINFER_ACCELERATOR:-nvidia}" in
+  nvidia)
+    BUILD_TARGET=CUDA
+    PACKAGES=(DeepSelect FlashMLA DeepEP DeepGEMM)
+    INSTALL=(deep_select flash_mla deep_ep deep_gemm)
+    ;;
+  ascend)
+    BUILD_TARGET=ASCEND
+    PACKAGES=(DeepSelect FlashMLA DeepEP-Ascend DeepGEMM-Ascend)
+    INSTALL=(deep_select flash_mla deep_ep deep_gemm)
+    ;;
+  *)
+    echo "Unsupported DEEPINFER_ACCELERATOR: '${DEEPINFER_ACCELERATOR}' (expect nvidia or ascend)" >&2
+    exit 1
+    ;;
+esac
+
+PLATFORM_TAG="${DEEPINFER_ACCELERATOR:-nvidia}-$(uname -m)"
+WHEEL_DIR="${WHEEL_DIR:-${REPO_ROOT}/build/${PLATFORM_TAG}/vendor-wheels}"
+
+GIT_REWRITE=(-c url."https://github.com/".insteadOf="git@github.com:")
+
+echo "📥 Initializing submodules ..."
+git -C "${REPO_ROOT}" "${GIT_REWRITE[@]}" submodule update --init
+
+target_var() {
+  case "$1" in
+    DeepSelect) echo DEEP_SELECT_BUILD_TARGET_PLATFORM ;;
+    FlashMLA) echo FLASH_MLA_BUILD_TARGET_PLATFORM ;;
+    *) echo "" ;;
+  esac
+}
+
+mkdir -p "${WHEEL_DIR}"
+for pkg in "${PACKAGES[@]}"; do
+  echo "🔧 Building ${pkg} (${BUILD_TARGET}) ..."
+  envs=()
+  var="$(target_var "${pkg}")"
+  [ -n "${var}" ] && envs+=("${var}=${BUILD_TARGET}")
+  [ "${pkg}" = DeepGEMM ] && envs+=(DG_FORCE_BUILD=1)
+  (
+    cd "${VENDOR_DIR}/${pkg}"
+    git "${GIT_REWRITE[@]}" submodule update --init --recursive
+    env "${envs[@]}" python setup.py bdist_wheel --dist-dir "${WHEEL_DIR}"
+  )
+done
+
+echo "📦 Installing wheels from ${WHEEL_DIR} ..."
+python -m pip install --no-deps --no-index --find-links "${WHEEL_DIR}" "${INSTALL[@]}"
+echo "✅ vendor operators installed"

scripts/format_csrc.sh

@@ -0,0 +1,73 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="${SCRIPT_DIR}/.."
+CSRC_DIR="${REPO_ROOT}/csrc"
+
+# clang-format executable: prefer the CLANG_FORMAT env var, otherwise look it up on PATH / in the pip package
+find_clang_format() {
+    if [ -n "${CLANG_FORMAT:-}" ]; then
+        echo "$CLANG_FORMAT"
+        return
+    fi
+    if command -v clang-format >/dev/null 2>&1; then
+        command -v clang-format
+        return
+    fi
+    # the binary shipped by pip install clang-format (not necessarily on PATH)
+    local pip_bin
+    pip_bin=$(python -c "import clang_format, os; print(os.path.join(os.path.dirname(clang_format.__file__), 'data', 'bin', 'clang-format'))" 2>/dev/null || true)
+    if [ -n "$pip_bin" ] && [ -x "$pip_bin" ]; then
+        echo "$pip_bin"
+        return
+    fi
+    echo "ERROR: clang-format not found. Install it or set CLANG_FORMAT=<path>." >&2
+    exit 1
+}
+
+CLANG_FORMAT_BIN=$(find_clang_format)
+
+# file extensions
+EXTENSIONS=("*.cpp" "*.cc" "*.h" "*.hpp" "*.cu" "*.cuh")
+
+# build the -name arguments for find
+NAME_ARGS=()
+for i in "${!EXTENSIONS[@]}"; do
+    if [ "$i" -gt 0 ]; then
+        NAME_ARGS+=("-o")
+    fi
+    NAME_ARGS+=("-name" "${EXTENSIONS[$i]}")
+done
+
+FILES=$(find "$CSRC_DIR" \( "${NAME_ARGS[@]}" \))
+
+if [ -z "$FILES" ]; then
+    echo "No files found to format."
+    exit 0
+fi
+
+FILE_COUNT=$(echo "$FILES" | wc -l)
+
+# mode: --check only checks and does not modify
+if [ "${1:-}" = "--check" ]; then
+    echo "Checking clang-format for $FILE_COUNT files in csrc/..."
+    FAILED=0
+    while IFS= read -r file; do
+        if ! "$CLANG_FORMAT_BIN" --dry-run --Werror "$file" 2>/dev/null; then
+            echo "  FAIL: $file"
+            FAILED=$((FAILED + 1))
+        fi
+    done <<< "$FILES"
+    if [ "$FAILED" -gt 0 ]; then
+        echo "$FAILED file(s) need formatting."
+        exit 1
+    else
+        echo "All files are correctly formatted."
+        exit 0
+    fi
+else
+    echo "Formatting $FILE_COUNT files in csrc/..."
+    echo "$FILES" | xargs "$CLANG_FORMAT_BIN" -i
+    echo "Done."
+fi

scripts/gen_stub.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+
+set -ex
+
+(
+    CUDA_STUBS_DIR=$(python -c 'from torch.utils.cpp_extension import CUDA_HOME; print(f"{CUDA_HOME}/lib64/stubs")')
+    CUDA_STUB_LIBRARY="${CUDA_STUBS_DIR}/libcuda.so"
+    if [ -e "${CUDA_STUB_LIBRARY}" ]; then
+        export LD_PRELOAD="${CUDA_STUB_LIBRARY}"
+    fi
+    python -c 'import torch
+from pybind11_stubgen import main
+main(["deepinfer._C", "-o", "."])
+'
+)
+
+# engine typing.Sequence -> Sequence
+sed -i 's/typing.Sequence/Sequence/g' deepinfer/_C/engine.pyi