diff --git a/.github/workflows/package_wheel_release.yml b/.github/workflows/package_wheel_release.yml
new file mode 100644
index 0000000..93e5f38
--- /dev/null
+++ b/.github/workflows/package_wheel_release.yml
@@ -0,0 +1,252 @@
+name: Build Wheels
+on:
+ workflow_dispatch:
+ inputs:
+ release:
+ description: 'Release? 1 = yes, 0 = no'
+ default: '0'
+ required: true
+ type: string
+jobs:
+ build_wheels:
+ name: ${{ matrix.os }} Python=${{ matrix.pyver }} CUDA=${{ matrix.cuda }} CPU_INSTRUCT=${{ matrix.instruct }} Torch=${{ matrix.torch }}
+ runs-on: ${{ matrix.os }}
+ strategy:
+ fail-fast: false
+ matrix:
+ include:
+ # Ubuntu
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.5.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.5.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.5.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.4.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.4.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.4.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.2.2', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.2.2', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.2.2', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.1.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.1.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.1.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.5.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.5.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.5.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.4.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.4.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.4.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.2.2', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.2.2', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.2.2', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.1.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.1.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.1.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.5.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.5.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.5.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.4.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.4.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.4.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.2.2', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.2.2', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.2.2', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.1.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.1.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.1.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.5.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.5.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.5.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.4.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.4.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.4.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.2.2', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.2.2', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.2.2', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.1.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.1.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.11', cuda: '12.1.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.5.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.5.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.5.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.4.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.4.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.4.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.2.2', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.2.2', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.2.2', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.1.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.1.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.1.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.5.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.5.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.5.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.4.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.4.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.4.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.2.2', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.2.2', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.2.2', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.1.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.1.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: ubuntu-20.04, pyver: '3.10', cuda: '12.1.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+
+ # Windows
+ - { os: windows-2022, pyver: '3.12', cuda: '12.5.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.12', cuda: '12.5.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.12', cuda: '12.4.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.12', cuda: '12.4.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.12', cuda: '12.2.2', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.12', cuda: '12.2.2', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.12', cuda: '12.1.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.12', cuda: '12.1.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.12', cuda: '12.5.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.12', cuda: '12.5.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.12', cuda: '12.4.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.12', cuda: '12.4.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.12', cuda: '12.2.2', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.12', cuda: '12.2.2', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.12', cuda: '12.1.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.12', cuda: '12.1.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.5.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.5.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.4.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.4.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.2.2', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.2.2', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.1.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.1.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.5.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.5.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.4.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.4.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.2.2', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.2.2', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.1.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.1.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.10', cuda: '12.5.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.10', cuda: '12.5.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.10', cuda: '12.4.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.10', cuda: '12.4.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.10', cuda: '12.2.2', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.10', cuda: '12.2.2', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.10', cuda: '12.1.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.10', cuda: '12.1.1', torch: '2.4.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.10', cuda: '12.5.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.10', cuda: '12.5.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.10', cuda: '12.4.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.10', cuda: '12.4.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+ - { os: windows-2022, pyver: '3.10', cuda: '12.2.2', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.10', cuda: '12.2.2', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.10', cuda: '12.1.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX512', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.10', cuda: '12.1.1', torch: '2.3.0', cudaarch: '8.0;8.6;8.7;8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '121'}
+
+ defaults:
+ run:
+ shell: pwsh
+
+ steps:
+ - uses: actions/checkout@v3
+
+ - name: Free Disk Space
+ uses: jlumbroso/free-disk-space@v1.3.1
+ if: runner.os == 'Linux'
+ with:
+ tool-cache: true
+ android: true
+ dotnet: true
+ haskell: true
+ large-packages: false
+ swap-storage: true
+
+ - uses: actions/setup-python@v4
+ with:
+ python-version: ${{ matrix.pyver }}
+
+ - name: check_space
+ run: |
+ if($IsLinux) {df -h}
+ if($IsWindows) {Get-PSDrive -PSProvider 'FileSystem'}
+
+ - uses: actions/setup-node@v4
+ with:
+ node-version: 20
+
+ - name: Setup Mamba
+ if: matrix.cuda != ''
+ uses: conda-incubator/setup-miniconda@v2.3.0
+ with:
+ activate-environment: "ktransformers"
+ python-version: ${{ matrix.pyver }}
+ miniforge-variant: Mambaforge
+ miniforge-version: latest
+ use-mamba: true
+ add-pip-as-python-dependency: true
+ auto-activate-base: false
+
+
+
+ - name: build web
+ run: |
+ cd ktransformers/website/
+ npm install
+ npm run build
+ cd ../../
+
+ - name: build for cuda
+ if: matrix.cuda != ''
+ run: |
+ git submodule init
+ git submodule update
+ if($IsWindows){
+ $originalPath = Get-Location
+ Import-Module 'C:\Program Files\Microsoft Visual Studio\2022\Enterprise\Common7\Tools\Microsoft.VisualStudio.DevShell.dll'
+ Enter-VsDevShell -VsInstallPath 'C:\Program Files\Microsoft Visual Studio\2022\Enterprise' -DevCmdArguments '-arch=x64 -host_arch=x64'
+ $env:DISTUTILS_USE_SDK=1
+ Set-Location $originalPath
+ }
+ $cudaVersion = '${{ matrix.cuda }}'
+ $env:MAMBA_NO_LOW_SPEED_LIMIT = 1
+ mamba install -y -c nvidia/label/cuda-$cudaVersion cuda-toolkit cuda-runtime
+ $env:CUDA_PATH = $env:CONDA_PREFIX
+ $env:CUDA_HOME = $env:CONDA_PREFIX
+ if ($IsLinux) {
+ $env:LD_LIBRARY_PATH = $env:CONDA_PREFIX + '/lib:' + $env:LD_LIBRARY_PATH
+ $env:LD_LIBRARY_PATH = $env:CONDA_PREFIX + '/lib/python${{ matrix.pyver }}/site-packages/nvidia/nvjitlink/lib:' + $env:LD_LIBRARY_PATH
+ if (!(Test-Path $env:CUDA_HOME/lib64)) {
+ New-Item -ItemType SymbolicLink -Path $env:CUDA_HOME/lib64 -Target $env:CUDA_HOME/lib
+ }
+ }
+ if ($IsWindows) {
+ $env:CUDA_PATH = "$env:CUDA_PATH/Library"
+ $env:CUDA_HOME = $env:CUDA_PATH
+ $env:PATH = "$env:CUDA_PATH/bin;" + $env:PATH
+ cp $env:CUDA_PATH/lib/*.lib $env:CUDA_PATH/lib/x64/
+ $env:INCLUDE =$env:CUDA_PATH + "/include/targets/x64;" + $env:INCLUDE
+
+ }
+ python -m pip install torch==${{ matrix.torch }} torchvision torchaudio --index-url https://download.pytorch.org/whl/cu${{ matrix.torch_cu }}
+ python -m pip install cpufeature build wheel ninja packaging setuptools
+ $env:KTRANSFORMERS_FORCE_BUILD = "TRUE"
+ $env:CPU_INSTRUCT = '${{ matrix.instruct }}'
+ $env:TORCH_CUDA_ARCH_LIST = '${{ matrix.cudaarch }}'
+ python -m build --no-isolation --verbose
+
+
+ - name: create Rlease dir
+ run: |
+ if ($IsWindows) {
+ $env:date = $(Get-Date -Format "yyyy-MM-dd")
+ New-Item -ItemType Directory -Force -Path "$Env:USERPROFILE\.ssh"
+ $Env:SSH_PATH = "$Env:USERPROFILE\.ssh\id_rsa"
+ Set-Content -Path $Env:SSH_PATH -Value "${{ secrets.SSH_PRIVATE_KEY }}"
+ (Get-Content -Path $Env:SSH_PATH).Replace("`r`n","`n") | Set-Content -Path $Env:SSH_PATH
+ chmod 600 $Env:SSH_PATH
+ }
+ if ($IsLinux) {
+ $env:date = $(date +%Y-%m-%d)
+ mkdir -p ~/.ssh/
+ echo "${{ secrets.SSH_PRIVATE_KEY }}" > ~/.ssh/id_rsa
+ chmod 600 ~/.ssh/id_rsa
+ }
+
+ ssh -p ${{ secrets.SSH_PORT }} -o StrictHostKeyChecking=no root@${{ secrets.SSH_SERVER }} "mkdir -p /mnt/data/release-$env:date"
+ scp -P ${{ secrets.SSH_PORT }} -o StrictHostKeyChecking=no dist/*.whl root@${{ secrets.SSH_SERVER }}:/mnt/data/release-$env:date/
\ No newline at end of file
diff --git a/.github/workflows/package_wheel_test.yml b/.github/workflows/package_wheel_test.yml
new file mode 100644
index 0000000..9fe82f8
--- /dev/null
+++ b/.github/workflows/package_wheel_test.yml
@@ -0,0 +1,132 @@
+name: Build Wheels
+on:
+ workflow_dispatch:
+ inputs:
+ release:
+ description: 'Release? 1 = yes, 0 = no'
+ default: '0'
+ required: true
+ type: string
+jobs:
+ build_wheels:
+ name: ${{ matrix.os }} Python=${{ matrix.pyver }} CUDA=${{ matrix.cuda }} CPU_INSTRUCT=${{ matrix.instruct }} Torch=${{ matrix.torch }}
+ runs-on: ${{ matrix.os }}
+ strategy:
+ fail-fast: false
+ matrix:
+ include:
+ # Ubuntu
+ - { os: ubuntu-20.04, pyver: '3.12', cuda: '12.2.2', torch: '2.3.0', cudaarch: '8.9;9.0+PTX', instruct: 'FANCY', torch_cu: '121'}
+ - { os: windows-2022, pyver: '3.11', cuda: '12.5.1', torch: '2.4.0', cudaarch: '8.9;9.0+PTX', instruct: 'AVX2', torch_cu: '124'}
+
+ defaults:
+ run:
+ shell: pwsh
+
+ steps:
+ - uses: actions/checkout@v3
+
+ - name: Free Disk Space
+ uses: jlumbroso/free-disk-space@v1.3.1
+ if: runner.os == 'Linux'
+ with:
+ tool-cache: true
+ android: true
+ dotnet: true
+ haskell: true
+ large-packages: false
+ swap-storage: true
+
+ - uses: actions/setup-python@v4
+ with:
+ python-version: ${{ matrix.pyver }}
+
+ - name: check_space
+ run: |
+ if($IsLinux) {df -h}
+ if($IsWindows) {Get-PSDrive -PSProvider 'FileSystem'}
+
+ - uses: actions/setup-node@v4
+ with:
+ node-version: 20
+
+ - name: Setup Mamba
+ if: matrix.cuda != ''
+ uses: conda-incubator/setup-miniconda@v2.3.0
+ with:
+ activate-environment: "ktransformers"
+ python-version: ${{ matrix.pyver }}
+ miniforge-variant: Mambaforge
+ miniforge-version: latest
+ use-mamba: true
+ add-pip-as-python-dependency: true
+ auto-activate-base: false
+
+
+
+ - name: build web
+ run: |
+ cd ktransformers/website/
+ npm install
+ npm run build
+ cd ../../
+
+ - name: build for cuda
+ if: matrix.cuda != ''
+ run: |
+ git submodule init
+ git submodule update
+ if($IsWindows){
+ $originalPath = Get-Location
+ Import-Module 'C:\Program Files\Microsoft Visual Studio\2022\Enterprise\Common7\Tools\Microsoft.VisualStudio.DevShell.dll'
+ Enter-VsDevShell -VsInstallPath 'C:\Program Files\Microsoft Visual Studio\2022\Enterprise' -DevCmdArguments '-arch=x64 -host_arch=x64'
+ $env:DISTUTILS_USE_SDK=1
+ Set-Location $originalPath
+ }
+ $cudaVersion = '${{ matrix.cuda }}'
+ $env:MAMBA_NO_LOW_SPEED_LIMIT = 1
+ mamba install -y -c nvidia/label/cuda-$cudaVersion cuda-toolkit cuda-runtime
+ $env:CUDA_PATH = $env:CONDA_PREFIX
+ $env:CUDA_HOME = $env:CONDA_PREFIX
+ if ($IsLinux) {
+ $env:LD_LIBRARY_PATH = $env:CONDA_PREFIX + '/lib:' + $env:LD_LIBRARY_PATH
+ $env:LD_LIBRARY_PATH = $env:CONDA_PREFIX + '/lib/python${{ matrix.pyver }}/site-packages/nvidia/nvjitlink/lib:' + $env:LD_LIBRARY_PATH
+ if (!(Test-Path $env:CUDA_HOME/lib64)) {
+ New-Item -ItemType SymbolicLink -Path $env:CUDA_HOME/lib64 -Target $env:CUDA_HOME/lib
+ }
+ }
+ if ($IsWindows) {
+ $env:CUDA_PATH = "$env:CUDA_PATH/Library"
+ $env:CUDA_HOME = $env:CUDA_PATH
+ $env:PATH = "$env:CUDA_PATH/bin;" + $env:PATH
+ cp $env:CUDA_PATH/lib/*.lib $env:CUDA_PATH/lib/x64/
+ $env:INCLUDE =$env:CUDA_PATH + "/include/targets/x64;" + $env:INCLUDE
+
+ }
+ python -m pip install torch==${{ matrix.torch }} torchvision torchaudio --index-url https://download.pytorch.org/whl/cu${{ matrix.torch_cu }}
+ python -m pip install cpufeature build wheel ninja packaging setuptools
+ $env:KTRANSFORMERS_FORCE_BUILD = "TRUE"
+ $env:CPU_INSTRUCT = '${{ matrix.instruct }}'
+ $env:TORCH_CUDA_ARCH_LIST = '${{ matrix.cudaarch }}'
+ python -m build --no-isolation --verbose
+
+
+ - name: create Rlease dir
+ run: |
+ if ($IsWindows) {
+ $env:date = $(Get-Date -Format "yyyy-MM-dd")
+ New-Item -ItemType Directory -Force -Path "$Env:USERPROFILE\.ssh"
+ $Env:SSH_PATH = "$Env:USERPROFILE\.ssh\id_rsa"
+ Set-Content -Path $Env:SSH_PATH -Value "${{ secrets.SSH_PRIVATE_KEY }}"
+ (Get-Content -Path $Env:SSH_PATH).Replace("`r`n","`n") | Set-Content -Path $Env:SSH_PATH
+ chmod 600 $Env:SSH_PATH
+ }
+ if ($IsLinux) {
+ $env:date = $(date +%Y-%m-%d)
+ mkdir -p ~/.ssh/
+ echo "${{ secrets.SSH_PRIVATE_KEY }}" > ~/.ssh/id_rsa
+ chmod 600 ~/.ssh/id_rsa
+ }
+
+ ssh -p ${{ secrets.SSH_PORT }} -o StrictHostKeyChecking=no root@${{ secrets.SSH_SERVER }} "mkdir -p /mnt/data/release-$env:date"
+ scp -P ${{ secrets.SSH_PORT }} -o StrictHostKeyChecking=no dist/*.whl root@${{ secrets.SSH_SERVER }}:/mnt/data/release-$env:date/
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
index 718ea55..1bb8666 100644
--- a/.gitignore
+++ b/.gitignore
@@ -14,4 +14,7 @@ node_modules
.DS_Store
compile_commands.json
*.egg-info*
-*dist/
\ No newline at end of file
+*dist/
+ktransformers/server/local_store/
+ktransformers/server_test1.db
+*.patch
\ No newline at end of file
diff --git a/README.md b/README.md
index a80fe67..8c5f505 100644
--- a/README.md
+++ b/README.md
@@ -268,7 +268,10 @@ In this example, the AutoModel is first initialized on the meta device to avoid
After injection, the original `generate` interface is available, but we also provide a compatible `prefill_and_generate` method, which enables further optimizations like CUDAGraph to improve generation speed.
-YAML Template
+How to custom your model
+
+A detailed tutorial of the injection and multi-GPU using DeepSeek-V2 as an example is given [here](doc/en/injection_tutorial.md).
+
Below is an example of a YAML template for replacing all original Linear modules with Marlin, an advanced 4-bit quantization kernel.
```yaml
@@ -287,7 +290,7 @@ Each rule in the YAML file has two parts: `match` and `replace`. The `match` par
You can find example rule templates for optimizing DeepSeek-V2 and Qwen2-57B-A14, two SOTA MoE models, in the [ktransformers/optimize/optimize_rules](ktransformers/optimize/optimize_rules) directory. These templates are used to power the `local_chat.py` demo.
-A detailed description of the injection using DeepSeek-V2 as an example is given [here](doc/en/deepseek-v2-injection.md).
+If you are interested in our design principles and the implementation of the injection framework, please refer to the [design document](doc/en/deepseek-v2-injection.md).
Acknowledgment and Contributors
diff --git a/doc/assets/deepseekv2_structure.png b/doc/assets/deepseekv2_structure.png
new file mode 100644
index 0000000..b9ced32
Binary files /dev/null and b/doc/assets/deepseekv2_structure.png differ
diff --git a/doc/assets/model_structure_guild.png b/doc/assets/model_structure_guild.png
new file mode 100644
index 0000000..d9f8e4a
Binary files /dev/null and b/doc/assets/model_structure_guild.png differ
diff --git a/doc/assets/multi_gpu.png b/doc/assets/multi_gpu.png
new file mode 100644
index 0000000..88c025b
Binary files /dev/null and b/doc/assets/multi_gpu.png differ
diff --git a/doc/en/deepseek-v2-injection.md b/doc/en/deepseek-v2-injection.md
index c1ccd39..e5dc1c2 100644
--- a/doc/en/deepseek-v2-injection.md
+++ b/doc/en/deepseek-v2-injection.md
@@ -90,7 +90,7 @@ The YAML rule is listed below.
- match:
name: "^model\\.layers\\..*\\.self_attn$" # regular expression
replace:
- class: ktransformers.operators.attention.DeepseekV2AttentionInjected # optimized MLA implementation
+ class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
```
As we can see, each rule in the YAML file has two parts: `match` and `replace`.
@@ -98,9 +98,9 @@ The match part specifies which module should be replaced, and the replace part s
Routed Experts
-For routed experts, the module we inject is a wrapper of CPUInfer, KTransformersMLPExpert. There are several implementations within a wrapper, and we need to specify keywords to tell the wrapper which implementation we want to use and how we intend to use it.
+For routed experts, the module we inject is a wrapper of CPUInfer, KTransformersExperts. There are several implementations within a wrapper, and we need to specify keywords to tell the wrapper which implementation we want to use and how we intend to use it.
-In KTransformers, some models exhibit different behaviors during prefilling and generation for better performance. KTransformersMLPExpert is one of them. All these special modules have a `device` keyword describing which device the module should be initialized on. Other keywords specify the behaviors during prefilling and generation and may be differ when using different injection modules. Here, we specify which implementation on which device we want to use during prefilling and generation, and which device the output should be on.
+In KTransformers, some models exhibit different behaviors during prefilling and generation for better performance. KTransformersExperts is one of them. All these special modules have a `device` keyword describing which device the module should be initialized on. Other keywords specify the behaviors during prefilling and generation and may be differ when using different injection modules. Here, we specify which implementation on which device we want to use during prefilling and generation, and which device the output should be on.
Note that we only use these parameters when layer-wise prefilling is enabled; otherwise, prefilling is conducted with the same configuration as generation.
In the original implementation of Transformers, MoE is implemented using `nn.ModuleList`. We don't want KTransformers to iterate through all the sub-modules in the list, so we set `recursive: False` in this rule to prevent recursive injection into submodules of the current module. Here is the YAML rule:
@@ -109,13 +109,13 @@ In the original implementation of Transformers, MoE is implemented using `nn.Mod
- match:
name: "^model\\.layers\\..*\\.mlp\\.experts$"
replace:
- class: ktransformers.operators.experts.KTransformersMLPExpert # custom MoE Kernel with expert parallelism
+ class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert parallelism
device: "cpu" # device to load this module on initialization
kwargs:
prefill_device: "cuda"
- prefill_mlp_type: "MLPExpertsTorch"
+ prefill_op: "KExpertsTorch"
generate_device: "cpu"
- generate_mlp_type: "MLPCPUExperts"
+ generate_op: "KExpertsCPU"
out_device: "cuda"
recursive: False # don't recursively inject submodules of this module
```
@@ -126,7 +126,7 @@ If we inject the expert list as a custom module, we can't use the interface in `
- match:
class: ktransformers.models.modeling_deepseek.DeepseekV2MoE
replace:
- class: ktransformers.operators.experts.DeepseekV2MoEInjected # MLP module with custom forward function
+ class: ktransformers.operators.experts.KDeepseekV2MoE # MLP module with custom forward function
```
Other Linear Modules
@@ -140,12 +140,12 @@ We also need to transfer some keywords similar to the injection of experts. Here
name: "^model\\.layers\\.(?!.*self_attn).*$" # regular expression
class: torch.nn.Linear # only match modules matching name and class simultaneously
replace:
- class: ktransformers.operators.linear.KTransformerLinear # optimized Kernel on quantized data types
+ class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types
kwargs:
generate_device: "cuda"
prefill_device: "cuda"
- generate_op: "QuantizedLinearMarlin"
- prefill_op: "QuantizedLinearTorch"
+ generate_op: "KLinearMarlin"
+ prefill_op: "KLinearTorch"
```
Pre-compute Buffers
diff --git a/doc/en/injection_tutorial.md b/doc/en/injection_tutorial.md
new file mode 100644
index 0000000..655163e
--- /dev/null
+++ b/doc/en/injection_tutorial.md
@@ -0,0 +1,328 @@
+# Tutorial: Inject Operator Step by Step
+
+> Author: Azure-Tang
+
+## TL;DR
+This tutorial will guide you through the process of injecting custom operators into a model using the KTransformers framework. We will use the DeepSeekV2-Chat model as an example to demonstrate how to inject custom operators into the model step by step. The tutorial will cover the following topics:
+* [How to write injection rules](#how-to-write-injection-rules)
+ * [Understanding the structure of the model](#understanding-model-structure)
+* [Multi-GPU](#muti-gpu)
+* [How to write a new operator and inject it into the model](#how-to-write-a-new-operator-and-inject-into-the-model)
+
+## How to Write Injection Rules
+The basic form of the injection rules for the Inject framework is as follows:
+```yaml
+- match:
+ name: "^model\\.layers\\..*\\.*$" # Target module name
+ class: torch.nn.Linear # Target module
+ replace:
+ class: "default"
+ kwargs:
+ generate_device: "cuda:0"
+ # your_op_param_1: 1234
+ # your_op_param_2: 5678
+ recursive: True
+```
+* match: This field marks the matching rules, which can appear in two forms, name and class. These two matching rules can appear together or separately; they only match when both criteria are met.
+* replace:
+ * class: Python class that can be imported to replace the target module. If no replacement is desired, set to default.
+ * kwargs: List of parameters needed for module initialization.
+ * generate_device: The device for this module, can be set to “cpu”, “cuda”, “cuda:1”, etc.
+* recursive: Whether to recursively inject this module’s submodules, default is True.
+
+For the recursive field: Some modules contain multiple submodules, such as the Self-attention module typically includes q/k/v/o four linear modules. If we replace the self-attention module but do not want the internal linear modules to be covered by other rules, set this rule to False.
+
+## Understanding Model Structure
+Using [deepseek-ai/DeepSeek-V2-Lite-Chat](https://huggingface.co/deepseek-ai/DeepSeek-V2-Lite-Chat) as an example, we can follow the above rules step by step to inject our custom module and run it. KTransformers offers a high degree of flexibility, allowing you to replace/experiment with basic operators. However, it also requires users to clearly understand the structure of the model they are running.
+
+Fortunately, knowing the structure of a model is very simple. Open the file list on the [deepseek-ai/DeepSeek-V2-Lite](https://huggingface.co/deepseek-ai/DeepSeek-V2-Lite-Chat/tree/main) homepage, and you can see the following files:
+
+
+ 
+
+
+
+From the `.saftensors` file, we can see the name of each layer’s weights, corresponding to the match.name attribute in the injection rules.
+From the `modeling_deepseek.py` file, we can see the specific implementation of each module class, with the class name corresponding to the match.class attribute in the injection rules.
+
+The structure of the DeepSeekV2 model from the `.saftensors` and `modeling_deepseek.py` files is as follows:
+
+
+ 
+
+
+
+Supported operators and their corresponding classes are as follows:
+
+| match | replace | backends | descriptions |
+| --------- | ---------------------- | ----------------------- | -------------------- |
+| Linear | KTransformersLinear | KLinearMarlin | Marlin as backend |
+| | | KLinearTorch | pytorch as backend |
+| | | KLinearCPUInfer | llamafile as backend |
+| experts | KTransformersExperts | KExpertsTorch | pytorch as backend |
+| | | KExpertsMarlin | Marlin as backend |
+| | | KExpertsCPU | llamafile as backend |
+| Attention | KDeepseekV2Attention | KDeepseekV2Attention | MLA implementation |
+| MoE | KMistralSparseMoEBlock | KQwen2MoeSparseMoeBlock | MoE for Qwen2 |
+| | KDeepseekV2MoE | KDeepseekV2MoE | MoE for DeepseekV2 |
+| Model | KQwen2MoeModel | KQwen2MoeModel | Model for Qwen2 |
+| | KDeepseekV2Model | KDeepseekV2Model | Model for DeepseekV2 |
+| RoPE | RotaryEmbedding | RotaryEmbedding | RoPE module |
+| | YarnRotaryEmbedding | YarnRotaryEmbedding | RoPE module |
+
+Then we start step-by-step injection of custom modules, our targets are:
+
+* Replace the linear module with custom Marlin linear module.
+* Replace the self-attention module with a custom Absorption-based MLA module.
+* Replace the experts module with a custom Experts module.
+* Replace the MoE module with a custom MoE module.
+* Replace the RoPE module with a custom RoPE module.
+* Set the running device for each module.
+
+The full implementation of the injection rules can be found in the [here](https://github.com/kvcache-ai/ktransformers/blob/main/ktransformers/optimize/optimize_rules/DeepSeek-V2-Chat.yaml).
+
+## Matrix Absorption-based MLA Injection
+
+For the injection of the Attention module, we only need to use a regular expression to match the module names used in transformers and replace them with our own MLA module implementation. The YAML injection rule is as follows:
+```yaml
+- match:
+ name: "^model\\.layers\\..*\\.self_attn$" # Regular expression
+ replace:
+ class: ktransformers.operators.attention.KDeepseekV2Attention # Optimized MLA implementation
+```
+As you can see, each rule in the YAML file has two parts: match and replace. The match part specifies the module to be replaced, and the replace part specifies the module to be injected into the model along with the initialization keywords.
+
+## Injection of Routed Experts
+For Routed Experts (corresponding to the exps in the diagram), the module we inject is CPUInfer, which is wrapped in the wrapper module KTransformersExperts. KTransformersExperts has multiple implementations, and we need to specify keywords to tell the wrapper module which implementation we want to use and how we plan to use it.
+
+In the source code of the transformer, MoE is implemented using nn.ModuleList. We do not want KTransformers to traverse all submodules in the list and inject them one by one, so in this rule, we set recursive: False to prevent recursive injection into the submodules of this module. The YAML rule is as follows:
+
+```yaml
+- match:
+ name: "^model\\.layers\\..*\\.mlp\\.experts$"
+ replace:
+ class: ktransformers.operators.experts.KTransformersExperts # Custom MoE kernel with expert parallelism
+ kwargs:
+ generate_device: "cpu"
+ generate_op: "MLPCPUExperts"
+ out_device: "cuda"
+ recursive: False # Don't recursively inject submodules of this module
+```
+
+If we inject Routed Experts as a custom module, we cannot use the interfaces in the original `nn.ModuleList`. Therefore, it is necessary to modify the forward function in the FFN module. The simplest method is to implement a new module with a custom forward function and inject it.
+```yaml
+- match:
+ class: ktransformers.models.modeling_deepseek.DeepseekV2MoE
+ replace:
+ class: ktransformers.operators.experts.KDeepseekV2MoE # MLP module with custom forward function
+```
+
+## Injection of Linear Layers
+
+For the remaining linear layer modules, we aim to use quantized operators to save storage space while improving performance. Since there is no current research on using MLA and quantization together, we do not want to inject linear into the MLA operator. Therefore, we can modify the regular expression and add a type check in the match part of the rule. Only modules that match both the name and class simultaneously will be injected. We also need to pass some keywords similar to the injection of Routed Experts. The YAML rule is as follows:
+
+```yaml
+- match:
+ name: "^model\\.layers\\.(?!.*self_attn).*$" # Regular expression
+ class: torch.nn.Linear # Only match modules matching name and class simultaneously
+ replace:
+ class: ktransformers.operators.linear.KTransformersLinear # Optimized kernel on quantized data types
+ kwargs:
+ generate_device: "cuda"
+ generate_op: "QuantizedLinearMarlin"
+```
+## Injection of Modules with Pre-calculated Buffers
+
+To avoid occupying resources when initializing the injected original model, we use torch’s meta device to initialize the original model. The RoPE module pre-calculates some buffers during initialization, but no calculations are performed when using the meta device. Therefore, we need to compensate for the calculation of the buffer when loading the model. Simply, we inject a custom module into the rotary embedding module, which performs pre-calculation during loading. The YAML rule is as follows:
+```yaml
+- match:
+ class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding
+ replace:
+ class: ktransformers.operators.RoPE.YarnRotaryEmbedding
+```
+
+## Specifying Running Devices for Modules
+
+Finally, we set a fallback basic attribute generate_device for all modules:
+```yaml
+- match:
+ name: "^model\\.layers\\..*\\.|^lm_head"
+ replace:
+ class: "default"
+ kwargs:
+ generate_device: "cuda"
+
+- match:
+ name: "^model.embed_tokens"
+ replace:
+ class: "default"
+ kwargs:
+ generate_device: "cpu"
+```
+Through these two rules, we place all previously unmatched layers (and their submodules) and lm_head on cuda, and the embedding on cpu. Note that the properties of a module will be determined by the first rule it matches. For example, if you later set a new replace.kwargs.generate_device in an injected module, the device set earlier will take precedence. If your computer has multiple cards, you can also configure the model to multiple cards.
+
+
+## Muti-GPU
+
+If you have multiple GPUs, you can set the device for each module to different GPUs.
+DeepseekV2-Chat got 60 layers, if we got 2 GPUs, we can allocate 30 layers to each GPU. Complete multi GPU rule examples [here](ktransformers/optimize/optimize_rules).
+
+
+
+
+ 
+
+
+
+First of all, for multi-GPU, we have to inject an new operator `KDeepseekV2Model`. And set division of the layers to different GPUs. For our case, we have to set the `transfer_map` in the `KDeepseekV2Model` operatoras as follows:
+
+```yaml
+- match:
+ name: "^model$"
+ replace:
+ class: "ktransformers.operators.models.KDeepseekV2Model"
+ kwargs:
+ transfer_map:
+ 30: "cuda:1"
+```
+
+And we have to set the device for each module in the model.
+
+For example, for `routed experts`, the yaml for one GPU is:
+```yaml
+- match:
+ name: "^model\\.layers\\..*\\.mlp\\.experts$"
+ replace:
+ class: ktransformers.operators.experts.KTransformersExperts # Custom MoE kernel with expert parallelism
+ kwargs:
+ generate_device: "cuda:0"
+ generate_op: "MLPCUDAExperts"
+ out_device: "cuda:0"
+ recursive: False # Don't recursively inject submodules of this module
+```
+But for two GPUs, we need to set the device for each module in the model.
+
+```yaml
+# allcate 0-29 layers‘s out_device to cuda:0
+- match:
+ name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.experts$"
+ replace:
+ class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism
+ kwargs:
+ generate_device: "cpu"
+ generate_op: "KExpertsCPU"
+ out_device: "cuda:0"
+ recursive: False # don't recursively inject submodules of this module
+
+# allocate 30-59 layers‘s out_device to cuda:1
+- match:
+ name: "^model\\.layers\\.([345][0-9])\\.mlp\\.experts$"
+ replace:
+ class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism
+ kwargs:
+ generate_device: "cpu"
+ generate_op: "KExpertsCPU"
+ out_device: "cuda:1"
+ recursive: False # don't recursively inject submodules of this module
+```
+For other modules, we can set the device in the same way.
+
+## How to Write a New Operator and Inject into the Model
+
+In this section, we will explain how to write an operator that can be injected, using the implementation of a new linear as an example.
+
+First, all injectable operators need to inherit from the BaseInjectedModule class, which inherits some attributes required by our injection framework. Its initialization function needs to meet the following basic format:
+
+```python
+class LinearTorchInject(BaseInjectedModule):
+ def __init__(
+ self,
+ key: str,
+ gguf_loader: GGUFLoader,
+ config: PretrainedConfig,
+ orig_module: nn.Module = None,
+ generate_device: str = "cuda",
+ **kwargs,
+ ):
+ super().__init__(key, gguf_loader, config, orig_module, generate_device, **kwargs)
+```
+If users have other parameters that need to be passed to this class, they can also be included in the init function and re-passed in the kwargs parameter in the yaml file. For example, if our operator wants to pass a parameter `my_param`, the init function can be written as:
+```python
+class LinearTorchInject(BaseInjectedModule):
+ def __init__(
+ self,
+ key: str,
+ gguf_loader: GGUFLoader,
+ config: PretrainedConfig,
+ orig_module: nn.Module = None,
+ generate_device: str = "cuda",
+ my_param: bool = True,
+ **kwargs,
+ ):
+ super().__init__(key, gguf_loader, config, orig_module, generate_device, **kwargs)
+ self.my_param = my_param
+```
+Then our injection rule can be written as:
+```yaml
+- match:
+ name: "^model\\.layers\\..*$" # Regular expression matches the module name.
+ class: torch.nn.Linear # Type restrictions can be added.
+ replace:
+ class: ktransformers.operators.linear.LinearTorchInject # Inject module path
+ kwargs: # Extra parameters
+ generate_device: "cuda"
+ my_param: True
+```
+For the linear module, it is also necessary to read weights from a gguf file. We provide the `KLinearBase` class to help users read weights from gguf files. Users only need to inherit and implement the load, unload, and forward functions. Therefore, a fully injectable linear class would look like this:
+```python
+class LinearTorchInject(BaseInjectedModule, KLinearBase):
+ def __init__(
+ self,
+ key: str,
+ gguf_loader: GGUFLoader,
+ config: PretrainedConfig,
+ orig_module: nn.Module = None,
+ generate_device: str = "cuda",
+ **kwargs,
+ ):
+ super().__init__(key, gguf_loader, config, orig_module, generate_device, **kwargs)
+ KLinearBase.__init__(self)
+ self.has_bias = False
+ self.dtype = torch.get_default_dtype()
+ self.w = None
+ self.has_bias = False
+
+ def load(self, w: dict | nn.Parameter | tuple | None = None, device: str|None = None):
+ if device is None: device = self.device
+ if w is None: w = self.load_weight(device=device)
+
+ if isinstance(w, nn.Parameter):
+ self.w = w.to(dtype=self.dtype).view(self.out_features, self.in_features).T
+ self.has_bias = False
+ elif isinstance(w, tuple):
+ self.w = w[0].to(dtype=self.dtype).view(self.out_features, self.in_features).T
+ self.bias = w[1].to(dtype=self.dtype)
+ self.has_bias = True
+ else:
+ raise ValueError("Invalid weight type")
+ self.w = self.w.to(device)
+ if self.has_bias:
+ self.bias = self.bias.to(device)
+
+ def unload(self):
+ if self.w is not None:
+ self.w = None
+ if self.has_bias:
+ self.bias = None
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ dtype = x.dtype
+ out_device = x.device
+ x = x.to(device=self.device, dtype=self.dtype)
+ x = x @ self.w
+ if self.has_bias:
+ x = x + self.bias
+ x = x.to(dtype=dtype, device=out_device)
+ return x
+```
+Note that the `self.load_weight` function is provided by the KLinearBase class to help users load weights from a gguf file into the module. The implementation details of KLinearBase can be found on [GITHUB](https://github.com/kvcache-ai/ktransformers/blob/44f57270c9514d79fab224186d90ccf61059331a/ktransformers/operators/linear.py#L31).
diff --git a/ktransformers/__init__.py b/ktransformers/__init__.py
index d1f2e39..48fef32 100644
--- a/ktransformers/__init__.py
+++ b/ktransformers/__init__.py
@@ -1 +1 @@
-__version__ = "0.1.1"
\ No newline at end of file
+__version__ = "0.1.2"
\ No newline at end of file
diff --git a/ktransformers/ktransformers_ext/CMakeLists.txt b/ktransformers/ktransformers_ext/CMakeLists.txt
index 89647a8..e6e0518 100644
--- a/ktransformers/ktransformers_ext/CMakeLists.txt
+++ b/ktransformers/ktransformers_ext/CMakeLists.txt
@@ -22,14 +22,13 @@ option(LLAMA_AVX2 "llama: enable AVX2"
option(LLAMA_AVX512 "llama: enable AVX512" OFF)
option(LLAMA_AVX512_VBMI "llama: enable AVX512-VBMI" OFF)
option(LLAMA_AVX512_VNNI "llama: enable AVX512-VNNI" OFF)
+option(LLAMA_AVX512_BF16 "llama: enable AVX512-BF16" OFF)
option(LLAMA_FMA "llama: enable FMA" OFF)
# in MSVC F16C is implied with AVX2/AVX512
if (NOT MSVC)
option(LLAMA_F16C "llama: enable F16C" OFF)
endif()
option(LLAMA_AVX512_FANCY_SIMD "llama: enable AVX512-VL, AVX512-BW, AVX512-DQ, AVX512-VNNI" OFF)
-option(LLAMA_AVX512_BF16 "llama: enable AVX512-BF16" OFF)
-
# Architecture specific
# TODO: probably these flags need to be tweaked on some architectures
diff --git a/ktransformers/ktransformers_ext/bench/bench_linear.py b/ktransformers/ktransformers_ext/bench/bench_linear.py
index 0a4de3a..3189afd 100644
--- a/ktransformers/ktransformers_ext/bench/bench_linear.py
+++ b/ktransformers/ktransformers_ext/bench/bench_linear.py
@@ -6,7 +6,7 @@
Date : 2024-07-25 10:31:59
Version : 1.0.0
LastEditors : chenht2022
-LastEditTime : 2024-07-25 10:32:51
+LastEditTime : 2024-08-06 10:35:35
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
'''
import os, sys
@@ -15,15 +15,18 @@
import cpuinfer_ext
import torch
+input_size = 16384
+output_size = 5120
+stride = 16
+group_max_len = 1024
+layer_num = 10
+qlen = 1
+CPUInfer = cpuinfer_ext.CPUInfer(64)
+warm_up_iter = 1000
+test_iter = 10000
+
def bench_linear(quant_mode: str):
with torch.inference_mode(mode=True):
- input_size = 16384
- output_size = 5120
- stride = 16
- layer_num = 10
- CPUInfer = cpuinfer_ext.CPUInfer(64)
- warm_up_iter = 1000
- test_iter = 10000
hidden_type = 30 # ggml_type::GGML_TYPE_BF16
if quant_mode == "fp32":
@@ -66,30 +69,37 @@ def bench_linear(quant_mode: str):
projs = []
for _ in range(layer_num):
proj = torch.randn((output_size, input_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
- config = cpuinfer_ext.linear.LinearConfig(input_size, output_size, stride, proj.data_ptr(), proj_type, hidden_type)
+ config = cpuinfer_ext.linear.LinearConfig(input_size, output_size, stride, group_max_len, proj.data_ptr(), proj_type, hidden_type)
linear = cpuinfer_ext.linear.Linear(config)
projs.append(proj)
linears.append(linear)
+ input = torch.randn((layer_num, qlen, input_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
+ output = torch.empty((layer_num, qlen, output_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
# warm up
for i in range(warm_up_iter):
- linear = linears[i % layer_num]
- input = torch.randn((1, input_size), dtype=torch.bfloat16).contiguous()
- output = torch.empty((1, output_size), dtype=torch.bfloat16).contiguous()
- CPUInfer.submit(linear.forward, input.data_ptr(), output.data_ptr())
+ CPUInfer.submit(
+ linears[i % layer_num].forward(
+ qlen,
+ input[i % layer_num].data_ptr(),
+ output[i % layer_num].data_ptr()
+ )
+ )
CPUInfer.sync()
# test
- total_time = 0
+ start = time.perf_counter()
for i in range(test_iter):
- linear = linears[i % layer_num]
- input = torch.randn((1, input_size), dtype=torch.bfloat16).contiguous()
- output = torch.empty((1, output_size), dtype=torch.bfloat16).contiguous()
- start = time.perf_counter()
- CPUInfer.submit(linear.forward, input.data_ptr(), output.data_ptr())
+ CPUInfer.submit(
+ linears[i % layer_num].forward(
+ qlen,
+ input[i % layer_num].data_ptr(),
+ output[i % layer_num].data_ptr()
+ )
+ )
CPUInfer.sync()
- end = time.perf_counter()
- total_time += end - start
+ end = time.perf_counter()
+ total_time = end - start
print('Quant mode: ', quant_mode)
print('Time(s): ', total_time)
print('Iteration: ', test_iter)
diff --git a/ktransformers/ktransformers_ext/bench/bench_linear_torch.py b/ktransformers/ktransformers_ext/bench/bench_linear_torch.py
index cb3e4ef..72e0e75 100644
--- a/ktransformers/ktransformers_ext/bench/bench_linear_torch.py
+++ b/ktransformers/ktransformers_ext/bench/bench_linear_torch.py
@@ -14,14 +14,17 @@
import torch
import torch.nn.quantized as nnq
+scale, zero_point = 0.1, 0 # Adjust scale and zero_point based on your dataset
+
+input_size = 16384
+output_size = 5120
+layer_num = 10
+qlen = 1
+warm_up_iter = 1000
+test_iter = 10000
+
def bench_linear(quant_mode: str):
with torch.inference_mode(mode=True):
- input_size = 16384
- output_size = 5120
- layer_num = 10
- warm_up_iter = 1000
- test_iter = 10000
-
if quant_mode == "fp32":
proj_type = torch.float32
bytes_per_elem = 4.000000
@@ -41,37 +44,32 @@ def bench_linear(quant_mode: str):
for _ in range(layer_num):
proj = torch.randn((output_size, input_size), dtype = torch.float32, device = "cuda").to("cpu").contiguous()
if quant_mode == "qint8":
- scale, zero_point = 0.1, 0 # Adjust scale and zero_point based on your dataset
proj_q = torch.quantize_per_tensor(proj, scale, zero_point, torch.qint8)
quantized_layer = nnq.Linear(input_size, output_size)
quantized_layer.set_weight_bias(proj_q, None)
projs.append(quantized_layer)
else:
projs.append(proj.to(proj_type))
+ input = torch.randn((layer_num, qlen, input_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
# warm up
for i in range(warm_up_iter):
- input = torch.randn((1, input_size), dtype=torch.float32).contiguous()
- if quant_mode == "qint8":
- input_q = torch.quantize_per_tensor(input, scale, zero_point, torch.quint8)
- quantized_layer = projs[i % layer_num]
- t_output = quantized_layer(input_q)
+ if isinstance(projs[i % layer_num], nnq.Linear):
+ input_q = torch.quantize_per_tensor(input[i % layer_num].to(torch.float32), scale, zero_point, torch.quint8)
+ t_output = projs[i % layer_num](input_q)
else:
- t_output = torch.mm(input.to(proj_type), projs[i % layer_num].t())
+ t_output = torch.mm(input[i % layer_num].to(proj_type), projs[i % layer_num].t())
# test
- total_time = 0
+ start = time.perf_counter()
for i in range(test_iter):
- input = torch.randn((1, input_size), dtype=torch.float32).contiguous()
- start = time.perf_counter()
- if quant_mode == "qint8":
- input_q = torch.quantize_per_tensor(input, scale, zero_point, torch.quint8)
- quantized_layer = projs[i % layer_num]
- t_output = quantized_layer(input_q)
+ if isinstance(projs[i % layer_num], nnq.Linear):
+ input_q = torch.quantize_per_tensor(input[i % layer_num].to(torch.float32), scale, zero_point, torch.quint8)
+ t_output = projs[i % layer_num](input_q)
else:
- t_output = torch.mm(input.to(proj_type), projs[i % layer_num].t())
- end = time.perf_counter()
- total_time += end - start
+ t_output = torch.mm(input[i % layer_num].to(proj_type), projs[i % layer_num].t())
+ end = time.perf_counter()
+ total_time = end - start
print('Quant mode: ', quant_mode)
print('Time(s): ', total_time)
print('Iteration: ', test_iter)
diff --git a/ktransformers/ktransformers_ext/bench/bench_mlp.py b/ktransformers/ktransformers_ext/bench/bench_mlp.py
index 5680a9b..690f9d9 100644
--- a/ktransformers/ktransformers_ext/bench/bench_mlp.py
+++ b/ktransformers/ktransformers_ext/bench/bench_mlp.py
@@ -6,7 +6,7 @@
Date : 2024-07-16 10:43:18
Version : 1.0.0
LastEditors : chenht2022
-LastEditTime : 2024-07-25 10:32:55
+LastEditTime : 2024-08-06 10:36:04
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
'''
import os, sys
@@ -15,15 +15,18 @@
import cpuinfer_ext
import torch
+hidden_size = 5120
+intermediate_size = 3072
+stride = 16
+group_max_len = 1024
+layer_num = 10
+qlen = 1
+CPUInfer = cpuinfer_ext.CPUInfer(64)
+warm_up_iter = 1000
+test_iter = 10000
+
def bench_mlp(quant_mode: str):
with torch.inference_mode(mode=True):
- hidden_size = 5120
- intermediate_size = 3072
- stride = 16
- layer_num = 10
- CPUInfer = cpuinfer_ext.CPUInfer(64)
- warm_up_iter = 1000
- test_iter = 10000
hidden_type = 30 # ggml_type::GGML_TYPE_BF16
if quant_mode == "fp32":
@@ -93,32 +96,39 @@ def bench_mlp(quant_mode: str):
gate_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
up_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
down_proj = torch.randn((hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
- config = cpuinfer_ext.mlp.MLPConfig(hidden_size, intermediate_size, stride, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type)
+ config = cpuinfer_ext.mlp.MLPConfig(hidden_size, intermediate_size, stride, group_max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type)
mlp = cpuinfer_ext.mlp.MLP(config)
gate_projs.append(gate_proj)
up_projs.append(up_proj)
down_projs.append(down_proj)
mlps.append(mlp)
+ input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
+ output = torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
# warm up
for i in range(warm_up_iter):
- mlp = mlps[i % layer_num]
- input = torch.randn((1, hidden_size), dtype=torch.bfloat16).contiguous()
- output = torch.empty((1, hidden_size), dtype=torch.bfloat16).contiguous()
- CPUInfer.submit(mlp.forward, input.data_ptr(), output.data_ptr())
+ CPUInfer.submit(
+ mlps[i % layer_num].forward(
+ qlen,
+ input[i % layer_num].data_ptr(),
+ output[i % layer_num].data_ptr()
+ )
+ )
CPUInfer.sync()
# test
- total_time = 0
+ start = time.perf_counter()
for i in range(test_iter):
- mlp = mlps[i % layer_num]
- input = torch.randn((1, hidden_size), dtype=torch.bfloat16).contiguous()
- output = torch.empty((1, hidden_size), dtype=torch.bfloat16).contiguous()
- start = time.perf_counter()
- CPUInfer.submit(mlp.forward, input.data_ptr(), output.data_ptr())
+ CPUInfer.submit(
+ mlps[i % layer_num].forward(
+ qlen,
+ input[i % layer_num].data_ptr(),
+ output[i % layer_num].data_ptr()
+ )
+ )
CPUInfer.sync()
- end = time.perf_counter()
- total_time += end - start
+ end = time.perf_counter()
+ total_time = end - start
print('Quant mode: ', quant_mode)
print('Time(s): ', total_time)
print('Iteration: ', test_iter)
diff --git a/ktransformers/ktransformers_ext/bench/bench_mlp_torch.py b/ktransformers/ktransformers_ext/bench/bench_mlp_torch.py
index 3aad58c..7b811d8 100644
--- a/ktransformers/ktransformers_ext/bench/bench_mlp_torch.py
+++ b/ktransformers/ktransformers_ext/bench/bench_mlp_torch.py
@@ -14,17 +14,38 @@
import torch
import torch.nn.quantized as nnq
+scale, zero_point = 0.1, 0 # Adjust scale and zero_point based on your dataset
+
+hidden_size = 5120
+intermediate_size = 3072
+layer_num = 10
+qlen = 1
+warm_up_iter = 1000
+test_iter = 10000
+
def act_fn(x):
return x / (1.0 + torch.exp(-x))
+def mlp_torch(input, gate_proj, up_proj, down_proj):
+ if isinstance(gate_proj, nnq.Linear):
+ input_q = torch.quantize_per_tensor(input.to(torch.float32), scale, zero_point, torch.quint8)
+ gate_buf = gate_proj(input_q)
+ up_buf = up_proj(input_q)
+ gate_buf = gate_buf.dequantize()
+ up_buf = up_buf.dequantize()
+ intermediate = act_fn(gate_buf) * up_buf
+ intermediate_q = torch.quantize_per_tensor(intermediate, scale, zero_point, torch.quint8)
+ expert_output = down_proj(intermediate_q)
+ ret = expert_output.dequantize()
+ else:
+ gate_buf = torch.mm(input.to(gate_proj.dtype), gate_proj.t())
+ up_buf = torch.mm(input.to(up_proj.dtype), up_proj.t())
+ intermediate = act_fn(gate_buf) * up_buf
+ ret = torch.mm(intermediate.to(down_proj.dtype), down_proj.t())
+ return ret
+
def bench_mlp(quant_mode: str):
with torch.inference_mode(mode=True):
- hidden_size = 5120
- intermediate_size = 3072
- layer_num = 10
- warm_up_iter = 1000
- test_iter = 10000
-
if quant_mode == "fp32":
proj_type = torch.float32
bytes_per_elem = 4.000000
@@ -48,7 +69,6 @@ def bench_mlp(quant_mode: str):
up_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
down_proj = torch.randn((hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
if quant_mode == "qint8":
- scale, zero_point = 0.1, 0 # Adjust scale and zero_point based on your dataset
gate_proj_q = torch.quantize_per_tensor(gate_proj, scale, zero_point, torch.qint8)
quantized_gate = nnq.Linear(hidden_size, intermediate_size)
quantized_gate.set_weight_bias(gate_proj_q, None)
@@ -65,58 +85,18 @@ def bench_mlp(quant_mode: str):
gate_projs.append(gate_proj.to(proj_type))
up_projs.append(up_proj.to(proj_type))
down_projs.append(down_proj.to(proj_type))
+ input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
# warm up
for i in range(warm_up_iter):
- input = torch.randn((1, hidden_size), dtype=torch.float32).contiguous()
- if quant_mode == "qint8":
- input_q = torch.quantize_per_tensor(input, scale, zero_point, torch.quint8)
- quantized_gate = gate_projs[i % layer_num]
- gate_buf = quantized_gate(input_q)
- quantized_up = up_projs[i % layer_num]
- up_buf = quantized_gate(input_q)
- gate_buf = gate_buf.dequantize()
- up_buf = up_buf.dequantize()
- intermediate = act_fn(gate_buf) * up_buf
- intermediate_q = torch.quantize_per_tensor(intermediate, scale, zero_point, torch.quint8)
- quantized_down = down_projs[i % layer_num]
- t_output = quantized_down(intermediate_q)
- else:
- gate_proj = gate_projs[i%layer_num]
- up_proj = up_projs[i%layer_num]
- down_proj = down_projs[i%layer_num]
- gate_buf = torch.mm(input.to(proj_type), gate_proj.t())
- up_buf = torch.mm(input.to(proj_type), up_proj.t())
- intermediate = act_fn(gate_buf) * up_buf
- t_output = torch.mm(intermediate.to(proj_type), down_proj.t())
+ mlp_torch(input[i % layer_num], gate_projs[i % layer_num], up_projs[i % layer_num], down_projs[i % layer_num])
# test
- total_time = 0
+ start = time.perf_counter()
for i in range(test_iter):
- input = torch.randn((1, hidden_size), dtype=torch.float32).contiguous()
- start = time.perf_counter()
- if quant_mode == "qint8":
- input_q = torch.quantize_per_tensor(input, scale, zero_point, torch.quint8)
- quantized_gate = gate_projs[i % layer_num]
- gate_buf = quantized_gate(input_q)
- quantized_up = up_projs[i % layer_num]
- up_buf = quantized_gate(input_q)
- gate_buf = gate_buf.dequantize()
- up_buf = up_buf.dequantize()
- intermediate = act_fn(gate_buf) * up_buf
- intermediate_q = torch.quantize_per_tensor(intermediate, scale, zero_point, torch.quint8)
- quantized_down = down_projs[i % layer_num]
- t_output = quantized_down(intermediate_q)
- else:
- gate_proj = gate_projs[i%layer_num]
- up_proj = up_projs[i%layer_num]
- down_proj = down_projs[i%layer_num]
- gate_buf = torch.mm(input.to(proj_type), gate_proj.t())
- up_buf = torch.mm(input.to(proj_type), up_proj.t())
- intermediate = act_fn(gate_buf) * up_buf
- t_output = torch.mm(intermediate.to(proj_type), down_proj.t())
- end = time.perf_counter()
- total_time += end - start
+ mlp_torch(input[i % layer_num], gate_projs[i % layer_num], up_projs[i % layer_num], down_projs[i % layer_num])
+ end = time.perf_counter()
+ total_time = end - start
print('Quant mode: ', quant_mode)
print('Time(s): ', total_time)
print('Iteration: ', test_iter)
diff --git a/ktransformers/ktransformers_ext/bench/bench_moe.py b/ktransformers/ktransformers_ext/bench/bench_moe.py
index 909f029..6d617b7 100644
--- a/ktransformers/ktransformers_ext/bench/bench_moe.py
+++ b/ktransformers/ktransformers_ext/bench/bench_moe.py
@@ -6,7 +6,7 @@
Date : 2024-07-25 10:32:05
Version : 1.0.0
LastEditors : chenht2022
-LastEditTime : 2024-07-25 10:33:00
+LastEditTime : 2024-08-06 10:41:28
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
'''
import os, sys
@@ -15,21 +15,21 @@
import cpuinfer_ext
import torch
+expert_num = 160
+hidden_size = 5120
+intermediate_size = 1536
+stride = 16
+group_min_len = 10
+group_max_len = 1024
+n_routed_experts = 6
+layer_num = 10
+qlen = 1
+CPUInfer = cpuinfer_ext.CPUInfer(64)
+warm_up_iter = 1000
+test_iter = 10000
+
def bench_moe(quant_mode: str):
with torch.inference_mode(mode=True):
- expert_num = 10
- hidden_size = 5120
- intermediate_size = 1536
- stride = 16
- group_min_len = 10
- group_max_len = 1024
- n_routed_experts = 6
- layer_num = 10
- qlen = 1
- CPUInfer = cpuinfer_ext.CPUInfer(64)
- warm_up_iter = 1000
- test_iter = 10000
-
hidden_type = 30 # ggml_type::GGML_TYPE_BF16
if quant_mode == "fp32":
gate_type = 0 # ggml_type::GGML_TYPE_F32
@@ -104,32 +104,38 @@ def bench_moe(quant_mode: str):
up_projs.append(up_proj)
down_projs.append(down_proj)
moes.append(moe)
- expert_ids = torch.randint(0, expert_num, (layer_num, qlen, n_routed_experts), dtype=torch.int64, device = "cuda").to("cpu").contiguous()
+ expert_ids = torch.stack([torch.stack([torch.randperm(expert_num, dtype=torch.int64, device = "cuda")[:n_routed_experts] for _ in range(qlen)]) for _ in range(layer_num)]).to("cpu").contiguous()
weights = torch.rand((layer_num, qlen, n_routed_experts), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
output = torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
# warm up
for i in range(warm_up_iter):
- CPUInfer.submit(moes[i % layer_num].forward,
- qlen,
- n_routed_experts,
- expert_ids[i % layer_num].data_ptr(),
- weights[i % layer_num].data_ptr(),
- input[i % layer_num].data_ptr(),
- output[i % layer_num].data_ptr())
+ CPUInfer.submit(
+ moes[i % layer_num].forward(
+ qlen,
+ n_routed_experts,
+ expert_ids[i % layer_num].data_ptr(),
+ weights[i % layer_num].data_ptr(),
+ input[i % layer_num].data_ptr(),
+ output[i % layer_num].data_ptr()
+ )
+ )
CPUInfer.sync()
# test
start = time.perf_counter()
for i in range(test_iter):
- CPUInfer.submit(moes[i % layer_num].forward,
- qlen,
- n_routed_experts,
- expert_ids[i % layer_num].data_ptr(),
- weights[i % layer_num].data_ptr(),
- input[i % layer_num].data_ptr(),
- output[i % layer_num].data_ptr())
+ CPUInfer.submit(
+ moes[i % layer_num].forward(
+ qlen,
+ n_routed_experts,
+ expert_ids[i % layer_num].data_ptr(),
+ weights[i % layer_num].data_ptr(),
+ input[i % layer_num].data_ptr(),
+ output[i % layer_num].data_ptr()
+ )
+ )
CPUInfer.sync()
end = time.perf_counter()
total_time = end - start
diff --git a/ktransformers/ktransformers_ext/bench/bench_moe_torch.py b/ktransformers/ktransformers_ext/bench/bench_moe_torch.py
index 5075636..1aecf40 100644
--- a/ktransformers/ktransformers_ext/bench/bench_moe_torch.py
+++ b/ktransformers/ktransformers_ext/bench/bench_moe_torch.py
@@ -14,19 +14,71 @@
import torch
import torch.nn.quantized as nnq
+scale, zero_point = 0.1, 0 # Adjust scale and zero_point based on your dataset
+
+expert_num = 160
+hidden_size = 5120
+intermediate_size = 1536
+n_routed_experts = 6
+layer_num = 10
+qlen = 1
+warm_up_iter = 1000
+test_iter = 10000
+
def act_fn(x):
return x / (1.0 + torch.exp(-x))
+def mlp_torch(input, gate_proj, up_proj, down_proj):
+ if isinstance(gate_proj, nnq.Linear):
+ input_q = torch.quantize_per_tensor(input.to(torch.float32), scale, zero_point, torch.quint8)
+ gate_buf = gate_proj(input_q)
+ up_buf = up_proj(input_q)
+ gate_buf = gate_buf.dequantize()
+ up_buf = up_buf.dequantize()
+ intermediate = act_fn(gate_buf) * up_buf
+ intermediate_q = torch.quantize_per_tensor(intermediate, scale, zero_point, torch.quint8)
+ expert_output = down_proj(intermediate_q)
+ ret = expert_output.dequantize()
+ else:
+ gate_buf = torch.mm(input.to(gate_proj.dtype), gate_proj.t())
+ up_buf = torch.mm(input.to(up_proj.dtype), up_proj.t())
+ intermediate = act_fn(gate_buf) * up_buf
+ ret = torch.mm(intermediate.to(down_proj.dtype), down_proj.t())
+ return ret
+
+def moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj):
+ cnts = expert_ids.new_zeros((expert_ids.shape[0], expert_num))
+ cnts.scatter_(1, expert_ids, 1)
+ tokens_per_expert = cnts.sum(dim=0)
+ idxs = expert_ids.view(-1).argsort()
+ sorted_tokens = input[idxs // expert_ids.shape[1]]
+
+ outputs = []
+ start_idx = 0
+ for i, num_tokens in enumerate(tokens_per_expert):
+ end_idx = start_idx + num_tokens
+ if num_tokens == 0:
+ continue
+ tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
+ expert_out = mlp_torch(tokens_for_this_expert, gate_proj[i], up_proj[i], down_proj[i])
+ outputs.append(expert_out)
+ start_idx = end_idx
+
+ outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
+
+ new_x = torch.empty_like(outs)
+ new_x[idxs] = outs
+ t_output = (
+ new_x.view(*expert_ids.shape, -1)
+ .type(weights.dtype)
+ .mul_(weights.unsqueeze(dim=-1))
+ .sum(dim=1)
+ .type(new_x.dtype)
+ )
+ return t_output
+
def bench_moe(quant_mode: str):
with torch.inference_mode(mode=True):
- expert_num = 10
- hidden_size = 5120
- intermediate_size = 1536
- n_routed_experts = 6
- layer_num = 10
- warm_up_iter = 1000
- test_iter = 10000
-
if quant_mode == "fp32":
proj_type = torch.float32
bytes_per_elem = 4.000000
@@ -50,7 +102,6 @@ def bench_moe(quant_mode: str):
up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
if quant_mode == "qint8":
- scale, zero_point = 0.1, 0 # Adjust scale and zero_point based on your dataset
quantized_gate_proj = []
quantized_up_proj = []
quantized_down_proj = []
@@ -74,82 +125,20 @@ def bench_moe(quant_mode: str):
gate_projs.append(gate_proj.to(proj_type))
up_projs.append(up_proj.to(proj_type))
down_projs.append(down_proj.to(proj_type))
+ expert_ids = torch.stack([torch.stack([torch.randperm(expert_num, dtype=torch.int64, device = "cuda")[:n_routed_experts] for _ in range(qlen)]) for _ in range(layer_num)]).to("cpu").contiguous()
+ weights = torch.rand((layer_num, qlen, n_routed_experts), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
+ input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
# warm up
for i in range(warm_up_iter):
- expert_ids = torch.randint(0, expert_num, (n_routed_experts,), dtype=torch.int64).contiguous()
- weights = torch.rand((n_routed_experts,), dtype=torch.float32).contiguous()
- input = torch.randn((1, hidden_size), dtype=torch.float32).contiguous()
- if quant_mode == "qint8":
- input_q = torch.quantize_per_tensor(input, scale, zero_point, torch.quint8)
- t_output = torch.zeros((1, hidden_size), dtype=torch.float32).contiguous()
- gate_proj = gate_projs[i%layer_num]
- up_proj = up_projs[i%layer_num]
- down_proj = down_projs[i%layer_num]
- for i, expert_id in enumerate(expert_ids):
- quantized_gate = gate_proj[expert_id]
- gate_buf = quantized_gate(input_q)
- quantized_up = up_proj[expert_id]
- up_buf = quantized_up(input_q)
- gate_buf = gate_buf.dequantize()
- up_buf = up_buf.dequantize()
- intermediate = act_fn(gate_buf) * up_buf
- intermediate_q = torch.quantize_per_tensor(intermediate, scale, zero_point, torch.quint8)
- quantized_down = down_proj[expert_id]
- expert_output = quantized_down(intermediate_q)
- expert_output = expert_output.dequantize()
- t_output += weights[i] * expert_output
- else:
- t_output = torch.zeros((1, hidden_size), dtype=proj_type).contiguous()
- gate_proj = gate_projs[i%layer_num]
- up_proj = up_projs[i%layer_num]
- down_proj = down_projs[i%layer_num]
- for i, expert_id in enumerate(expert_ids):
- gate_buf = torch.mm(input.to(proj_type), gate_proj[expert_id].t())
- up_buf = torch.mm(input.to(proj_type), up_proj[expert_id].t())
- intermediate = act_fn(gate_buf) * up_buf
- expert_output = torch.mm(intermediate.to(proj_type), down_proj[expert_id].t())
- t_output += weights[i] * expert_output
+ moe_torch(input[i % layer_num], expert_ids[i % layer_num], weights[i % layer_num], gate_projs[i % layer_num], up_projs[i % layer_num], down_projs[i % layer_num])
# test
- total_time = 0
+ start = time.perf_counter()
for i in range(test_iter):
- expert_ids = torch.randint(0, expert_num, (n_routed_experts,), dtype=torch.int64).contiguous()
- weights = torch.rand((n_routed_experts,), dtype=torch.float32).contiguous()
- input = torch.randn((1, hidden_size), dtype=torch.float32).contiguous()
- start = time.perf_counter()
- if quant_mode == "qint8":
- input_q = torch.quantize_per_tensor(input, scale, zero_point, torch.quint8)
- t_output = torch.zeros((1, hidden_size), dtype=torch.float32).contiguous()
- gate_proj = gate_projs[i%layer_num]
- up_proj = up_projs[i%layer_num]
- down_proj = down_projs[i%layer_num]
- for i, expert_id in enumerate(expert_ids):
- quantized_gate = gate_proj[expert_id]
- gate_buf = quantized_gate(input_q)
- quantized_up = up_proj[expert_id]
- up_buf = quantized_up(input_q)
- gate_buf = gate_buf.dequantize()
- up_buf = up_buf.dequantize()
- intermediate = act_fn(gate_buf) * up_buf
- intermediate_q = torch.quantize_per_tensor(intermediate, scale, zero_point, torch.quint8)
- quantized_down = down_proj[expert_id]
- expert_output = quantized_down(intermediate_q)
- expert_output = expert_output.dequantize()
- t_output += weights[i] * expert_output
- else:
- t_output = torch.zeros((1, hidden_size), dtype=proj_type).contiguous()
- gate_proj = gate_projs[i%layer_num]
- up_proj = up_projs[i%layer_num]
- down_proj = down_projs[i%layer_num]
- for i, expert_id in enumerate(expert_ids):
- gate_buf = torch.mm(input.to(proj_type), gate_proj[expert_id].t())
- up_buf = torch.mm(input.to(proj_type), up_proj[expert_id].t())
- intermediate = act_fn(gate_buf) * up_buf
- expert_output = torch.mm(intermediate.to(proj_type), down_proj[expert_id].t())
- t_output += weights[i] * expert_output
- end = time.perf_counter()
- total_time += end - start
+ moe_torch(input[i % layer_num], expert_ids[i % layer_num], weights[i % layer_num], gate_projs[i % layer_num], up_projs[i % layer_num], down_projs[i % layer_num])
+ end = time.perf_counter()
+ total_time = end - start
print('Quant mode: ', quant_mode)
print('Time(s): ', total_time)
print('Iteration: ', test_iter)
diff --git a/ktransformers/ktransformers_ext/cpu_backend/cpuinfer.h b/ktransformers/ktransformers_ext/cpu_backend/cpuinfer.h
index eae6f90..9618e6b 100644
--- a/ktransformers/ktransformers_ext/cpu_backend/cpuinfer.h
+++ b/ktransformers/ktransformers_ext/cpu_backend/cpuinfer.h
@@ -1,12 +1,12 @@
/**
- * @Description :
+ * @Description :
* @Author : chenht2022
* @Date : 2024-07-16 10:43:18
* @Version : 1.0.0
- * @LastEditors : chenht2022
- * @LastEditTime : 2024-07-25 10:33:42
- * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
-**/
+ * @LastEditors : chenht2022
+ * @LastEditTime : 2024-08-07 09:47:43
+ * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
+ **/
#ifndef CPUINFER_CPUINFER_H
#define CPUINFER_CPUINFER_H
@@ -17,6 +17,7 @@
#include
#include
#include
+#include "cuda_runtime.h"
#include "backend.h"
#include "task_queue.h"
@@ -39,16 +40,39 @@ class CPUInfer {
}
template
- void submit(Func f, Obj* obj, Args... args) {
+ void enqueue(Func f, Obj* obj, Args... args) {
task_queue_->enqueue([=]() {
std::invoke(f, *obj, args..., backend_);
});
}
+ void submit(std::pair params) {
+ void (*func)(void*) = (void (*)(void*))params.first;
+ void* args = (void*)params.second;
+ *((CPUInfer**)args) = this;
+ func(args);
+ }
+
void sync() {
task_queue_->sync();
}
+ void submit_with_cuda_stream(intptr_t user_cuda_stream, std::pair params) {
+ void (*func)(void*) = (void (*)(void*))params.first;
+ void* args = (void*)params.second;
+ *((CPUInfer**)args) = this;
+ cudaLaunchHostFunc((cudaStream_t)user_cuda_stream, (cudaHostFn_t)func, args);
+ }
+
+ static void sync_(void* cpu_infer_ptr) {
+ CPUInfer* cpuinfer = (CPUInfer*)cpu_infer_ptr;
+ cpuinfer->sync();
+ }
+
+ void sync_with_cuda_stream(intptr_t user_cuda_stream) {
+ cudaLaunchHostFunc((cudaStream_t)user_cuda_stream, (cudaHostFn_t)&sync_, (void*)this);
+ }
+
public:
Backend* backend_;
TaskQueue* task_queue_;
diff --git a/ktransformers/ktransformers_ext/cpu_backend/task_queue.h b/ktransformers/ktransformers_ext/cpu_backend/task_queue.h
index a633a40..13836b7 100644
--- a/ktransformers/ktransformers_ext/cpu_backend/task_queue.h
+++ b/ktransformers/ktransformers_ext/cpu_backend/task_queue.h
@@ -4,7 +4,7 @@
* @Date : 2024-07-16 10:43:18
* @Version : 1.0.0
* @LastEditors : chenxl
- * @LastEditTime : 2024-08-08 04:23:51
+ * @LastEditTime : 2024-08-12 12:28:25
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
**/
#ifndef CPUINFER_TASKQUEUE_H
diff --git a/ktransformers/ktransformers_ext/cuda/binding.cpp b/ktransformers/ktransformers_ext/cuda/binding.cpp
index 2d5da68..06ec5f3 100644
--- a/ktransformers/ktransformers_ext/cuda/binding.cpp
+++ b/ktransformers/ktransformers_ext/cuda/binding.cpp
@@ -3,8 +3,8 @@
* @Author : Azure-Tang
* @Date : 2024-07-25 13:38:30
* @Version : 1.0.0
- * @LastEditors : Azure
- * @LastEditTime : 2024-07-26 08:36:03
+ * @LastEditors : kkk1nak0
+ * @LastEditTime : 2024-08-12 03:05:04
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
**/
@@ -23,8 +23,14 @@ PYBIND11_MODULE(KTransformersOps, m) {
py::arg("data"), py::arg("blk_size"), py::arg("device"));
m.def("dequantize_q6_k", &dequantize_q6_k, "Function to dequantize q6_k data.",
py::arg("data"), py::arg("blk_size"), py::arg("device"));
+ m.def("dequantize_q5_k", &dequantize_q5_k, "Function to dequantize q5_k data.",
+ py::arg("data"), py::arg("blk_size"), py::arg("device"));
m.def("dequantize_q4_k", &dequantize_q4_k, "Function to dequantize q4_k data.",
py::arg("data"), py::arg("blk_size"), py::arg("device"));
+ m.def("dequantize_q3_k", &dequantize_q3_k, "Function to dequantize q3_k data.",
+ py::arg("data"), py::arg("blk_size"), py::arg("device"));
+ m.def("dequantize_q2_k", &dequantize_q2_k, "Function to dequantize q2_k data.",
+ py::arg("data"), py::arg("blk_size"), py::arg("device"));
m.def("gptq_marlin_gemm", &gptq_marlin_gemm, "Function to perform GEMM using Marlin quantization.",
py::arg("a"), py::arg("b_q_weight"), py::arg("b_scales"), py::arg("g_idx"),
py::arg("perm"), py::arg("workspace"), py::arg("num_bits"), py::arg("size_m"),
diff --git a/ktransformers/ktransformers_ext/cuda/custom_gguf/binding.cpp b/ktransformers/ktransformers_ext/cuda/custom_gguf/binding.cpp
index ea52e8f..70fc606 100644
--- a/ktransformers/ktransformers_ext/cuda/custom_gguf/binding.cpp
+++ b/ktransformers/ktransformers_ext/cuda/custom_gguf/binding.cpp
@@ -12,14 +12,22 @@ int test(){
}
torch::Tensor dequantize_q6_k(torch::Tensor data, int blk_size, torch::Device device);
+torch::Tensor dequantize_q5_k(torch::Tensor data, int blk_size, torch::Device device);
+torch::Tensor dequantize_q2_k(torch::Tensor data, int blk_size, torch::Device device);
PYBIND11_MODULE(cudaops, m) {
m.def("dequantize_q8_0", &dequantize_q8_0, "Function to dequantize q8_0 data.",
py::arg("data"), py::arg("blk_size"), py::arg("device"));
m.def("dequantize_q6_k", &dequantize_q6_k, "Function to dequantize q6_k data.",
py::arg("data"), py::arg("blk_size"), py::arg("device"));
+ m.def("dequantize_q5_k", &dequantize_q5_k, "Function to dequantize q5_k data.",
+ py::arg("data"), py::arg("blk_size"), py::arg("device"));
m.def("dequantize_q4_k", &dequantize_q4_k, "Function to dequantize q4_k data.",
py::arg("data"), py::arg("blk_size"), py::arg("device"));
+ m.def("dequantize_q3_k", &dequantize_q3_k, "Function to dequantize q3_k data.",
+ py::arg("data"), py::arg("blk_size"), py::arg("device"));
+ m.def("dequantize_q2_k", &dequantize_q2_k, "Function to dequantize q2_k data.",
+ py::arg("data"), py::arg("blk_size"), py::arg("device"));
m.def("test", &test, "Function to test.");
}
diff --git a/ktransformers/ktransformers_ext/cuda/custom_gguf/custom_ggml.h b/ktransformers/ktransformers_ext/cuda/custom_gguf/custom_ggml.h
deleted file mode 100644
index 333dc69..0000000
--- a/ktransformers/ktransformers_ext/cuda/custom_gguf/custom_ggml.h
+++ /dev/null
@@ -1,39 +0,0 @@
-
-
-
-#include
-
-
-__device__ float ggml_compute_fp16_to_fp32(uint16_t h) {
- return __uint2float_rd(h);
-}
-
-static inline float ggml_compute_fp16_to_fp32(uint16_t h) {
- uint16_t tmp;
- memcpy(&tmp, &h, sizeof(ggml_fp16_t));
- return (float)tmp;
-}
-
-// define the global table for fp16 to fp32 conversion
-__device__ float ggml_table_f32_f16[1 << 16];
-
-// CUDA Kernel to init the table
-__global__ void init_fp16_to_fp32_table() {
- int idx = blockIdx.x * blockDim.x + threadIdx.x;
- for (auto blk_id = idx; blk_id<(1 << 16); blk_id+=blockDim.x * gridDim.x){
- ggml_table_f32_f16[blk_id] = GGML_COMPUTE_FP16_TO_FP32(blk_id);
- }
-}
-
-#define GGML_COMPUTE_FP16_TO_FP32(x) ggml_compute_fp16_to_fp32(x)
-
-extern __device__ float ggml_table_f32_f16[1 << 16]; // Declare as __device__ if used within device code
-
-// This version of the function is designed to be called from within a CUDA kernel
-#if !defined(GGML_FP16_TO_FP32)
-__device__ float ggml_lookup_fp16_to_fp32(uint16_t f) {
- return ggml_table_f32_f16[f];
-}
-
-#define GGML_FP16_TO_FP32(x) ggml_lookup_fp16_to_fp32(x)
-#endif
\ No newline at end of file
diff --git a/ktransformers/ktransformers_ext/cuda/custom_gguf/dequant.cu b/ktransformers/ktransformers_ext/cuda/custom_gguf/dequant.cu
index 38f4842..cc5552b 100644
--- a/ktransformers/ktransformers_ext/cuda/custom_gguf/dequant.cu
+++ b/ktransformers/ktransformers_ext/cuda/custom_gguf/dequant.cu
@@ -3,8 +3,8 @@
* @Author : Azure-Tang, Boxin Zhang
* @Date : 2024-07-25 13:38:30
* @Version : 1.0.0
- * @LastEditors : Azure
- * @LastEditTime : 2024-07-26 11:58:50
+ * @LastEditors : kkk1nak0
+ * @LastEditTime : 2024-08-12 04:18:04
* Adapted from https://github.com/ggerganov/ggml/blob/fca1caafea7de9fbd7efc733b9818f9cf2da3050/src/ggml-quants.c
* Copyright (c) 2023-2024 The ggml authors
* Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
@@ -14,6 +14,7 @@
#include
#include
#include
+#include
__global__ void dequantize_q8_0_kernel(float* output, const float* scales, const int8_t* qs, int num_blocks, int blk_size) {
int global_idx = blockIdx.x * blockDim.x + threadIdx.x;
@@ -35,6 +36,97 @@ __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t * __restrict_
}
}
+__global__ void dequantize_q2_k_kernel(int8_t* data, float* output, int blk_size, int num_blocks) {
+ int global_idx = blockIdx.x * blockDim.x + threadIdx.x;
+ for (auto block_id=global_idx; block_id(data + block_id * blk_size + 80)));
+ const float min = __half2float(*(reinterpret_cast(data + block_id * blk_size + 82)));
+
+ const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 16);
+
+ int is = 0;
+ float dl, ml;
+
+ for (int n = 0; n < 256; n += 128) {
+ int shift = 0;
+ for (int j = 0; j < 4; ++j) {
+ uint8_t* scales = (uint8_t*)(data + block_id * blk_size + (is++));
+ uint8_t sc = *scales;
+ dl = d * (sc & 0xF); ml = min * (sc >> 4);
+ for (int l = 0; l < 16; ++l) *output_blk++ = dl * ((int8_t)((q[l] >> shift) & 3)) - ml;
+
+ scales = (uint8_t*)(data + block_id * blk_size + (is++));
+ sc = *scales;
+
+ dl = d * (sc & 0xF); ml = min * (sc >> 4);
+ for (int l = 0; l < 16; ++l) *output_blk++ = dl * ((int8_t)((q[l+16] >> shift) & 3)) - ml;
+
+ shift += 2;
+ }
+ q += 32;
+ }
+ }
+}
+
+__global__ void dequantize_q3_k_kernel(int8_t* data, float* output, int blk_size, int num_blocks) {
+
+ int global_idx = blockIdx.x * blockDim.x + threadIdx.x;
+ const uint32_t kmask1 = 0x03030303;
+ const uint32_t kmask2 = 0x0f0f0f0f;
+ for (auto block_id=global_idx; block_id(data + block_id * blk_size + 108)));
+
+ const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 32);
+ const uint8_t * __restrict__ hm = (uint8_t*)(data + block_id * blk_size + 0);
+ uint8_t m = 1;
+
+
+ uint8_t* block_scales = (uint8_t*)(data + block_id * blk_size + 96);
+
+ for (int i = 0; i < 3; i++) {
+ aux[i] = 0;
+ for (int j = 0; j < 4; j++) {
+ aux[i] |= ((uint32_t)block_scales[i * 4 + j]) << (j * 8);
+ }
+ }
+
+ uint32_t tmp = aux[2];
+ aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4);
+ aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4);
+ aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4);
+ aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4);
+
+ int is = 0;
+ float dl;
+ for (int n = 0; n < 256; n += 128) {
+ int shift = 0;
+ for (int j = 0; j < 4; ++j) {
+
+ dl = d_all * (scales[is++] - 32);
+ for (int l = 0; l < 16; ++l) {
+ *output_blk++ = dl * ((int8_t)((q[l+ 0] >> shift) & 3) - ((hm[l+ 0] & m) ? 0 : 4));
+ }
+
+ dl = d_all * (scales[is++] - 32);
+ for (int l = 0; l < 16; ++l) {
+ *output_blk++ = dl * ((int8_t)((q[l+16] >> shift) & 3) - ((hm[l+16] & m) ? 0 : 4));
+ }
+
+ shift += 2;
+ m <<= 1;
+ }
+ q += 32;
+ }
+ }
+}
+
+
__global__ void dequantize_q4_k_kernel(int8_t* data, float* output, int blk_size, int num_blocks) {
int global_idx = blockIdx.x * blockDim.x + threadIdx.x;
for (auto block_id=global_idx; block_id(data + block_id * blk_size + 0)));
+ const float min = __half2float(*(reinterpret_cast(data + block_id * blk_size + 2)));
+
+ const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 16);
+ const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size + 48);
+
+ int is = 0;
+ uint8_t sc, m;
+ uint8_t u1 = 1, u2 = 2;
+ uint8_t* scales = (uint8_t*)(data + block_id * blk_size + 4);
+
+ for (int j = 0; j < 256; j += 64) {
+ get_scale_min_k4(is + 0, scales, &sc, &m);
+ const float d1 = d * sc; const float m1 = min * m;
+ get_scale_min_k4(is + 1, scales, &sc, &m);
+ const float d2 = d * sc; const float m2 = min * m;
+ for (int l = 0; l < 32; ++l) *output_blk++ = d1 * ((ql[l] & 0xF) + (qh[l] & u1 ? 16 : 0)) - m1;
+ for (int l = 0; l < 32; ++l) *output_blk++ = d2 * ((ql[l] >> 4) + (qh[l] & u2 ? 16 : 0)) - m2;
+ ql += 32; is += 2;
+ u1 <<= 2; u2 <<= 2;
+ }
+ }
+}
+
__global__ void dequantize_q6_k_kernel(int8_t* data, float* output, int blk_size, int num_blocks) {
int global_idx = blockIdx.x * blockDim.x + threadIdx.x;
for (auto block_id=global_idx; block_id>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, num_blocks);
+
+ cudaDeviceSynchronize();
+ return output;
+}
+
torch::Tensor dequantize_q4_k(torch::Tensor data, int blk_size, torch::Device device) {
// data.numel%blk_size should be 0, else raise err
int num_blocks = data.numel() / blk_size;
+ const at::cuda::OptionalCUDAGuard device_guard(device);
auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous);
auto data_gpu = torch::empty({data.numel()}, options);
@@ -162,3 +304,39 @@ torch::Tensor dequantize_q4_k(torch::Tensor data, int blk_size, torch::Device de
cudaDeviceSynchronize();
return output;
}
+
+torch::Tensor dequantize_q3_k(torch::Tensor data, int blk_size, torch::Device device) {
+ int num_blocks = data.numel() / blk_size;
+
+ auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous);
+ auto data_gpu = torch::empty({data.numel()}, options);
+
+ data_gpu.copy_(data, false);
+
+ // Create output tensor
+ auto output = torch::zeros({num_blocks, 256}, torch::dtype(torch::kFloat32).device(device));
+
+ // Launch kernel
+ dequantize_q3_k_kernel<<< 512, 256 >>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, num_blocks);
+
+ cudaDeviceSynchronize();
+ return output;
+}
+
+torch::Tensor dequantize_q2_k(torch::Tensor data, int blk_size, torch::Device device) {
+ int num_blocks = data.numel() / blk_size;
+
+ auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous);
+ auto data_gpu = torch::empty({data.numel()}, options);
+
+ data_gpu.copy_(data, false);
+
+ // Create output tensor
+ auto output = torch::zeros({num_blocks, 256}, torch::dtype(torch::kFloat32).device(device));
+
+ // Launch kernel
+ dequantize_q2_k_kernel<<< 512, 256 >>>(data_gpu.data_ptr(), output.data_ptr(), blk_size, num_blocks);
+
+ cudaDeviceSynchronize();
+ return output;
+}
\ No newline at end of file
diff --git a/ktransformers/ktransformers_ext/cuda/custom_gguf/ops.h b/ktransformers/ktransformers_ext/cuda/custom_gguf/ops.h
index 9af8f30..5196f88 100644
--- a/ktransformers/ktransformers_ext/cuda/custom_gguf/ops.h
+++ b/ktransformers/ktransformers_ext/cuda/custom_gguf/ops.h
@@ -3,8 +3,8 @@
* @Author : Azure-Tang
* @Date : 2024-07-22 09:27:55
* @Version : 1.0.0
- * @LastEditors : Azure
- * @LastEditTime : 2024-07-26 08:38:20
+ * @LastEditors : kkk1nak0
+ * @LastEditTime : 2024-08-12 03:48:46
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
**/
#pragma once
@@ -15,4 +15,7 @@
torch::Tensor dequantize_q8_0(torch::Tensor data, int blk_size, torch::Device device);
torch::Tensor dequantize_q6_k(torch::Tensor data, int blk_size, torch::Device device);
-torch::Tensor dequantize_q4_k(torch::Tensor data, int blk_size, torch::Device device);
\ No newline at end of file
+torch::Tensor dequantize_q5_k(torch::Tensor data, int blk_size, torch::Device device);
+torch::Tensor dequantize_q4_k(torch::Tensor data, int blk_size, torch::Device device);
+torch::Tensor dequantize_q3_k(torch::Tensor data, int blk_size, torch::Device device);
+torch::Tensor dequantize_q2_k(torch::Tensor data, int blk_size, torch::Device device);
\ No newline at end of file
diff --git a/ktransformers/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cu b/ktransformers/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cu
index 9205d3b..54e538a 100644
--- a/ktransformers/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cu
+++ b/ktransformers/ktransformers_ext/cuda/gptq_marlin/gptq_marlin.cu
@@ -23,7 +23,7 @@
*/
#include "gptq_marlin.cuh"
#include "gptq_marlin_dtypes.cuh"
-
+#include
#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \
static_assert(std::is_same::value || \
std::is_same::value, \
@@ -1774,6 +1774,7 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
torch::Tensor& perm, torch::Tensor& workspace,
int64_t num_bits, int64_t size_m, int64_t size_n,
int64_t size_k, bool is_k_full) {
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
// Verify num_bits
TORCH_CHECK(num_bits == 4 || num_bits == 8,
"num_bits must be 4 or 8. Got = ", num_bits);
@@ -1816,7 +1817,6 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
// Alloc buffers
- const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
torch::Tensor c = torch::empty({size_m, size_n}, options);
torch::Tensor a_tmp = torch::empty({size_m, size_k}, options);
diff --git a/ktransformers/ktransformers_ext/cuda/setup.py b/ktransformers/ktransformers_ext/cuda/setup.py
index 7ccf9ee..156bb0e 100644
--- a/ktransformers/ktransformers_ext/cuda/setup.py
+++ b/ktransformers/ktransformers_ext/cuda/setup.py
@@ -2,17 +2,25 @@
from setuptools import setup, Extension
from torch.utils import cpp_extension
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-
-# setup marlin gemm
-setup(name='KTransformersOps',
- ext_modules=[
- CUDAExtension('KTransformersOps', [
+setup(
+ name='KTransformersOps',
+ ext_modules=[
+ CUDAExtension(
+ 'KTransformersOps', [
'custom_gguf/dequant.cu',
'binding.cpp',
'gptq_marlin/gptq_marlin.cu',
- # 'gptq_marlin_repack.cu',
- ])
- ],
- cmdclass={'build_ext': BuildExtension
-})
-
+ # 'gptq_marlin_repack.cu',
+ ],
+ extra_compile_args={
+ 'cxx': ['-O3'],
+ 'nvcc': [
+ '-O3',
+ '--use_fast_math',
+ '-Xcompiler', '-fPIC',
+ ]
+ },
+ )
+ ],
+ cmdclass={'build_ext': BuildExtension}
+)
\ No newline at end of file
diff --git a/ktransformers/ktransformers_ext/examples/test_linear.py b/ktransformers/ktransformers_ext/examples/test_linear.py
index 6cb8d0c..7a331db 100644
--- a/ktransformers/ktransformers_ext/examples/test_linear.py
+++ b/ktransformers/ktransformers_ext/examples/test_linear.py
@@ -6,7 +6,7 @@
Date : 2024-07-25 10:32:05
Version : 1.0.0
LastEditors : chenht2022
-LastEditTime : 2024-07-25 10:34:00
+LastEditTime : 2024-08-06 10:36:59
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
'''
import os, sys
@@ -15,23 +15,23 @@
import cpuinfer_ext
import torch
-with torch.inference_mode(mode=True):
- input_size = 16384
- output_size = 5120
- stride = 32
- proj_type = 1 # ggml_type::GGML_TYPE_F16
- hidden_type = 1 # ggml_type::GGML_TYPE_F16
- layer_num = 10
- CPUInfer = cpuinfer_ext.CPUInfer(48)
- validation_iter = 100
- warm_up_iter = 1000
- test_iter = 10000
+input_size = 16384
+output_size = 5120
+stride = 32
+group_max_len = 1024
+proj_type = 1 # ggml_type::GGML_TYPE_F16
+hidden_type = 1 # ggml_type::GGML_TYPE_F16
+qlen = 30
+layer_num = 10
+CPUInfer = cpuinfer_ext.CPUInfer(48)
+validation_iter = 100
+with torch.inference_mode(mode=True):
linears = []
projs = []
for _ in range(layer_num):
proj = torch.randn((output_size, input_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous()
- config = cpuinfer_ext.linear.LinearConfig(input_size, output_size, stride, proj.data_ptr(), proj_type, hidden_type)
+ config = cpuinfer_ext.linear.LinearConfig(input_size, output_size, stride, group_max_len, proj.data_ptr(), proj_type, hidden_type)
linear = cpuinfer_ext.linear.Linear(config)
projs.append(proj)
linears.append(linear)
@@ -39,11 +39,17 @@
# validation
for i in range(validation_iter):
linear = linears[i % layer_num]
- input = torch.randn((1, input_size), dtype=torch.float16).contiguous()
- output = torch.empty((1, output_size), dtype=torch.float16).contiguous()
+ input = torch.randn((qlen, input_size), dtype=torch.float16).contiguous()
+ output = torch.empty((qlen, output_size), dtype=torch.float16).contiguous()
input = input / 100
- CPUInfer.submit(linear.forward, input.data_ptr(), output.data_ptr())
+ CPUInfer.submit(
+ linear.forward(
+ qlen,
+ input.data_ptr(),
+ output.data_ptr()
+ )
+ )
CPUInfer.sync()
# print('cpuinfer output', output)
@@ -54,30 +60,3 @@
diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output))
print('diff = ', diff)
assert(diff < 0.001)
-
- # warm up
- for i in range(warm_up_iter):
- linear = linears[i % layer_num]
- input = torch.randn((1, input_size), dtype=torch.float16).contiguous()
- output = torch.empty((1, output_size), dtype=torch.float16).contiguous()
- input = input / 100
- CPUInfer.submit(linear.forward, input.data_ptr(), output.data_ptr())
- CPUInfer.sync()
-
- # test
- total_time = 0
- for i in range(test_iter):
- linear = linears[i % layer_num]
- input = torch.randn((1, input_size), dtype=torch.float16).contiguous()
- output = torch.empty((1, output_size), dtype=torch.float16).contiguous()
- input = input / 100
- start = time.perf_counter()
- CPUInfer.submit(linear.forward, input.data_ptr(), output.data_ptr())
- CPUInfer.sync()
- end = time.perf_counter()
- total_time += end - start
- print('Time: ', total_time)
- print('Iteration: ', test_iter)
- print('Time per iteration: ', total_time / test_iter)
- print('Bandwidth: ', input_size * output_size * 2 * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s')
- print("All tasks completed.")
\ No newline at end of file
diff --git a/ktransformers/ktransformers_ext/examples/test_mlp.py b/ktransformers/ktransformers_ext/examples/test_mlp.py
index d965877..9805e72 100644
--- a/ktransformers/ktransformers_ext/examples/test_mlp.py
+++ b/ktransformers/ktransformers_ext/examples/test_mlp.py
@@ -6,7 +6,7 @@
Date : 2024-07-25 10:32:05
Version : 1.0.0
LastEditors : chenht2022
-LastEditTime : 2024-07-25 10:34:03
+LastEditTime : 2024-08-06 10:37:28
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
'''
import os, sys
@@ -15,20 +15,30 @@
import cpuinfer_ext
import torch
-with torch.inference_mode(mode=True):
- hidden_size = 5120
- intermediate_size = 3072
- stride = 32
- gate_type = 1 # ggml_type::GGML_TYPE_F16
- up_type = 1 # ggml_type::GGML_TYPE_F16
- down_type = 1 # ggml_type::GGML_TYPE_F16
- hidden_type = 1 # ggml_type::GGML_TYPE_F16
- layer_num = 10
- CPUInfer = cpuinfer_ext.CPUInfer(48)
- validation_iter = 100
- warm_up_iter = 1000
- test_iter = 10000
+hidden_size = 5120
+intermediate_size = 3072
+stride = 32
+group_max_len = 1024
+gate_type = 1 # ggml_type::GGML_TYPE_F16
+up_type = 1 # ggml_type::GGML_TYPE_F16
+down_type = 1 # ggml_type::GGML_TYPE_F16
+hidden_type = 1 # ggml_type::GGML_TYPE_F16
+qlen = 30
+layer_num = 10
+CPUInfer = cpuinfer_ext.CPUInfer(48)
+validation_iter = 100
+
+def act_fn(x):
+ return x / (1.0 + torch.exp(-x))
+
+def mlp_torch(input, gate_proj, up_proj, down_proj):
+ gate_buf = torch.mm(input, gate_proj.t())
+ up_buf = torch.mm(input, up_proj.t())
+ intermediate = act_fn(gate_buf) * up_buf
+ ret = torch.mm(intermediate, down_proj.t())
+ return ret
+with torch.inference_mode(mode=True):
mlps = []
gate_projs = []
up_projs = []
@@ -37,7 +47,7 @@
gate_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous()
up_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous()
down_proj = torch.randn((hidden_size, intermediate_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous()
- config = cpuinfer_ext.mlp.MLPConfig(hidden_size, intermediate_size, stride, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type)
+ config = cpuinfer_ext.mlp.MLPConfig(hidden_size, intermediate_size, stride, group_max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type)
mlp = cpuinfer_ext.mlp.MLP(config)
gate_projs.append(gate_proj)
up_projs.append(up_proj)
@@ -47,52 +57,26 @@
# validation
for i in range(validation_iter):
mlp = mlps[i % layer_num]
- input = torch.randn((1, hidden_size), dtype=torch.float16).contiguous()
- output = torch.empty((1, hidden_size), dtype=torch.float16).contiguous()
+ input = torch.randn((qlen, hidden_size), dtype=torch.float16).contiguous()
+ output = torch.empty((qlen, hidden_size), dtype=torch.float16).contiguous()
input = input / 100
- CPUInfer.submit(mlp.forward, input.data_ptr(), output.data_ptr())
+ CPUInfer.submit(
+ mlp.forward(
+ qlen,
+ input.data_ptr(),
+ output.data_ptr()
+ )
+ )
CPUInfer.sync()
# print('cpuinfer output', output)
- def act_fn(x):
- return x / (1.0 + torch.exp(-x))
gate_proj = gate_projs[i%layer_num]
up_proj = up_projs[i%layer_num]
down_proj = down_projs[i%layer_num]
- gate_buf = torch.mm(input, gate_proj.t())
- up_buf = torch.mm(input, up_proj.t())
- intermediate = act_fn(gate_buf) * up_buf
- t_output = torch.mm(intermediate, down_proj.t())
+ t_output = mlp_torch(input, gate_proj, up_proj, down_proj)
# print('torch output', t_output)
diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output))
print('diff = ', diff)
assert(diff < 0.001)
-
- # warm up
- for i in range(warm_up_iter):
- mlp = mlps[i % layer_num]
- input = torch.randn((1, hidden_size), dtype=torch.float16).contiguous()
- output = torch.empty((1, hidden_size), dtype=torch.float16).contiguous()
- input = input / 100
- CPUInfer.submit(mlp.forward, input.data_ptr(), output.data_ptr())
- CPUInfer.sync()
-
- # test
- total_time = 0
- for i in range(test_iter):
- mlp = mlps[i % layer_num]
- input = torch.randn((1, hidden_size), dtype=torch.float16).contiguous()
- output = torch.empty((1, hidden_size), dtype=torch.float16).contiguous()
- input = input / 100
- start = time.time()
- CPUInfer.submit(mlp.forward, input.data_ptr(), output.data_ptr())
- CPUInfer.sync()
- end = time.time()
- total_time += end - start
- print('Time: ', total_time)
- print('Iteration: ', test_iter)
- print('Time per iteration: ', total_time / test_iter)
- print('Bandwidth: ', hidden_size * intermediate_size * 3 * 2 * test_iter / total_time / 1024 / 1024 / 1024, 'GB/s')
- print("All tasks completed.")
\ No newline at end of file
diff --git a/ktransformers/ktransformers_ext/examples/test_moe.py b/ktransformers/ktransformers_ext/examples/test_moe.py
index 0597811..3fa4dbd 100644
--- a/ktransformers/ktransformers_ext/examples/test_moe.py
+++ b/ktransformers/ktransformers_ext/examples/test_moe.py
@@ -6,7 +6,7 @@
Date : 2024-07-25 10:32:05
Version : 1.0.0
LastEditors : chenht2022
-LastEditTime : 2024-07-25 10:34:06
+LastEditTime : 2024-08-06 10:38:05
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
'''
import os, sys
@@ -15,25 +15,64 @@
import cpuinfer_ext
import torch
-with torch.inference_mode(mode=True):
- expert_num = 10
- hidden_size = 5120
- intermediate_size = 1536
- stride = 32
- group_min_len = 10
- group_max_len = 1024
- gate_type = 1 # ggml_type::GGML_TYPE_F16
- up_type = 1 # ggml_type::GGML_TYPE_F16
- down_type = 1 # ggml_type::GGML_TYPE_F16
- hidden_type = 1 # ggml_type::GGML_TYPE_F16
- n_routed_experts = 6
- qlen = 30
- layer_num = 10
- CPUInfer = cpuinfer_ext.CPUInfer(48)
- validation_iter = 100
- warm_up_iter = 1000
- test_iter = 10000
+expert_num = 160
+hidden_size = 5120
+intermediate_size = 1536
+stride = 32
+group_min_len = 10
+group_max_len = 1024
+gate_type = 1 # ggml_type::GGML_TYPE_F16
+up_type = 1 # ggml_type::GGML_TYPE_F16
+down_type = 1 # ggml_type::GGML_TYPE_F16
+hidden_type = 1 # ggml_type::GGML_TYPE_F16
+n_routed_experts = 6
+qlen = 30
+layer_num = 10
+CPUInfer = cpuinfer_ext.CPUInfer(48)
+validation_iter = 100
+
+def act_fn(x):
+ return x / (1.0 + torch.exp(-x))
+
+def mlp_torch(input, gate_proj, up_proj, down_proj):
+ gate_buf = torch.mm(input, gate_proj.t())
+ up_buf = torch.mm(input, up_proj.t())
+ intermediate = act_fn(gate_buf) * up_buf
+ ret = torch.mm(intermediate, down_proj.t())
+ return ret
+
+def moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj):
+ cnts = expert_ids.new_zeros((expert_ids.shape[0], expert_num))
+ cnts.scatter_(1, expert_ids, 1)
+ tokens_per_expert = cnts.sum(dim=0)
+ idxs = expert_ids.view(-1).argsort()
+ sorted_tokens = input[idxs // expert_ids.shape[1]]
+
+ outputs = []
+ start_idx = 0
+ for i, num_tokens in enumerate(tokens_per_expert):
+ end_idx = start_idx + num_tokens
+ if num_tokens == 0:
+ continue
+ tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
+ expert_out = mlp_torch(tokens_for_this_expert, gate_proj[i], up_proj[i], down_proj[i])
+ outputs.append(expert_out)
+ start_idx = end_idx
+
+ outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
+
+ new_x = torch.empty_like(outs)
+ new_x[idxs] = outs
+ t_output = (
+ new_x.view(*expert_ids.shape, -1)
+ .type(weights.dtype)
+ .mul_(weights.unsqueeze(dim=-1))
+ .sum(dim=1)
+ .type(new_x.dtype)
+ )
+ return t_output
+with torch.inference_mode(mode=True):
moes = []
gate_projs = []
up_projs = []
@@ -51,63 +90,32 @@
# validation
for i in range(validation_iter):
- moe = moes[i % layer_num]
- expert_ids = torch.randint(0, expert_num, (qlen, n_routed_experts), dtype=torch.int64).contiguous()
+ expert_ids = torch.stack([torch.randperm(expert_num)[:n_routed_experts] for _ in range(qlen)]).contiguous()
weights = torch.rand((qlen, n_routed_experts), dtype=torch.float32).contiguous()
- input = torch.randn((qlen, 1, hidden_size), dtype=torch.float16).contiguous()
- output = torch.empty((qlen, 1, hidden_size), dtype=torch.float16).contiguous()
+ input = torch.randn((qlen, hidden_size), dtype=torch.float16).contiguous()
+ output = torch.empty((qlen, hidden_size), dtype=torch.float16).contiguous()
input = input / 100
- CPUInfer.submit(moe.forward, qlen, n_routed_experts, expert_ids.data_ptr(), weights.data_ptr(), input.data_ptr(), output.data_ptr())
+ moe = moes[i % layer_num]
+ CPUInfer.submit(
+ moe.forward(
+ qlen,
+ n_routed_experts,
+ expert_ids.data_ptr(),
+ weights.data_ptr(),
+ input.data_ptr(),
+ output.data_ptr()
+ )
+ )
CPUInfer.sync()
# print('cpuinfer output', output)
- def act_fn(x):
- return x / (1.0 + torch.exp(-x))
- t_output = torch.zeros((qlen, 1, hidden_size), dtype=torch.float32).contiguous()
gate_proj = gate_projs[i%layer_num]
up_proj = up_projs[i%layer_num]
down_proj = down_projs[i%layer_num]
- for token_idx in range(qlen):
- for i, expert_id in enumerate(expert_ids[token_idx]):
- gate_buf = torch.mm(input[token_idx], gate_proj[expert_id].t())
- up_buf = torch.mm(input[token_idx], up_proj[expert_id].t())
- intermediate = act_fn(gate_buf) * up_buf
- expert_output = torch.mm(intermediate, down_proj[expert_id].t())
- t_output[token_idx] += weights[token_idx][i] * expert_output
+ t_output = moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj)
# print('torch output', t_output)
diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output))
print('diff = ', diff)
assert(diff < 0.001)
-
- # warm up
- for i in range(warm_up_iter):
- moe = moes[i % layer_num]
- expert_ids = torch.randint(0, expert_num, (qlen, n_routed_experts), dtype=torch.int64).contiguous()
- weights = torch.rand((qlen, n_routed_experts), dtype=torch.float32).contiguous()
- input = torch.randn((qlen, hidden_size), dtype=torch.float16).contiguous()
- output = torch.empty((qlen, hidden_size), dtype=torch.float16).contiguous()
- input = input / 100
- CPUInfer.submit(moe.forward, qlen, n_routed_experts, expert_ids.data_ptr(), weights.data_ptr(), input.data_ptr(), output.data_ptr())
- CPUInfer.sync()
-
- # test
- total_time = 0
- for i in range(test_iter):
- moe = moes[i % layer_num]
- expert_ids = torch.randint(0, expert_num, (qlen, n_routed_experts), dtype=torch.int64).contiguous()
- weights = torch.rand((qlen, n_routed_experts), dtype=torch.float32).contiguous()
- input = torch.randn((qlen, hidden_size), dtype=torch.float16).contiguous()
- output = torch.empty((qlen, hidden_size), dtype=torch.float16).contiguous()
- input = input / 100
- start = time.perf_counter()
- CPUInfer.submit(moe.forward, qlen, n_routed_experts, expert_ids.data_ptr(), weights.data_ptr(), input.data_ptr(), output.data_ptr())
- CPUInfer.sync()
- end = time.perf_counter()
- total_time += end - start
- print('Time: ', total_time)
- print('Iteration: ', test_iter)
- print('Time per iteration: ', total_time / test_iter)
- print('Bandwidth: ', hidden_size * intermediate_size * 3 * n_routed_experts * 2 * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s')
- print("All tasks completed.")
\ No newline at end of file
diff --git a/ktransformers/ktransformers_ext/ext_bindings.cpp b/ktransformers/ktransformers_ext/ext_bindings.cpp
index 0aeead3..c220a9b 100644
--- a/ktransformers/ktransformers_ext/ext_bindings.cpp
+++ b/ktransformers/ktransformers_ext/ext_bindings.cpp
@@ -3,8 +3,8 @@
* @Author : chenht2022
* @Date : 2024-07-22 02:03:22
* @Version : 1.0.0
- * @LastEditors : chenht2022
- * @LastEditTime : 2024-07-25 10:34:23
+ * @LastEditors : chenht2022
+ * @LastEditTime : 2024-08-07 10:39:37
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
**/
// Python bindings
@@ -12,7 +12,6 @@
#include
#include
#include "cpu_backend/cpuinfer.h"
-#include "cuda_runtime.h"
#include "device_launch_parameters.h"
#include "llamafile/flags.h"
#include "operators/llamafile/linear.h"
@@ -26,239 +25,155 @@
namespace py = pybind11;
using namespace pybind11::literals;
-// Binding functions for the Linear class
class LinearBindings {
public:
- static void bind_forward(CPUInfer& cpuinfer, Linear* linear, py::args args, py::kwargs kwargs) {
- auto input = args[0].cast();
- auto output = args[1].cast();
- cpuinfer.submit(&Linear::forward, linear,
- (const void*)input, (void*)output);
- }
-
- static void bind_warm_up(CPUInfer& cpuinfer, Linear* linear, py::args args, py::kwargs kwargs) {
- cpuinfer.submit(&Linear::warm_up, linear);
- }
-
- static void bind_functions(CPUInfer& cpuinfer, py::object func, py::args args, py::kwargs kwargs) {
- auto linear = func.attr("__self__").cast();
- std::string func_name = py::str(func.attr("__func__").attr("__name__"));
-
- if (func_name == "forward") {
- bind_forward(cpuinfer, linear, args, kwargs);
- } else if (func_name == "warm_up") {
- bind_warm_up(cpuinfer, linear, args, kwargs);
- } else {
- throw py::value_error("Unsupported function: " +
- std::string(func_name));
+ class WarmUpBindinds {
+ public:
+ struct Args {
+ CPUInfer* cpuinfer;
+ Linear* linear;
+ };
+ static void inner(void* args) {
+ Args* args_ = (Args*)args;
+ args_->cpuinfer->enqueue(&Linear::warm_up, args_->linear);
+ }
+ static std::pair cpuinfer_interface(Linear& linear) {
+ Args* args = new Args{nullptr, &linear};
+ return std::make_pair((intptr_t)&inner, (intptr_t)args);
+ }
+ };
+ class ForwardBindings {
+ public:
+ struct Args {
+ CPUInfer* cpuinfer;
+ Linear* linear;
+ int qlen;
+ const void* input;
+ void* output;
+ };
+ static void inner(void* args) {
+ Args* args_ = (Args*)args;
+ args_->cpuinfer->enqueue(&Linear::forward, args_->linear, args_->qlen, args_->input, args_->output);
}
- }
+ static std::pair cpuinfer_interface(Linear& linear, int qlen, intptr_t input, intptr_t output) {
+ Args* args = new Args{nullptr, &linear, qlen, (const void*)input, (void*)output};
+ return std::make_pair((intptr_t)&inner, (intptr_t)args);
+ }
+ };
};
-// Binding functions for the MLP class
class MLPBindings {
public:
- static void bind_forward(CPUInfer& cpuinfer, MLP* mlp, py::args args, py::kwargs kwargs) {
- auto input = args[0].cast();
- auto output = args[1].cast();
- cpuinfer.submit(&MLP::forward, mlp,
- (const void*)input, (void*)output);
- }
-
- static void bind_warm_up(CPUInfer& cpuinfer, MLP* mlp, py::args args, py::kwargs kwargs) {
- cpuinfer.submit(&MLP::warm_up, mlp);
- }
-
- static void bind_functions(CPUInfer& cpuinfer, py::object func, py::args args, py::kwargs kwargs) {
- auto mlp = func.attr("__self__").cast();
- std::string func_name = py::str(func.attr("__func__").attr("__name__"));
-
- if (func_name == "forward") {
- bind_forward(cpuinfer, mlp, args, kwargs);
- } else if (func_name == "warm_up") {
- bind_warm_up(cpuinfer, mlp, args, kwargs);
- } else {
- throw py::value_error("Unsupported function: " +
- std::string(func_name));
+ class WarmUpBindinds {
+ public:
+ struct Args {
+ CPUInfer* cpuinfer;
+ MLP* mlp;
+ };
+ static void inner(void* args) {
+ Args* args_ = (Args*)args;
+ args_->cpuinfer->enqueue(&MLP::warm_up, args_->mlp);
}
- }
+ static std::pair cpuinfer_interface(MLP& mlp) {
+ Args* args = new Args{nullptr, &mlp};
+ return std::make_pair((intptr_t)&inner, (intptr_t)args);
+ }
+ };
+ class ForwardBindings {
+ public:
+ struct Args {
+ CPUInfer* cpuinfer;
+ MLP* mlp;
+ int qlen;
+ const void* input;
+ void* output;
+ };
+ static void inner(void* args) {
+ Args* args_ = (Args*)args;
+ args_->cpuinfer->enqueue(&MLP::forward, args_->mlp, args_->qlen, args_->input, args_->output);
+ }
+ static std::pair cpuinfer_interface(MLP& mlp, int qlen, intptr_t input, intptr_t output) {
+ Args* args = new Args{nullptr, &mlp, qlen, (const void*)input, (void*)output};
+ return std::make_pair((intptr_t)&inner, (intptr_t)args);
+ }
+ };
};
-// Binding functions for the MOE class
class MOEBindings {
public:
- static void bind_forward(CPUInfer& cpuinfer, MOE* moe, py::args args, py::kwargs kwargs) {
- int qlen = args[0].cast();
- int k = args[1].cast();
- auto expert_ids = args[2].cast();
- auto weights = args[3].cast();
- auto input = args[4].cast();
- auto output = args[5].cast();
- cpuinfer.submit(&MOE::forward, moe,
- qlen, k, (const uint64_t*)expert_ids, (const float*)weights, (const void*)input, (void*)output);
- }
-
- static void bind_warm_up(CPUInfer& cpuinfer, MOE* moe, py::args args, py::kwargs kwargs) {
- cpuinfer.submit(&MOE::warm_up, moe);
- }
-
- static void bind_functions(CPUInfer& cpuinfer, py::object func, py::args args, py::kwargs kwargs) {
- auto moe = func.attr("__self__").cast();
- std::string func_name = py::str(func.attr("__func__").attr("__name__"));
-
- if (func_name == "forward") {
- bind_forward(cpuinfer, moe, args, kwargs);
- } else if (func_name == "warm_up") {
- bind_warm_up(cpuinfer, moe, args, kwargs);
- } else {
- throw py::value_error("Unsupported function: " +
- std::string(func_name));
+ class WarmUpBindinds {
+ public:
+ struct Args {
+ CPUInfer* cpuinfer;
+ MOE* moe;
+ };
+ static void inner(void* args) {
+ Args* args_ = (Args*)args;
+ args_->cpuinfer->enqueue(&MOE::warm_up, args_->moe);
}
- }
-};
-
-struct MOEForwardArgs {
- CPUInfer* cpuinfer;
- MOE* moe;
- int qlen;
- int k;
- uint64_t* expert_ids;
- float* weights;
- void* input;
- void* output;
+ static std::pair cpuinfer_interface(MOE& moe) {
+ Args* args = new Args{nullptr, &moe};
+ return std::make_pair((intptr_t)&inner, (intptr_t)args);
+ }
+ };
+ class ForwardBindings {
+ public:
+ struct Args {
+ CPUInfer* cpuinfer;
+ MOE* moe;
+ int qlen;
+ int k;
+ const uint64_t* expert_ids;
+ const float* weights;
+ const void* input;
+ void* output;
+ };
+ static void inner(void* args) {
+ Args* args_ = (Args*)args;
+ args_->cpuinfer->enqueue(&MOE::forward, args_->moe, args_->qlen, args_->k, args_->expert_ids, args_->weights, args_->input, args_->output);
+ }
+ static std::pair cpuinfer_interface(MOE& moe, int qlen, int k, intptr_t expert_ids, intptr_t weights, intptr_t input, intptr_t output) {
+ Args* args = new Args{nullptr, &moe, qlen, k, (const uint64_t*)expert_ids, (const float*)weights, (const void*)input, (void*)output};
+ return std::make_pair((intptr_t)&inner, (intptr_t)args);
+ }
+ };
};
-void submit_moe_forward_with_host_args_ptr(void* host_args_ptr) {
- MOEForwardArgs* host_args = (MOEForwardArgs*)host_args_ptr;
- host_args->cpuinfer->submit(&MOE::forward, host_args->moe,
- host_args->qlen, host_args->k, host_args->expert_ids, host_args->weights, host_args->input, host_args->output);
-}
-
-void cpuinfer_sync(void* host_args_ptr) {
- CPUInfer* cpuinfer = (CPUInfer*)host_args_ptr;
- cpuinfer->sync();
-}
-
PYBIND11_MODULE(cpuinfer_ext, m) {
- auto linear_module = m.def_submodule("linear");
+ py::class_(m, "CPUInfer")
+ .def(py::init())
+ .def("submit", &CPUInfer::submit)
+ .def("submit_with_cuda_stream", &CPUInfer::submit_with_cuda_stream)
+ .def("sync", &CPUInfer::sync)
+ .def("sync_with_cuda_stream", &CPUInfer::sync_with_cuda_stream);
+ auto linear_module = m.def_submodule("linear");
py::class_(linear_module, "LinearConfig")
- .def(py::init([](int hidden_size, int intermediate_size, int stride, intptr_t proj, int proj_type, int hidden_type) {
- return LinearConfig(hidden_size, intermediate_size, stride, (void*)proj, (ggml_type)proj_type, (ggml_type)hidden_type);
+ .def(py::init([](int hidden_size, int intermediate_size, int stride, int group_max_len, intptr_t proj, int proj_type, int hidden_type) {
+ return LinearConfig(hidden_size, intermediate_size, stride, group_max_len, (void*)proj, (ggml_type)proj_type, (ggml_type)hidden_type);
}));
-
py::class_(linear_module, "Linear")
.def(py::init())
- .def("warm_up", [](Linear& linear) {
- throw std::runtime_error("!!! Doing nothing, please use CPUInfer.submit to call it!!!\n");
- })
- .def("forward", [](Linear& linear, intptr_t input, intptr_t output) {
- throw std::runtime_error("!!! Doing nothing, please use CPUInfer.submit to call it!!!\n");
- });
+ .def("warm_up", &LinearBindings::WarmUpBindinds::cpuinfer_interface)
+ .def("forward", &LinearBindings::ForwardBindings::cpuinfer_interface);
auto mlp_module = m.def_submodule("mlp");
-
py::class_(mlp_module, "MLPConfig")
- .def(py::init([](int hidden_size, int intermediate_size, int stride, intptr_t gate_proj, intptr_t up_proj, intptr_t down_proj, int gate_type, int up_type, int down_type, int hidden_type) {
- return MLPConfig(hidden_size, intermediate_size, stride, (void*)gate_proj, (void*)up_proj, (void*)down_proj, (ggml_type)gate_type, (ggml_type)up_type, (ggml_type)down_type, (ggml_type)hidden_type);
+ .def(py::init([](int hidden_size, int intermediate_size, int stride, int group_max_len, intptr_t gate_proj, intptr_t up_proj, intptr_t down_proj, int gate_type, int up_type, int down_type, int hidden_type) {
+ return MLPConfig(hidden_size, intermediate_size, stride, group_max_len, (void*)gate_proj, (void*)up_proj, (void*)down_proj, (ggml_type)gate_type, (ggml_type)up_type, (ggml_type)down_type, (ggml_type)hidden_type);
}));
-
py::class_(mlp_module, "MLP")
.def(py::init())
- .def("warm_up", [](MLP& mlp) {
- throw std::runtime_error("!!! Doing nothing, please use CPUInfer.submit to call it!!!\n");
- })
- .def("forward", [](MLP& mlp, intptr_t input, intptr_t output) {
- throw std::runtime_error("!!! Doing nothing, please use CPUInfer.submit to call it!!!\n");
- });
+ .def("warm_up", &MLPBindings::WarmUpBindinds::cpuinfer_interface)
+ .def("forward", &MLPBindings::ForwardBindings::cpuinfer_interface);
auto moe_module = m.def_submodule("moe");
-
py::class_(moe_module, "MOEConfig")
.def(py::init([](int expert_num, int routed_expert_num, int hidden_size, int intermediate_size, int stride, int group_min_len, int group_max_len, intptr_t gate_proj, intptr_t up_proj, intptr_t down_proj, int gate_type, int up_type, int down_type, int hidden_type) {
return MOEConfig(expert_num, routed_expert_num, hidden_size, intermediate_size, stride, group_min_len, group_max_len, (void*)gate_proj, (void*)up_proj, (void*)down_proj, (ggml_type)gate_type, (ggml_type)up_type, (ggml_type)down_type, (ggml_type)hidden_type);
}));
-
py::class_(moe_module, "MOE")
.def(py::init())
- .def("warm_up", [](MOE& moe) {
- throw std::runtime_error("!!! Doing nothing, please use CPUInfer.submit to call it!!!\n");
- })
- .def("forward", [](MOE& moe, int k, uint64_t expert_ids, intptr_t weights, intptr_t input, intptr_t output) {
- throw std::runtime_error("!!! Doing nothing, please use CPUInfer.submit to call it!!!\n");
- });
-
- py::class_(m, "CPUInfer")
- .def(py::init())
- .def("submit",
- [linear_module, mlp_module, moe_module](CPUInfer& cpuinfer, py::object func, py::args args, py::kwargs kwargs) {
- if (py::hasattr(func, "__self__") &&
- py::hasattr(func, "__func__")) {
- std::string class_name = py::str(func.attr("__self__")
- .attr("__class__")
- .attr("__name__"));
- if (class_name == "Linear") {
- LinearBindings::bind_functions(cpuinfer, func,
- args, kwargs);
- } else if (class_name == "MLP") {
- MLPBindings::bind_functions(cpuinfer, func,
- args, kwargs);
- } else if (class_name == "MOE") {
- MOEBindings::bind_functions(cpuinfer, func,
- args, kwargs);
- } else {
- // handle other classes
- throw py::type_error("Unsupported class type: " +
- class_name);
- }
- } else {
- // handle cases where func does not have __self__ or
- // __func__
- throw py::type_error(
- "Invalid function object: missing "
- "__self__ or __func__ attribute.");
- }
- })
- .def("submit_with_cuda_stream",
- [linear_module, mlp_module, moe_module](CPUInfer& cpuinfer, intptr_t user_cuda_stream, py::object func, py::args args, py::kwargs kwargs) {
- if (py::hasattr(func, "__self__") &&
- py::hasattr(func, "__func__")) {
- std::string class_name = py::str(func.attr("__self__")
- .attr("__class__")
- .attr("__name__"));
- if (class_name == "MOE") {
- std::string func_name = py::str(func.attr("__func__").attr("__name__"));
- if (func_name == "forward") {
- auto moe = func.attr("__self__").cast();
- int qlen = args[0].cast();
- int k = args[1].cast();
- auto expert_ids = args[2].cast();
- auto weights = args[3].cast();
- auto input = args[4].cast();
- auto output = args[5].cast();
- MOEForwardArgs* moe_forward_args = new MOEForwardArgs{&cpuinfer, moe, qlen, k, (uint64_t*)expert_ids, (float*)weights, (void*)input, (void*)output};
- // submit_moe_forward_with_host_args_ptr(moe_forward_args);
- cudaLaunchHostFunc((cudaStream_t)user_cuda_stream, (cudaHostFn_t)submit_moe_forward_with_host_args_ptr, moe_forward_args);
- } else {
- throw py::value_error("Unsupported function: " +
- std::string(func_name));
- }
- } else {
- // handle other classes
- throw py::type_error("Unsupported class type: " +
- class_name);
- }
- } else {
- // handle cases where func does not have __self__ or
- // __func__
- throw py::type_error(
- "Invalid function object: missing "
- "__self__ or __func__ attribute.");
- }
- })
- .def("sync_with_cuda_stream", [](CPUInfer& cpuinfer, intptr_t user_cuda_stream) {
- // cpuinfer_sync((void*)(&cpuinfer));
- cudaLaunchHostFunc((cudaStream_t)user_cuda_stream, (cudaHostFn_t)cpuinfer_sync, (void*)(&cpuinfer));
- })
- .def("sync", &CPUInfer::sync);
+ .def("warm_up", &MOEBindings::WarmUpBindinds::cpuinfer_interface)
+ .def("forward", &MOEBindings::ForwardBindings::cpuinfer_interface);
}
diff --git a/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/gptq.py b/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/gptq.py
deleted file mode 100644
index cda3e7a..0000000
--- a/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/gptq.py
+++ /dev/null
@@ -1,206 +0,0 @@
-import math
-import os
-import time
-from logging import getLogger
-
-import torch
-import torch.nn as nn
-import transformers
-
-from .quantizer import Quantizer
-
-
-logger = getLogger(__name__)
-
-torch.backends.cuda.matmul.allow_tf32 = False
-torch.backends.cudnn.allow_tf32 = False
-
-
-class GPTQ:
- def __init__(self, layer):
- self.layer = layer
- self.dev = self.layer.weight.device
- W = layer.weight.data.clone()
- if isinstance(self.layer, nn.Conv2d):
- W = W.flatten(1)
- if isinstance(self.layer, transformers.pytorch_utils.Conv1D):
- W = W.t()
- self.rows = W.shape[0]
- self.columns = W.shape[1]
- self.H = torch.zeros((self.columns, self.columns), device=self.dev)
- self.nsamples = 0
- self.quantizer = Quantizer()
-
- def add_batch(self, inp, out):
- if os.environ.get("DEBUG"):
- self.inp1 = inp
- self.out1 = out
- if len(inp.shape) == 2:
- inp = inp.unsqueeze(0)
- tmp = inp.shape[0]
- if isinstance(self.layer, nn.Linear) or isinstance(self.layer, transformers.Conv1D):
- if len(inp.shape) == 3:
- inp = inp.reshape((-1, inp.shape[-1]))
- inp = inp.t()
- if isinstance(self.layer, nn.Conv2d):
- unfold = nn.Unfold(
- self.layer.kernel_size,
- dilation=self.layer.dilation,
- padding=self.layer.padding,
- stride=self.layer.stride,
- )
- inp = unfold(inp)
- inp = inp.permute([1, 0, 2])
- inp = inp.flatten(1)
- self.H *= self.nsamples / (self.nsamples + tmp)
- self.nsamples += tmp
- # inp = inp.float()
- inp = math.sqrt(2 / self.nsamples) * inp.float()
- # self.H += 2 / self.nsamples * inp.matmul(inp.t())
- self.H += inp.matmul(inp.t())
-
- def fasterquant(
- self,
- blocksize=128,
- percdamp=0.01,
- group_size=-1,
- actorder=False,
- static_groups=False,
- ):
- W = self.layer.weight.data.clone()
- if isinstance(self.layer, nn.Conv2d):
- W = W.flatten(1)
- if isinstance(self.layer, transformers.Conv1D):
- W = W.t()
- W = W.float()
-
- tick = time.time()
-
- if not self.quantizer.ready():
- self.quantizer.find_params(W, weight=True)
-
- H = self.H
- del self.H
- dead = torch.diag(H) == 0
- H[dead, dead] = 1
- W[:, dead] = 0
-
- g_idx = []
- scale = []
- zero = []
- now_idx = 1
-
- if static_groups:
- import copy
-
- groups = []
- for i in range(0, self.columns, group_size):
- quantizer = copy.deepcopy(self.quantizer)
- quantizer.find_params(W[:, i : (i + group_size)], weight=True)
- scale.append(quantizer.scale)
- zero.append(quantizer.zero)
- groups.append(quantizer)
-
- if actorder:
- perm = torch.argsort(torch.diag(H), descending=True)
- W = W[:, perm]
- H = H[perm][:, perm]
- invperm = torch.argsort(perm)
-
- Losses = torch.zeros_like(W)
- Q = torch.zeros_like(W)
-
- damp = percdamp * torch.mean(torch.diag(H))
- diag = torch.arange(self.columns, device=self.dev)
- H[diag, diag] += damp
- H = torch.linalg.cholesky(H)
- H = torch.cholesky_inverse(H)
- H = torch.linalg.cholesky(H, upper=True)
- Hinv = H
-
- for i1 in range(0, self.columns, blocksize):
- i2 = min(i1 + blocksize, self.columns)
- count = i2 - i1
-
- W1 = W[:, i1:i2].clone()
- Q1 = torch.zeros_like(W1)
- Err1 = torch.zeros_like(W1)
- Losses1 = torch.zeros_like(W1)
- Hinv1 = Hinv[i1:i2, i1:i2]
-
- for i in range(count):
- w = W1[:, i]
- d = Hinv1[i, i]
-
- if group_size != -1:
- if not static_groups:
- if (i1 + i) % group_size == 0:
- self.quantizer.find_params(W[:, (i1 + i) : (i1 + i + group_size)], weight=True)
-
- if ((i1 + i) // group_size) - now_idx == -1:
- scale.append(self.quantizer.scale)
- zero.append(self.quantizer.zero)
- now_idx += 1
- else:
- idx = i1 + i
- if actorder:
- idx = perm[idx]
- self.quantizer = groups[idx // group_size]
-
- q = self.quantizer.quantize(w.unsqueeze(1)).flatten()
- Q1[:, i] = q
- Losses1[:, i] = (w - q) ** 2 / d**2
-
- err1 = (w - q) / d
- W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
- Err1[:, i] = err1
-
- Q[:, i1:i2] = Q1
- Losses[:, i1:i2] = Losses1 / 2
-
- W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])
-
- if os.environ.get("DEBUG"):
- self.layer.weight.data[:, :i2] = Q[:, :i2]
- self.layer.weight.data[:, i2:] = W[:, i2:]
- logger.debug(torch.sum((self.layer(self.inp1) - self.out1) ** 2))
- logger.debug(torch.sum(Losses))
-
- torch.cuda.synchronize()
- logger.info(f"duration: {(time.time() - tick)}")
- logger.info(f"avg loss: {torch.sum(Losses).item() / self.nsamples}")
-
- group_size = group_size if group_size != -1 else self.columns
- if static_groups and actorder:
- g_idx = [perm[i] // group_size for i in range(self.columns)]
- else:
- g_idx = [i // group_size for i in range(self.columns)]
- g_idx = torch.tensor(g_idx, dtype=torch.int32, device=Q.device)
- if actorder:
- Q = Q[:, invperm]
- g_idx = g_idx[invperm]
-
- if isinstance(self.layer, transformers.Conv1D):
- Q = Q.t()
- self.layer.weight.data = Q.reshape(self.layer.weight.shape).type_as(self.layer.weight.data)
- if os.environ.get("DEBUG"):
- logger.debug(torch.sum((self.layer(self.inp1) - self.out1) ** 2))
-
- if scale == []:
- scale.append(self.quantizer.scale)
- zero.append(self.quantizer.zero)
- scale = torch.cat(scale, dim=1)
- zero = torch.cat(zero, dim=1)
- return scale, zero, g_idx
-
- def free(self):
- if os.environ.get("DEBUG"):
- self.inp1 = None
- self.out1 = None
- self.H = None
- self.Losses = None
- self.Trace = None
- torch.cuda.empty_cache()
-
-
-__all__ = ["GPTQ"]
diff --git a/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/gptq_marlin.py b/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/gptq_marlin.py
deleted file mode 100644
index 599070f..0000000
--- a/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/gptq_marlin.py
+++ /dev/null
@@ -1,458 +0,0 @@
-import enum
-from enum import Enum
-from typing import Any, Dict, List, Optional
-
-import torch
-from torch.nn.parameter import Parameter
-
-from vllm import _custom_ops as ops
-from vllm.logger import init_logger
-from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
- set_weight_attrs)
-from vllm.model_executor.layers.quantization.base_config import (
- QuantizationConfig)
-
-logger = init_logger(__name__)
-
-GPTQ_MARLIN_TILE = 16
-GPTQ_MARLIN_MIN_THREAD_N = 64
-GPTQ_MARLIN_MIN_THREAD_K = 128
-GPTQ_MARLIN_MAX_PARALLEL = 16
-
-GPTQ_MARLIN_SUPPORTED_NUM_BITS = [4, 8]
-GPTQ_MARLIN_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128]
-GPTQ_MARLIN_SUPPORTED_SYM = [True]
-
-
-# Permutations for Marlin scale shuffling
-def get_scale_perms(num_bits: int):
- scale_perm: List[int] = []
- for i in range(8):
- scale_perm.extend([i + 8 * j for j in range(8)])
- scale_perm_single: List[int] = []
- for i in range(4):
- scale_perm_single.extend(
- [2 * i + j for j in [0, 1, 8, 9, 16, 17, 24, 25]])
- return scale_perm, scale_perm_single
-
-
-def get_pack_factor(num_bits: int):
- assert (num_bits in GPTQ_MARLIN_SUPPORTED_NUM_BITS
- ), f"Unsupported num_bits = {num_bits}"
- return 32 // num_bits
-
-
-def marlin_permute_scales(s: torch.Tensor, size_k: int, size_n: int,
- group_size: int, num_bits: int):
- scale_perm, scale_perm_single = get_scale_perms(num_bits)
- if group_size < size_k and group_size != -1:
- s = s.reshape((-1, len(scale_perm)))[:, scale_perm]
- else:
- s = s.reshape((-1, len(scale_perm_single)))[:, scale_perm_single]
- s = s.reshape((-1, size_n)).contiguous()
-
- return s
-
-
-class GPTQMarlinConfig(QuantizationConfig):
- """Config class for GPTQ Marlin"""
-
- def __init__(self, weight_bits: int, group_size: int, desc_act: bool,
- is_sym: bool) -> None:
- if desc_act and group_size == -1:
- # In this case, act_order == True is the same as act_order == False
- # (since we have only one group per output channel)
- desc_act = False
-
- self.weight_bits = weight_bits
- self.group_size = group_size
- self.desc_act = desc_act
- self.is_sym = is_sym
-
- # Verify
- if self.weight_bits not in GPTQ_MARLIN_SUPPORTED_NUM_BITS:
- raise ValueError(
- f"Marlin does not support weight_bits = {self.weight_bits}. "
- f"Only weight_bits = {GPTQ_MARLIN_SUPPORTED_NUM_BITS} "
- "are supported.")
- if self.group_size not in GPTQ_MARLIN_SUPPORTED_GROUP_SIZES:
- raise ValueError(
- f"Marlin does not support group_size = {self.group_size}. "
- f"Only group_sizes = {GPTQ_MARLIN_SUPPORTED_GROUP_SIZES} "
- "are supported.")
- if self.is_sym not in GPTQ_MARLIN_SUPPORTED_SYM:
- raise ValueError(
- f"Marlin does not support is_sym = {self.is_sym}. "
- f"Only sym = {GPTQ_MARLIN_SUPPORTED_SYM} are supported.")
-
- # Init
- self.pack_factor = get_pack_factor(weight_bits)
- self.tile_size = GPTQ_MARLIN_TILE
- self.min_thread_n = GPTQ_MARLIN_MIN_THREAD_N
- self.min_thread_k = GPTQ_MARLIN_MIN_THREAD_K
- self.max_parallel = GPTQ_MARLIN_MAX_PARALLEL
-
- def __repr__(self) -> str:
- return (f"GPTQMarlinConfig(weight_bits={self.weight_bits}, "
- f"group_size={self.group_size}, "
- f"desc_act={self.desc_act})")
-
- @classmethod
- def get_name(cls) -> str:
- return "gptq_marlin"
-
- @classmethod
- def get_supported_act_dtypes(cls) -> List[torch.dtype]:
- return [torch.half, torch.bfloat16]
-
- @classmethod
- def get_min_capability(cls) -> int:
- return 80
-
- @classmethod
- def get_config_filenames(cls) -> List[str]:
- return ["quantize_config.json"]
-
- @classmethod
- def from_config(cls, config: Dict[str, Any]) -> "GPTQMarlinConfig":
- weight_bits = cls.get_from_keys(config, ["bits"])
- group_size = cls.get_from_keys(config, ["group_size"])
- desc_act = cls.get_from_keys(config, ["desc_act"])
- is_sym = cls.get_from_keys(config, ["sym"])
- return cls(weight_bits, group_size, desc_act, is_sym)
-
- @classmethod
- def override_quantization_method(cls, hf_quant_cfg,
- user_quant) -> Optional[str]:
- can_convert = cls.is_marlin_compatible(hf_quant_cfg)
-
- is_valid_user_quant = (user_quant is None or user_quant == "marlin")
-
- if can_convert and is_valid_user_quant:
- msg = ("The model is convertible to {} during runtime."
- " Using {} kernel.".format(cls.get_name(), cls.get_name()))
- logger.info(msg)
- return cls.get_name()
-
- if can_convert and user_quant == "gptq":
- logger.info("Detected that the model can run with gptq_marlin"
- ", however you specified quantization=gptq explicitly,"
- " so forcing gptq. Use quantization=gptq_marlin for"
- " faster inference")
- return None
-
- def get_quant_method(
- self,
- layer: torch.nn.Module) -> Optional["GPTQMarlinLinearMethod"]:
- if isinstance(layer, LinearBase):
- return GPTQMarlinLinearMethod(self)
- return None
-
- def get_scaled_act_names(self) -> List[str]:
- return []
-
- @classmethod
- def is_marlin_compatible(cls, quant_config: Dict[str, Any]):
- # Extract data from quant config.
- num_bits = quant_config.get("bits", None)
- group_size = quant_config.get("group_size", None)
- sym = quant_config.get("sym", None)
- desc_act = quant_config.get("desc_act", None)
-
- # If we cannot find the info needed in the config, cannot convert.
- if (num_bits is None or group_size is None or sym is None
- or desc_act is None):
- return False
-
- # If the capability of the device is too low, cannot convert.
- major, minor = torch.cuda.get_device_capability()
- device_capability = major * 10 + minor
- if device_capability < cls.get_min_capability():
- return False
-
- # Otherwise, can convert if model satisfies marlin constraints.
- return (num_bits in GPTQ_MARLIN_SUPPORTED_NUM_BITS
- and group_size in GPTQ_MARLIN_SUPPORTED_GROUP_SIZES
- and sym in GPTQ_MARLIN_SUPPORTED_SYM)
-
-
-class GPTQMarlinState(Enum):
- REPACK = enum.auto()
- READY = enum.auto()
-
-
-class GPTQMarlinLinearMethod(LinearMethodBase):
- """Linear method for GPTQ Marlin.
-
- Args:
- quant_config: The GPTQ Marlin quantization config.
- """
-
- def __init__(self, quant_config: GPTQMarlinConfig) -> None:
- self.quant_config = quant_config
-
- def create_weights(
- self,
- layer: torch.nn.Module,
- input_size_per_partition: int,
- output_partition_sizes: List[int],
- input_size: int,
- output_size: int,
- params_dtype: torch.dtype,
- **extra_weight_attrs,
- ) -> None:
- del output_size
-
- # Normalize group_size
- if self.quant_config.group_size != -1:
- group_size = self.quant_config.group_size
- else:
- group_size = input_size
-
- # Validate dtype
- if params_dtype not in [torch.float16, torch.bfloat16]:
- raise ValueError(f"The params dtype must be float16 "
- f"or bfloat16, but got {params_dtype}")
-
- # Validate output_size_per_partition
- output_size_per_partition = sum(output_partition_sizes)
- if output_size_per_partition % self.quant_config.min_thread_n != 0:
- raise ValueError(
- f"Weight output_size_per_partition = "
- f"{output_size_per_partition} is not divisible by "
- f" min_thread_n = {self.quant_config.min_thread_n}.")
-
- # Validate input_size_per_partition
- if input_size_per_partition % self.quant_config.min_thread_k != 0:
- raise ValueError(
- f"Weight input_size_per_partition = "
- f"{input_size_per_partition} is not divisible "
- f"by min_thread_k = {self.quant_config.min_thread_k}.")
-
- if (group_size < input_size
- and input_size_per_partition % group_size != 0):
- raise ValueError(
- f"Weight input_size_per_partition = {input_size_per_partition}"
- f" is not divisible by group_size = {group_size}.")
-
- # Detect sharding of scales/zp
-
- # By default, no sharding over "input dim"
- scales_and_zp_size = input_size // group_size
- scales_and_zp_input_dim = None
-
- if self.quant_config.desc_act:
- # Act-order case
- assert self.quant_config.group_size != -1
-
- is_k_full = input_size_per_partition == input_size
-
- else:
- # No act-order case
-
- # K is always full due to full alignment with
- # group-size and shard of scales/zp
- is_k_full = True
-
- # If this is a row-parallel case, then shard scales/zp
- if (input_size != input_size_per_partition
- and self.quant_config.group_size != -1):
- scales_and_zp_size = input_size_per_partition // group_size
- scales_and_zp_input_dim = 0
-
- # Init buffers
-
- # Quantized weights
- qweight = Parameter(
- torch.empty(
- input_size_per_partition // self.quant_config.pack_factor,
- output_size_per_partition,
- dtype=torch.int32,
- ),
- requires_grad=False,
- )
- set_weight_attrs(
- qweight,
- {
- **extra_weight_attrs,
- "input_dim": 0,
- "output_dim": 1,
- "packed_dim": 0,
- "pack_factor": self.quant_config.pack_factor,
- },
- )
-
- # Activation order
- g_idx = Parameter(
- torch.empty(
- input_size_per_partition,
- dtype=torch.int32,
- ),
- requires_grad=False,
- )
- # Ignore warning from fused linear layers such as QKVParallelLinear.
- set_weight_attrs(
- g_idx,
- {
- **extra_weight_attrs, "input_dim": 0,
- "ignore_warning": True
- },
- )
-
- g_idx_sort_indices = torch.empty(
- g_idx.shape,
- dtype=torch.int32,
- )
-
- # Scales
- scales = Parameter(
- torch.empty(
- scales_and_zp_size,
- output_size_per_partition,
- dtype=params_dtype,
- ),
- requires_grad=False,
- )
- set_weight_attrs(
- scales,
- {
- **extra_weight_attrs,
- "input_dim": scales_and_zp_input_dim,
- "output_dim": 1,
- },
- )
-
- # Quantized zero-points
- qzeros = Parameter(
- torch.empty(
- scales_and_zp_size,
- output_size_per_partition // self.quant_config.pack_factor,
- dtype=torch.int32,
- device="meta",
- ),
- requires_grad=False,
- )
- set_weight_attrs(
- qzeros,
- {
- **extra_weight_attrs,
- "input_dim": scales_and_zp_input_dim,
- "output_dim": 1,
- "packed_dim": 1,
- "pack_factor": self.quant_config.pack_factor,
- },
- )
-
- # Allocate marlin workspace
- max_workspace_size = (
- output_size_per_partition //
- self.quant_config.min_thread_n) * self.quant_config.max_parallel
- workspace = torch.zeros(max_workspace_size,
- dtype=torch.int,
- requires_grad=False)
-
- layer.register_parameter("qweight", qweight)
- layer.register_parameter("g_idx", g_idx)
- layer.register_parameter("scales", scales)
- layer.register_parameter("qzeros", qzeros)
- layer.g_idx_sort_indices = g_idx_sort_indices
- layer.workspace = workspace
- layer.input_size_per_partition = input_size_per_partition
- layer.output_size_per_partition = output_size_per_partition
- layer.input_size = input_size
- layer.is_k_full = is_k_full
- layer.marlin_state = GPTQMarlinState.REPACK
-
- def apply(
- self,
- layer: torch.nn.Module,
- x: torch.Tensor,
- bias: Optional[torch.Tensor] = None,
- ) -> torch.Tensor:
- reshaped_x = x.reshape(-1, x.shape[-1])
-
- size_m = reshaped_x.shape[0]
- part_size_n = layer.output_size_per_partition
- part_size_k = layer.input_size_per_partition
- full_size_k = layer.input_size
-
- out_shape = x.shape[:-1] + (part_size_n, )
-
- if layer.marlin_state == GPTQMarlinState.REPACK:
- layer.marlin_state = GPTQMarlinState.READY
-
- # Newly generated tensors need to replace existing tensors that are
- # already registered as parameters by vLLM (and won't be freed)
- def replace_tensor(name, new_t):
- # It is important to use resize_() here since it ensures
- # the same buffer is reused
- getattr(layer, name).resize_(new_t.shape)
- getattr(layer, name).copy_(new_t)
- del new_t
-
- cur_device = layer.qweight.device
-
- # Process act_order
- if self.quant_config.desc_act:
- # Get sorting based on g_idx
- g_idx_sort_indices = torch.argsort(layer.g_idx).to(torch.int)
-
- sorted_g_idx = layer.g_idx[g_idx_sort_indices]
-
- replace_tensor("g_idx", sorted_g_idx)
- replace_tensor("g_idx_sort_indices", g_idx_sort_indices)
-
- else:
- # Reset g_idx related tensors
- layer.g_idx = Parameter(
- torch.empty(0, dtype=torch.int, device=cur_device),
- requires_grad=False,
- )
- layer.g_idx_sort_indices = Parameter(
- torch.empty(0, dtype=torch.int, device=cur_device),
- requires_grad=False,
- )
-
- # Repack weights
- marlin_qweight = ops.gptq_marlin_repack(
- layer.qweight,
- layer.g_idx_sort_indices,
- part_size_k,
- part_size_n,
- self.quant_config.weight_bits,
- )
- replace_tensor("qweight", marlin_qweight)
-
- # Permute scales
- scales_size_k = part_size_k
- scales_size_n = part_size_n
- if self.quant_config.desc_act:
- scales_size_k = full_size_k
-
- marlin_scales = marlin_permute_scales(
- layer.scales,
- scales_size_k,
- scales_size_n,
- self.quant_config.group_size,
- self.quant_config.weight_bits,
- )
- replace_tensor("scales", marlin_scales)
-
- output = ops.gptq_marlin_gemm(
- reshaped_x,
- layer.qweight,
- layer.scales,
- layer.g_idx,
- layer.g_idx_sort_indices,
- layer.workspace,
- self.quant_config.weight_bits,
- size_m,
- part_size_n,
- part_size_k,
- layer.is_k_full,
- )
-
- if bias is not None:
- output.add_(bias) # In-place add
-
- return output.reshape(out_shape)
diff --git a/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/quantizer.py b/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/quantizer.py
deleted file mode 100644
index e945a70..0000000
--- a/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/quantizer.py
+++ /dev/null
@@ -1,140 +0,0 @@
-from logging import getLogger
-
-import torch
-import torch.nn as nn
-
-
-logger = getLogger(__name__)
-
-
-def quantize(x, scale, zero, maxq):
- if maxq < 0:
- return (x > scale / 2).float() * scale + (x < zero / 2).float() * zero
- q = torch.clamp(torch.round(x / scale) + zero, 0, maxq)
- return scale * (q - zero)
-
-
-class Quantizer(nn.Module):
- def __init__(self, shape=1):
- super(Quantizer, self).__init__()
- self.register_buffer("maxq", torch.tensor(0))
- self.register_buffer("scale", torch.zeros(shape))
- self.register_buffer("zero", torch.zeros(shape))
-
- def configure(
- self,
- bits,
- perchannel=False,
- sym=True,
- mse=False,
- norm=2.4,
- grid=100,
- maxshrink=0.8,
- trits=False,
- ):
- self.maxq = torch.tensor(2**bits - 1)
- self.perchannel = perchannel
- self.sym = sym
- self.mse = mse
- self.norm = norm
- self.grid = grid
- self.maxshrink = maxshrink
- if trits:
- self.maxq = torch.tensor(-1)
-
- def find_params(self, x, weight=False):
- dev = x.device
- self.maxq = self.maxq.to(dev)
-
- shape = x.shape
- if self.perchannel:
- if weight:
- x = x.flatten(1)
- else:
- if len(shape) == 4:
- x = x.permute([1, 0, 2, 3])
- x = x.flatten(1)
- if len(shape) == 3:
- x = x.reshape((-1, shape[-1])).t()
- if len(shape) == 2:
- x = x.t()
- else:
- x = x.flatten().unsqueeze(0)
-
- tmp = torch.zeros(x.shape[0], device=dev)
- xmin = torch.minimum(x.min(1)[0], tmp)
- xmax = torch.maximum(x.max(1)[0], tmp)
-
- if self.sym:
- xmax = torch.maximum(torch.abs(xmin), xmax)
- tmp = xmin < 0
- if torch.any(tmp):
- xmin[tmp] = -xmax[tmp]
- tmp = (xmin == 0) & (xmax == 0)
- xmin[tmp] = -1
- xmax[tmp] = +1
-
- if self.maxq < 0:
- self.scale = xmax
- self.zero = xmin
- else:
- self.scale = (xmax - xmin) / self.maxq
- if self.sym:
- self.zero = torch.full_like(self.scale, (self.maxq + 1) / 2)
- else:
- self.zero = torch.round(-xmin / self.scale)
-
- if self.mse:
- best = torch.full([x.shape[0]], float("inf"), device=dev)
- for i in range(int(self.maxshrink * self.grid)):
- p = 1 - i / self.grid
- xmin1 = p * xmin
- xmax1 = p * xmax
- scale1 = (xmax1 - xmin1) / self.maxq
- zero1 = torch.round(-xmin1 / scale1) if not self.sym else self.zero
- q = quantize(x, scale1.unsqueeze(1), zero1.unsqueeze(1), self.maxq)
- q -= x
- q.abs_()
- q.pow_(self.norm)
- err = torch.sum(q, 1)
- tmp = err < best
- if torch.any(tmp):
- best[tmp] = err[tmp]
- self.scale[tmp] = scale1[tmp]
- self.zero[tmp] = zero1[tmp]
- if not self.perchannel:
- if weight:
- tmp = shape[0]
- else:
- tmp = shape[1] if len(shape) != 3 else shape[2]
- self.scale = self.scale.repeat(tmp)
- self.zero = self.zero.repeat(tmp)
-
- if weight:
- shape = [-1] + [1] * (len(shape) - 1)
- self.scale = self.scale.reshape(shape)
- self.zero = self.zero.reshape(shape)
- return
- if len(shape) == 4:
- self.scale = self.scale.reshape((1, -1, 1, 1))
- self.zero = self.zero.reshape((1, -1, 1, 1))
- if len(shape) == 3:
- self.scale = self.scale.reshape((1, 1, -1))
- self.zero = self.zero.reshape((1, 1, -1))
- if len(shape) == 2:
- self.scale = self.scale.unsqueeze(0)
- self.zero = self.zero.unsqueeze(0)
-
- def quantize(self, x):
- if self.ready():
- return quantize(x, self.scale, self.zero, self.maxq)
- return x
-
- def enabled(self):
- return self.maxq > 0
-
- def ready(self):
- return torch.all(self.scale != 0)
-
-
-__all__ = ["Quantizer"]
diff --git a/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/repack.py b/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/repack.py
deleted file mode 100644
index 987f05b..0000000
--- a/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/repack.py
+++ /dev/null
@@ -1,99 +0,0 @@
-import torch
-import enum
-from enum import Enum
-from typing import Any, Dict, List, Optional
-from torch.nn.parameter import Parameter
-
-def apply(
- self,
- layer: torch.nn.Module,
- x: torch.Tensor,
- bias: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
- reshaped_x = x.reshape(-1, x.shape[-1])
-
- size_m = reshaped_x.shape[0]
- part_size_n = layer.output_size_per_partition
- part_size_k = layer.input_size_per_partition
- full_size_k = layer.input_size
-
- out_shape = x.shape[:-1] + (part_size_n, )
-
- if layer.marlin_state == GPTQMarlinState.REPACK:
- layer.marlin_state = GPTQMarlinState.READY
-
- # Newly generated tensors need to replace existing tensors that are
- # already registered as parameters by vLLM (and won't be freed)
- def replace_tensor(name, new_t):
- # It is important to use resize_() here since it ensures
- # the same buffer is reused
- getattr(layer, name).resize_(new_t.shape)
- getattr(layer, name).copy_(new_t)
- del new_t
-
- cur_device = layer.qweight.device
-
- # Process act_order
- if self.quant_config.desc_act:
- # Get sorting based on g_idx
- g_idx_sort_indices = torch.argsort(layer.g_idx).to(torch.int)
-
- sorted_g_idx = layer.g_idx[g_idx_sort_indices]
-
- replace_tensor("g_idx", sorted_g_idx)
- replace_tensor("g_idx_sort_indices", g_idx_sort_indices)
-
- else:
- # Reset g_idx related tensors
- layer.g_idx = Parameter(
- torch.empty(0, dtype=torch.int, device=cur_device),
- requires_grad=False,
- )
- layer.g_idx_sort_indices = Parameter(
- torch.empty(0, dtype=torch.int, device=cur_device),
- requires_grad=False,
- )
-
- # Repack weights
- marlin_qweight = ops.gptq_marlin_repack(
- layer.qweight,
- layer.g_idx_sort_indices,
- part_size_k,
- part_size_n,
- self.quant_config.weight_bits,
- )
- replace_tensor("qweight", marlin_qweight)
-
- # Permute scales
- scales_size_k = part_size_k
- scales_size_n = part_size_n
- if self.quant_config.desc_act:
- scales_size_k = full_size_k
-
- marlin_scales = marlin_permute_scales(
- layer.scales,
- scales_size_k,
- scales_size_n,
- self.quant_config.group_size,
- self.quant_config.weight_bits,
- )
- replace_tensor("scales", marlin_scales)
-
- output = ops.gptq_marlin_gemm(
- reshaped_x,
- layer.qweight,
- layer.scales,
- layer.g_idx,
- layer.g_idx_sort_indices,
- layer.workspace,
- self.quant_config.weight_bits,
- size_m,
- part_size_n,
- part_size_k,
- layer.is_k_full,
- )
-
- if bias is not None:
- output.add_(bias) # In-place add
-
- return output.reshape(out_shape)
diff --git a/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/marlin_utils.py b/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/marlin_utils.py
index 7b0398f..accbc00 100644
--- a/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/marlin_utils.py
+++ b/ktransformers/ktransformers_ext/operators/custom_marlin/quantize/utils/marlin_utils.py
@@ -220,7 +220,7 @@ def compute_max_diff(output, output_ref):
class MarlinWorkspace:
- def __init__(self, out_features, min_thread_n, max_parallel):
+ def __init__(self, out_features, min_thread_n, max_parallel, device):
assert (out_features % min_thread_n == 0), (
"out_features = {} is undivisible by min_thread_n = {}".format(
out_features, min_thread_n))
@@ -229,4 +229,4 @@ def __init__(self, out_features, min_thread_n, max_parallel):
self.scratch = torch.zeros(max_workspace_size,
dtype=torch.int,
- device="cuda")
+ device=device)
diff --git a/ktransformers/ktransformers_ext/operators/llamafile/linear.cpp b/ktransformers/ktransformers_ext/operators/llamafile/linear.cpp
index bf1935e..81e5006 100644
--- a/ktransformers/ktransformers_ext/operators/llamafile/linear.cpp
+++ b/ktransformers/ktransformers_ext/operators/llamafile/linear.cpp
@@ -3,7 +3,7 @@
* @Author : chenht2022
* @Date : 2024-07-12 10:07:58
* @Version : 1.0.0
- * @LastEditors : chenht2022
+ * @LastEditors : chenht2022
* @LastEditTime : 2024-07-25 10:34:58
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
**/
@@ -13,9 +13,15 @@ Linear::Linear(LinearConfig config) {
config_ = config;
proj_ = config_.proj;
- input_fp32_.resize(config_.input_size);
- proj_input_.resize(config_.input_size * 4);
- proj_output_.resize(config_.output_size);
+ std::vector> mem_requests;
+ mem_requests.push_back({(void**)&input_fp32_, sizeof(float) * config_.group_max_len * config_.input_size});
+ mem_requests.push_back({(void**)&proj_input_, config_.group_max_len * config_.input_size * ggml_type_size(ggml_internal_get_type_traits(config_.proj_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.proj_type).vec_dot_type)});
+ mem_requests.push_back({(void**)&proj_output_, sizeof(float) * config_.group_max_len * config_.output_size});
+ shared_mem_buffer.alloc(this, mem_requests);
+}
+
+Linear::~Linear() {
+ shared_mem_buffer.dealloc(this);
}
void Linear::warm_up(Backend* backend) {
@@ -26,22 +32,42 @@ void Linear::warm_up(Backend* backend) {
input_fp32[i] = 0;
}
from_float(input_fp32.data(), input.data(), config_.input_size, config_.hidden_type);
- forward(input.data(), output.data(), backend);
+ forward_many(1, input.data(), output.data(), backend);
}
-void Linear::forward(const void* input, void* output, Backend* backend) {
+void Linear::forward_many(int qlen, const void* input, void* output, Backend* backend) {
const void* proj_input_ptr;
if (config_.hidden_type == ggml_internal_get_type_traits(config_.proj_type).vec_dot_type) {
proj_input_ptr = input;
} else {
- to_float(input, input_fp32_.data(), config_.input_size, config_.hidden_type);
- from_float(input_fp32_.data(), proj_input_.data(), config_.input_size, ggml_internal_get_type_traits(config_.proj_type).vec_dot_type);
- proj_input_ptr = proj_input_.data();
+ to_float(input, input_fp32_, qlen * config_.input_size, config_.hidden_type);
+ from_float(input_fp32_, proj_input_, qlen * config_.input_size, ggml_internal_get_type_traits(config_.proj_type).vec_dot_type);
+ proj_input_ptr = proj_input_;
}
int nth = config_.output_size / config_.stride;
backend->do_work_stealing_job(nth, [&](int task_id) {
- int ith = task_id % nth;
- llamafile_sgemm(config_.output_size, 1, config_.input_size / ggml_blck_size(config_.proj_type), proj_, config_.input_size / ggml_blck_size(config_.proj_type), proj_input_ptr, config_.input_size / ggml_blck_size(config_.proj_type), proj_output_.data(), config_.output_size, ith, nth, GGML_TASK_TYPE_COMPUTE, config_.proj_type, ggml_internal_get_type_traits(config_.proj_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT);
+ int ith = task_id;
+ void* proj_ptr = (uint8_t*)proj_ + ith * config_.stride * config_.input_size * ggml_type_size(config_.proj_type) / ggml_blck_size(config_.proj_type);
+ float* proj_output_ptr = proj_output_ + ith * config_.stride;
+ llamafile_sgemm(config_.stride, qlen, config_.input_size / ggml_blck_size(config_.proj_type), proj_ptr, config_.input_size / ggml_blck_size(config_.proj_type), proj_input_ptr, config_.input_size / ggml_blck_size(config_.proj_type), proj_output_ptr, config_.output_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.proj_type, ggml_internal_get_type_traits(config_.proj_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT);
+ if (config_.stride % ggml_blck_size(config_.hidden_type) == 0) {
+ for (int i = 0; i < qlen; i++) {
+ float* output_fp32_ptr = proj_output_ + i * config_.output_size + ith * config_.stride;
+ void* output_ptr = (uint8_t*)output + i * config_.output_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type) + ith * config_.stride * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type);
+ from_float(output_fp32_ptr, output_ptr, config_.stride, config_.hidden_type);
+ }
+ }
});
- from_float(proj_output_.data(), output, config_.output_size, config_.hidden_type);
+ if (config_.stride % ggml_blck_size(config_.hidden_type) != 0) {
+ from_float(proj_output_, output, qlen * config_.output_size, config_.hidden_type);
+ }
+}
+
+void Linear::forward(int qlen, const void* input, void* output, Backend* backend) {
+ if (qlen <= 0) {
+ return;
+ }
+ int forward_len = std::min(qlen, config_.group_max_len);
+ forward_many(forward_len, input, output, backend);
+ forward(qlen - forward_len, (uint8_t*)input + forward_len * config_.input_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), (uint8_t*)output + forward_len * config_.output_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), backend);
}
\ No newline at end of file
diff --git a/ktransformers/ktransformers_ext/operators/llamafile/linear.h b/ktransformers/ktransformers_ext/operators/llamafile/linear.h
index 4285551..fd856f9 100644
--- a/ktransformers/ktransformers_ext/operators/llamafile/linear.h
+++ b/ktransformers/ktransformers_ext/operators/llamafile/linear.h
@@ -3,7 +3,7 @@
* @Author : chenht2022
* @Date : 2024-07-12 10:07:58
* @Version : 1.0.0
- * @LastEditors : chenht2022
+ * @LastEditors : chenht2022
* @LastEditTime : 2024-07-25 10:35:00
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
**/
@@ -22,34 +22,38 @@
#include "llama.cpp/ggml-quants.h"
#include "llama.cpp/ggml.h"
#include "llamafile/sgemm.h"
+#include "shared_mem_buffer.h"
struct LinearConfig {
int input_size;
int output_size;
int stride;
+ int group_max_len;
void* proj;
ggml_type proj_type;
ggml_type hidden_type;
LinearConfig() {}
- LinearConfig(int input_size, int output_size, int stride, void* proj, ggml_type proj_type, ggml_type hidden_type)
- : input_size(input_size), output_size(output_size), stride(stride), proj(proj), proj_type(proj_type), hidden_type(hidden_type) {}
+ LinearConfig(int input_size, int output_size, int stride, int group_max_len, void* proj, ggml_type proj_type, ggml_type hidden_type)
+ : input_size(input_size), output_size(output_size), stride(stride), group_max_len(group_max_len), proj(proj), proj_type(proj_type), hidden_type(hidden_type) {}
};
class Linear {
public:
Linear(LinearConfig);
+ ~Linear();
void warm_up(Backend* backend);
- void forward(const void* input, void* output, Backend* backend);
+ void forward_many(int qlen, const void* input, void* output, Backend* backend);
+ void forward(int qlen, const void* input, void* output, Backend* backend);
private:
LinearConfig config_;
void* proj_; // [output_size * input_size ( /32 if quantized)]
- std::vector input_fp32_; // [input_size]
- std::vector proj_input_; // [input_size * 4]
- std::vector proj_output_; // [output_size]
+ float* input_fp32_; // [group_max_len * input_size]
+ uint8_t* proj_input_; // [group_max_len * input_size * ggml_type_size(ggml_internal_get_type_traits(proj_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(proj_type).vec_dot_type)]
+ float* proj_output_; // [group_max_len * output_size]
};
#endif
\ No newline at end of file
diff --git a/ktransformers/ktransformers_ext/operators/llamafile/mlp.cpp b/ktransformers/ktransformers_ext/operators/llamafile/mlp.cpp
index 979af5e..abad01e 100644
--- a/ktransformers/ktransformers_ext/operators/llamafile/mlp.cpp
+++ b/ktransformers/ktransformers_ext/operators/llamafile/mlp.cpp
@@ -3,7 +3,7 @@
* @Author : chenht2022
* @Date : 2024-07-16 10:43:18
* @Version : 1.0.0
- * @LastEditors : chenht2022
+ * @LastEditors : chenht2022
* @LastEditTime : 2024-07-25 10:35:04
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
**/
@@ -15,14 +15,20 @@ MLP::MLP(MLPConfig config) {
up_proj_ = config_.up_proj;
down_proj_ = config_.down_proj;
- input_fp32_.resize(config_.hidden_size);
- gate_input_.resize(config_.hidden_size * 4);
- up_input_.resize(config_.hidden_size * 4);
- gate_output_.resize(config_.intermediate_size);
- up_output_.resize(config_.intermediate_size);
- intermediate_fp32_.resize(config_.intermediate_size);
- down_input_.resize(config_.intermediate_size * 4);
- down_output_.resize(config_.hidden_size);
+ std::vector> mem_requests;
+ mem_requests.push_back({(void**)&input_fp32_, sizeof(float) * config_.group_max_len * config_.hidden_size});
+ mem_requests.push_back({(void**)&gate_input_, config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)});
+ mem_requests.push_back({(void**)&up_input_, config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)});
+ mem_requests.push_back({(void**)&gate_output_, sizeof(float) * config_.group_max_len * config_.intermediate_size});
+ mem_requests.push_back({(void**)&up_output_, sizeof(float) * config_.group_max_len * config_.intermediate_size});
+ mem_requests.push_back({(void**)&intermediate_fp32_, sizeof(float) * config_.group_max_len * config_.intermediate_size});
+ mem_requests.push_back({(void**)&down_input_, config_.group_max_len * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type)});
+ mem_requests.push_back({(void**)&down_output_, sizeof(float) * config_.group_max_len * config_.hidden_size});
+ shared_mem_buffer.alloc(this, mem_requests);
+}
+
+MLP::~MLP() {
+ shared_mem_buffer.dealloc(this);
}
void MLP::warm_up(Backend* backend) {
@@ -33,33 +39,33 @@ void MLP::warm_up(Backend* backend) {
input_fp32[i] = 0;
}
from_float(input_fp32.data(), input.data(), config_.hidden_size, config_.hidden_type);
- forward(input.data(), output.data(), backend);
+ forward_many(1, input.data(), output.data(), backend);
}
static float act_fn(float x) {
return x / (1.0f + expf(-x));
}
-void MLP::forward(const void* input, void* output, Backend* backend) {
+void MLP::forward_many(int qlen, const void* input, void* output, Backend* backend) {
const void* gate_input_ptr;
const void* up_input_ptr;
if (config_.hidden_type == ggml_internal_get_type_traits(config_.gate_type).vec_dot_type && config_.hidden_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) {
gate_input_ptr = up_input_ptr = input;
} else {
- to_float(input, input_fp32_.data(), config_.hidden_size, config_.hidden_type);
+ to_float(input, input_fp32_, qlen * config_.hidden_size, config_.hidden_type);
if (ggml_internal_get_type_traits(config_.gate_type).vec_dot_type == ggml_internal_get_type_traits(config_.up_type).vec_dot_type) {
- from_float(input_fp32_.data(), gate_input_.data(), config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type);
- gate_input_ptr = up_input_ptr = gate_input_.data();
+ from_float(input_fp32_, gate_input_, qlen * config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type);
+ gate_input_ptr = up_input_ptr = gate_input_;
} else {
if (config_.hidden_type != ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) {
- from_float(input_fp32_.data(), gate_input_.data(), config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type);
- gate_input_ptr = gate_input_.data();
+ from_float(input_fp32_, gate_input_, qlen * config_.hidden_size, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type);
+ gate_input_ptr = gate_input_;
} else {
gate_input_ptr = input;
}
if (config_.hidden_type != ggml_internal_get_type_traits(config_.up_type).vec_dot_type) {
- from_float(input_fp32_.data(), up_input_.data(), config_.hidden_size, ggml_internal_get_type_traits(config_.up_type).vec_dot_type);
- up_input_ptr = up_input_.data();
+ from_float(input_fp32_, up_input_, qlen * config_.hidden_size, ggml_internal_get_type_traits(config_.up_type).vec_dot_type);
+ up_input_ptr = up_input_;
} else {
up_input_ptr = input;
}
@@ -69,35 +75,49 @@ void MLP::forward(const void* input, void* output, Backend* backend) {
backend->do_work_stealing_job(nth, [&](int task_id) {
int ith = task_id;
void* gate_proj_ptr = (uint8_t*)gate_proj_ + ith * config_.stride * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type);
- float* gate_output_ptr = gate_output_.data() + ith * config_.stride;
- llamafile_sgemm(config_.stride, 1, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_proj_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_input_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_output_ptr, config_.stride, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.gate_type, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT);
+ float* gate_output_ptr = gate_output_ + ith * config_.stride;
+ llamafile_sgemm(config_.stride, qlen, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_proj_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_input_ptr, config_.hidden_size / ggml_blck_size(config_.gate_type), gate_output_ptr, config_.intermediate_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.gate_type, ggml_internal_get_type_traits(config_.gate_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT);
void* up_proj_ptr = (uint8_t*)up_proj_ + ith * config_.stride * config_.hidden_size * ggml_type_size(config_.up_type) / ggml_blck_size(config_.up_type);
- float* up_output_ptr = up_output_.data() + ith * config_.stride;
- llamafile_sgemm(config_.stride, 1, config_.hidden_size / ggml_blck_size(config_.up_type), up_proj_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_input_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_output_ptr, config_.stride, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.up_type, ggml_internal_get_type_traits(config_.up_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT);
- for (int i = ith * config_.stride; i < (ith + 1) * config_.stride; i++) {
- intermediate_fp32_[i] = act_fn(gate_output_[i]) * up_output_[i];
- }
- if (config_.stride % ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) == 0) {
- float* intermediate_fp32_ptr = intermediate_fp32_.data() + ith * config_.stride;
- void* down_input_ptr = (uint8_t*)down_input_.data() + ith * config_.stride * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type);
- from_float(intermediate_fp32_ptr, down_input_ptr, config_.stride, ggml_internal_get_type_traits(config_.down_type).vec_dot_type);
+ float* up_output_ptr = up_output_ + ith * config_.stride;
+ llamafile_sgemm(config_.stride, qlen, config_.hidden_size / ggml_blck_size(config_.up_type), up_proj_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_input_ptr, config_.hidden_size / ggml_blck_size(config_.up_type), up_output_ptr, config_.intermediate_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.up_type, ggml_internal_get_type_traits(config_.up_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT);
+ for (int i = 0; i < qlen; i++) {
+ for (int j = ith * config_.stride; j < (ith + 1) * config_.stride; j++) {
+ intermediate_fp32_[i * config_.intermediate_size + j] = act_fn(gate_output_[i * config_.intermediate_size + j]) * up_output_[i * config_.intermediate_size + j];
+ }
+ if (config_.stride % ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) == 0) {
+ float* intermediate_fp32_ptr = intermediate_fp32_ + i * config_.intermediate_size + ith * config_.stride;
+ void* down_input_ptr = (uint8_t*)down_input_ + i * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) + ith * config_.stride * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type);
+ from_float(intermediate_fp32_ptr, down_input_ptr, config_.stride, ggml_internal_get_type_traits(config_.down_type).vec_dot_type);
+ }
}
});
if (config_.stride % ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) != 0) {
- from_float(intermediate_fp32_.data(), down_input_.data(), config_.intermediate_size, ggml_internal_get_type_traits(config_.down_type).vec_dot_type);
+ from_float(intermediate_fp32_, down_input_, qlen * config_.intermediate_size, ggml_internal_get_type_traits(config_.down_type).vec_dot_type);
}
nth = config_.hidden_size / config_.stride;
backend->do_work_stealing_job(nth, [&](int task_id) {
int ith = task_id;
void* down_proj_ptr = (uint8_t*)down_proj_ + ith * config_.stride * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type);
- float* down_output_ptr = down_output_.data() + ith * config_.stride;
- llamafile_sgemm(config_.stride, 1, config_.intermediate_size / ggml_blck_size(config_.down_type), down_proj_ptr, config_.intermediate_size / ggml_blck_size(config_.down_type), down_input_.data(), config_.intermediate_size / ggml_blck_size(config_.down_type), down_output_ptr, config_.stride, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.down_type, ggml_internal_get_type_traits(config_.down_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT);
+ float* down_output_ptr = down_output_ + ith * config_.stride;
+ llamafile_sgemm(config_.stride, qlen, config_.intermediate_size / ggml_blck_size(config_.down_type), down_proj_ptr, config_.intermediate_size / ggml_blck_size(config_.down_type), down_input_, config_.intermediate_size / ggml_blck_size(config_.down_type), down_output_ptr, config_.hidden_size, 0, 1, GGML_TASK_TYPE_COMPUTE, config_.down_type, ggml_internal_get_type_traits(config_.down_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT);
if (config_.stride % ggml_blck_size(config_.hidden_type) == 0) {
- void* output_ptr = (uint8_t*)output + ith * config_.stride * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type);
- from_float(down_output_ptr, output_ptr, config_.stride, config_.hidden_type);
+ for (int i = 0; i < qlen; i++) {
+ float* output_fp32_ptr = down_output_ + i * config_.hidden_size + ith * config_.stride;
+ void* output_ptr = (uint8_t*)output + i * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type) + ith * config_.stride * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type);
+ from_float(output_fp32_ptr, output_ptr, config_.stride, config_.hidden_type);
+ }
}
});
if (config_.stride % ggml_blck_size(config_.hidden_type) != 0) {
- from_float(down_output_.data(), output, config_.hidden_size, config_.hidden_type);
+ from_float(down_output_, output, qlen * config_.hidden_size, config_.hidden_type);
}
}
+
+void MLP::forward(int qlen, const void* input, void* output, Backend* backend) {
+ if (qlen <= 0) {
+ return;
+ }
+ int forward_len = std::min(qlen, config_.group_max_len);
+ forward_many(forward_len, input, output, backend);
+ forward(qlen - forward_len, (uint8_t*)input + forward_len * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), (uint8_t*)output + forward_len * config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type), backend);
+}
\ No newline at end of file
diff --git a/ktransformers/ktransformers_ext/operators/llamafile/mlp.h b/ktransformers/ktransformers_ext/operators/llamafile/mlp.h
index 604db77..eb93294 100644
--- a/ktransformers/ktransformers_ext/operators/llamafile/mlp.h
+++ b/ktransformers/ktransformers_ext/operators/llamafile/mlp.h
@@ -3,7 +3,7 @@
* @Author : chenht2022
* @Date : 2024-07-12 10:07:58
* @Version : 1.0.0
- * @LastEditors : chenht2022
+ * @LastEditors : chenht2022
* @LastEditTime : 2024-07-25 10:35:06
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
**/
@@ -22,11 +22,13 @@
#include "llama.cpp/ggml-quants.h"
#include "llama.cpp/ggml.h"
#include "llamafile/sgemm.h"
+#include "shared_mem_buffer.h"
struct MLPConfig {
int hidden_size;
int intermediate_size;
int stride;
+ int group_max_len;
void* gate_proj;
void* up_proj;
void* down_proj;
@@ -37,15 +39,17 @@ struct MLPConfig {
MLPConfig() {}
- MLPConfig(int hidden_size, int intermediate_size, int stride, void* gate_proj, void* up_proj, void* down_proj, ggml_type gate_type, ggml_type up_type, ggml_type down_type, ggml_type hidden_type)
- : hidden_size(hidden_size), intermediate_size(intermediate_size), stride(stride), gate_proj(gate_proj), up_proj(up_proj), down_proj(down_proj), gate_type(gate_type), up_type(up_type), down_type(down_type), hidden_type(hidden_type) {}
+ MLPConfig(int hidden_size, int intermediate_size, int stride, int group_max_len, void* gate_proj, void* up_proj, void* down_proj, ggml_type gate_type, ggml_type up_type, ggml_type down_type, ggml_type hidden_type)
+ : hidden_size(hidden_size), intermediate_size(intermediate_size), stride(stride), group_max_len(group_max_len), gate_proj(gate_proj), up_proj(up_proj), down_proj(down_proj), gate_type(gate_type), up_type(up_type), down_type(down_type), hidden_type(hidden_type) {}
};
class MLP {
public:
MLP(MLPConfig);
+ ~MLP();
void warm_up(Backend* backend);
- void forward(const void* input, void* output, Backend* backend);
+ void forward_many(int qlen, const void* input, void* output, Backend* backend);
+ void forward(int qlen, const void* input, void* output, Backend* backend);
private:
MLPConfig config_;
@@ -53,14 +57,14 @@ class MLP {
void* up_proj_; // [intermediate_size * hidden_size ( /32 if quantized)]
void* down_proj_; // [hidden_size * intermediate_size ( /32 if quantized)]
- std::vector input_fp32_; // [hidden_size]
- std::vector gate_input_; // [hidden_size * 4]
- std::vector up_input_; // [hidden_size * 4]
- std::vector gate_output_; // [intermediate_size]
- std::vector up_output_; // [intermediate_size]
- std::vector intermediate_fp32_; // [intermediate_size]
- std::vector down_input_; // [intermediate_size * 4]
- std::vector down_output_; // [hidden_size]
+ float* input_fp32_; // [group_max_len * hidden_size]
+ uint8_t* gate_input_; // [group_max_len * hidden_size * ggml_type_size(ggml_internal_get_type_traits(gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(gate_type).vec_dot_type)]
+ uint8_t* up_input_; // [group_max_len * hidden_size * ggml_type_size(ggml_internal_get_type_traits(up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(up_type).vec_dot_type)]
+ float* gate_output_; // [group_max_len * intermediate_size]
+ float* up_output_; // [group_max_len * intermediate_size]
+ float* intermediate_fp32_; // [group_max_len * intermediate_size]
+ uint8_t* down_input_; // [group_max_len * intermediate_size * ggml_type_size(ggml_internal_get_type_traits(down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(down_type).vec_dot_type)]
+ float* down_output_; // [group_max_len * hidden_size]
};
#endif
\ No newline at end of file
diff --git a/ktransformers/ktransformers_ext/operators/llamafile/moe.cpp b/ktransformers/ktransformers_ext/operators/llamafile/moe.cpp
index 3a5c852..d75db65 100644
--- a/ktransformers/ktransformers_ext/operators/llamafile/moe.cpp
+++ b/ktransformers/ktransformers_ext/operators/llamafile/moe.cpp
@@ -1,97 +1,62 @@
/**
- * @Description :
+ * @Description :
* @Author : chenht2022
* @Date : 2024-07-22 02:03:22
* @Version : 1.0.0
- * @LastEditors : chenht2022
+ * @LastEditors : chenht2022
* @LastEditTime : 2024-07-25 10:35:07
- * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
-**/
+ * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
+ **/
#include "moe.h"
#include
#include
-uint8_t* MOE::buffer_ = nullptr;
-
MOE::MOE(MOEConfig config) {
config_ = config;
gate_proj_ = config_.gate_proj;
up_proj_ = config_.up_proj;
down_proj_ = config_.down_proj;
- if (MOE::buffer_ == nullptr) {
- uint64_t buffer_size = 0;
- buffer_size += sizeof(float) * config_.group_max_len * config_.hidden_size;
- buffer_size += config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type);
- buffer_size += config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type);
- buffer_size += config_.routed_expert_num * config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type);
- buffer_size += config_.routed_expert_num * config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type);
- buffer_size += sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size;
- buffer_size += sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size;
- buffer_size += sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size;
- buffer_size += config_.routed_expert_num * config_.group_max_len * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type);
- buffer_size += sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.hidden_size;
- buffer_size += sizeof(float) * config_.group_max_len * config_.hidden_size;
- buffer_ = (uint8_t*)malloc(buffer_size);
- }
-
- uint64_t offset = 0;
- s_input_fp32_ = (float*)(buffer_ + offset);
- offset += sizeof(float) * config_.hidden_size;
- s_gate_input_ = (uint8_t*)(buffer_ + offset);
- offset += config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type);
- s_up_input_ = (uint8_t*)(buffer_ + offset);
- offset += config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type);
+ std::vector> s_mem_requests;
+ s_mem_requests.push_back({(void**)&s_input_fp32_, sizeof(float) * config_.hidden_size});
+ s_mem_requests.push_back({(void**)&s_gate_input_, config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)});
+ s_mem_requests.push_back({(void**)&s_up_input_, config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)});
s_gate_output_.resize(config_.routed_expert_num);
s_up_output_.resize(config_.routed_expert_num);
s_intermediate_fp32_.resize(config_.routed_expert_num);
s_down_input_.resize(config_.routed_expert_num);
s_down_output_.resize(config_.routed_expert_num);
for (int i = 0; i < config_.routed_expert_num; i++) {
- s_gate_output_[i] = (float*)(buffer_ + offset);
- offset += sizeof(float) * config_.intermediate_size;
- s_up_output_[i] = (float*)(buffer_ + offset);
- offset += sizeof(float) * config_.intermediate_size;
- s_intermediate_fp32_[i] = (float*)(buffer_ + offset);
- offset += sizeof(float) * config_.intermediate_size;
- s_down_input_[i] = (uint8_t*)(buffer_ + offset);
- offset += config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type);
- s_down_output_[i] = (float*)(buffer_ + offset);
- offset += sizeof(float) * config_.hidden_size;
+ s_mem_requests.push_back({(void**)&s_gate_output_[i], sizeof(float) * config_.intermediate_size});
+ s_mem_requests.push_back({(void**)&s_up_output_[i], sizeof(float) * config_.intermediate_size});
+ s_mem_requests.push_back({(void**)&s_intermediate_fp32_[i], sizeof(float) * config_.intermediate_size});
+ s_mem_requests.push_back({(void**)&s_down_input_[i], config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type)});
+ s_mem_requests.push_back({(void**)&s_down_output_[i], sizeof(float) * config_.hidden_size});
}
- s_output_fp32_ = (float*)(buffer_ + offset);
+ s_mem_requests.push_back({(void**)&s_output_fp32_, sizeof(float) * config_.hidden_size});
+ shared_mem_buffer.alloc(this, s_mem_requests);
- offset = 0;
+ std::vector> m_mem_requests;
m_input_fp32_.resize(config_.group_max_len);
m_gate_input_.resize(config_.group_max_len);
m_up_input_.resize(config_.group_max_len);
for (int i = 0; i < config_.group_max_len; i++) {
- m_input_fp32_[i] = (float*)(buffer_ + offset);
- offset += sizeof(float) * config_.hidden_size;
- m_gate_input_[i] = (uint8_t*)(buffer_ + offset);
- offset += config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type);
- m_up_input_[i] = (uint8_t*)(buffer_ + offset);
- offset += config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type);
+ m_mem_requests.push_back({(void**)&m_input_fp32_[i], sizeof(float) * config_.hidden_size});
+ m_mem_requests.push_back({(void**)&m_gate_input_[i], config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)});
+ m_mem_requests.push_back({(void**)&m_up_input_[i], config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)});
}
- m_local_gate_input_ = (uint8_t*)(buffer_ + offset);
- offset += config_.routed_expert_num * config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type);
- m_local_up_input_ = (uint8_t*)(buffer_ + offset);
- offset += config_.routed_expert_num * config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type);
- m_local_gate_output_ = (float*)(buffer_ + offset);
- offset += sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size;
- m_local_up_output_ = (float*)(buffer_ + offset);
- offset += sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size;
- m_local_intermediate_fp32_ = (float*)(buffer_ + offset);
- offset += sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size;
- m_local_down_input_ = (uint8_t*)(buffer_ + offset);
- offset += config_.routed_expert_num * config_.group_max_len * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type);
- m_local_down_output_ = (float*)(buffer_ + offset);
- offset += sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.hidden_size;
+ m_mem_requests.push_back({(void**)&m_local_gate_input_, config_.routed_expert_num * config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.gate_type).vec_dot_type)});
+ m_mem_requests.push_back({(void**)&m_local_up_input_, config_.routed_expert_num * config_.group_max_len * config_.hidden_size * ggml_type_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.up_type).vec_dot_type)});
+ m_mem_requests.push_back({(void**)&m_local_gate_output_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size});
+ m_mem_requests.push_back({(void**)&m_local_up_output_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size});
+ m_mem_requests.push_back({(void**)&m_local_intermediate_fp32_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.intermediate_size});
+ m_mem_requests.push_back({(void**)&m_local_down_input_, config_.routed_expert_num * config_.group_max_len * config_.intermediate_size * ggml_type_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type) / ggml_blck_size(ggml_internal_get_type_traits(config_.down_type).vec_dot_type)});
+ m_mem_requests.push_back({(void**)&m_local_down_output_, sizeof(float) * config_.routed_expert_num * config_.group_max_len * config_.hidden_size});
m_output_fp32_.resize(config_.group_max_len);
for (int i = 0; i < config_.group_max_len; i++) {
- m_output_fp32_[i] = (float*)(buffer_ + offset);
- offset += sizeof(float) * config_.hidden_size;
+ m_mem_requests.push_back({(void**)&m_output_fp32_[i], sizeof(float) * config_.hidden_size});
}
+ shared_mem_buffer.alloc(this, m_mem_requests);
m_local_pos_.resize(config_.group_max_len);
for (int i = 0; i < config_.group_max_len; i++) {
@@ -107,6 +72,10 @@ MOE::MOE(MOEConfig config) {
m_local_down_output_ptr_.resize(config_.expert_num);
}
+MOE::~MOE() {
+ shared_mem_buffer.dealloc(this);
+}
+
void MOE::warm_up(Backend* backend) {
std::vector input_fp32(config_.hidden_size);
std::vector input(config_.hidden_size * ggml_type_size(config_.hidden_type) / ggml_blck_size(config_.hidden_type));
diff --git a/ktransformers/ktransformers_ext/operators/llamafile/moe.h b/ktransformers/ktransformers_ext/operators/llamafile/moe.h
index 588a829..a1470aa 100644
--- a/ktransformers/ktransformers_ext/operators/llamafile/moe.h
+++ b/ktransformers/ktransformers_ext/operators/llamafile/moe.h
@@ -3,7 +3,7 @@
* @Author : chenht2022
* @Date : 2024-07-22 02:03:22
* @Version : 1.0.0
- * @LastEditors : chenht2022
+ * @LastEditors : chenht2022
* @LastEditTime : 2024-07-25 10:35:10
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
**/
@@ -22,6 +22,7 @@
#include "llama.cpp/ggml-quants.h"
#include "llama.cpp/ggml.h"
#include "llamafile/sgemm.h"
+#include "shared_mem_buffer.h"
struct MOEConfig {
int expert_num;
@@ -48,13 +49,13 @@ struct MOEConfig {
class MOE {
public:
MOE(MOEConfig);
+ ~MOE();
void warm_up(Backend* backend);
void forward_one(int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend);
void forward_many(int qlen, int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend);
void forward(int qlen, int k, const uint64_t* expert_ids, const float* weights, const void* input, void* output, Backend* backend);
private:
- static uint8_t* buffer_;
MOEConfig config_;
void* gate_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)]
void* up_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)]
diff --git a/ktransformers/ktransformers_ext/operators/llamafile/shared_mem_buffer.cpp b/ktransformers/ktransformers_ext/operators/llamafile/shared_mem_buffer.cpp
new file mode 100644
index 0000000..dc2d65d
--- /dev/null
+++ b/ktransformers/ktransformers_ext/operators/llamafile/shared_mem_buffer.cpp
@@ -0,0 +1,55 @@
+/**
+ * @Description :
+ * @Author : chenht2022
+ * @Date : 2024-08-05 04:49:08
+ * @Version : 1.0.0
+ * @LastEditors : chenht2022
+ * @LastEditTime : 2024-08-05 09:21:29
+ * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
+ **/
+#include "shared_mem_buffer.h"
+#include
+
+SharedMemBuffer::SharedMemBuffer() {
+ buffer_ = nullptr;
+ size_ = 0;
+}
+
+SharedMemBuffer::~SharedMemBuffer() {
+ if (buffer_) {
+ free(buffer_);
+ }
+}
+
+void SharedMemBuffer::alloc(void* object, std::vector> requests) {
+ uint64_t size = 0;
+ for (auto& request : requests) {
+ size += request.second;
+ }
+ if (size > size_) {
+ if (buffer_) {
+ free(buffer_);
+ }
+ buffer_ = malloc(size);
+ size_ = size;
+ for (auto& obj_requests : hist_requests_) {
+ for (auto& requests : obj_requests.second) {
+ arrange(requests);
+ }
+ }
+ }
+ arrange(requests);
+ hist_requests_[object].push_back(requests);
+}
+
+void SharedMemBuffer::dealloc(void* object) {
+ hist_requests_.erase(object);
+}
+
+void SharedMemBuffer::arrange(std::vector> requests) {
+ uint64_t offset = 0;
+ for (auto& request : requests) {
+ *(request.first) = (uint8_t*)buffer_ + offset;
+ offset += request.second;
+ }
+}
diff --git a/ktransformers/ktransformers_ext/operators/llamafile/shared_mem_buffer.h b/ktransformers/ktransformers_ext/operators/llamafile/shared_mem_buffer.h
new file mode 100644
index 0000000..eeaccd4
--- /dev/null
+++ b/ktransformers/ktransformers_ext/operators/llamafile/shared_mem_buffer.h
@@ -0,0 +1,37 @@
+/**
+ * @Description :
+ * @Author : chenht2022
+ * @Date : 2024-08-05 04:49:08
+ * @Version : 1.0.0
+ * @LastEditors : chenht2022
+ * @LastEditTime : 2024-08-05 06:36:41
+ * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
+ **/
+
+#ifndef CPUINFER_SHAREDMEMBUFFER_H
+#define CPUINFER_SHAREDMEMBUFFER_H
+
+#include
+#include
+#include