vllm 0.17 cuda: adding quack-kernels, troch-c-dlpack-ext

Author: d-goldin

pkgs/development/python-modules/apache-tvm-ffi/default.nix

@@ -9,6 +9,8 @@
   ninja,
   scikit-build-core,
   setuptools-scm,
+  python,
+  build,
 
   # dependencies
   typing-extensions,
@@ -43,6 +45,10 @@ buildPythonPackage (finalAttrs: {
 
   dependencies = [
     typing-extensions
+    # for the extension we technically need torch, but this
+    # causes some circular import issue during build, so the
+    # user/dependant needs to ensure that for now
+    # torch
   ];
 
   optional-dependencies = {
@@ -59,6 +65,27 @@ buildPythonPackage (finalAttrs: {
     writableTmpDirAsHomeHook
   ];
 
+  # Building the extension already requires an installed
+  # tvm-ffi, so we're just doing this "temporarily".
+  # Would still be nice to break out the c-ext separately...
+  # Maybe multiple outputs derivation?
+  buildPhase = ''
+    pypaBuildPhase
+
+    # pre-installing wheel for extension build
+    TMP_INSTALL=$(mktemp -d)
+    mkdir -p $TMP_INSTALL
+    ${python}/bin/python -m installer --prefix "$TMP_INSTALL" dist/*.whl;
+    pushd addons/torch_c_dlpack_ext
+    PYTHONPATH="$TMP_INSTALL/${python.sitePackages}:$PYTHONPATH" \
+      ${build}/bin/pyproject-build --no-isolation --outdir ../../dist/ --wheel
+    popd
+  '';
+
+  # needed so we can work around the torch<->ext build circular
+  # module init
+  dontCheckRuntimeDeps = true;
+
   meta = {
     description = "Open ABI and FFI for Machine Learning Systems";
     changelog = "https://github.com/apache/tvm-ffi/releases/tag/${finalAttrs.src.tag}";

pkgs/development/python-modules/flashinfer/default.nix

@@ -85,6 +85,8 @@ buildPythonPackage (finalAttrs: {
 
   pythonRemoveDeps = [
     "nvidia-cudnn-frontend"
+    # dsl is used in quack-kernels, so could be
+    # made available separately
     "nvidia-cutlass-dsl"
   ];
   dependencies = [

pkgs/development/python-modules/quack-kernels/default.nix

@@ -0,0 +1,84 @@
+{
+  fetchFromGitHub,
+  buildPythonPackage,
+  packaging,
+  setuptools,
+  setuptools-scm,
+  torch,
+  cudaPackages,
+  apache-tvm-ffi,
+  lib,
+  ...
+}:
+let
+  # Not clear if anyone else really needs this
+  cutlass-dsl = buildPythonPackage rec {
+    pname = "nvidia-cutlass-dsl";
+    # TODO: Should this be aligned with the actual
+    # cutlass version? Seems to work fine for now.
+    version = "4.4.2";
+    pyproject = true;
+
+    src = fetchFromGitHub {
+      owner = "NVIDIA";
+      repo = "cutlass";
+      tag = "v${version}";
+      hash = "sha256-0q9Ad0Z6E/rO2PdM4uQc8H0E0qs9uKc3reHepiHhjEc=";
+    };
+
+    patchPhase = ''
+      echo ${version} > VERSION.EDITABLE
+    '';
+
+    sourceRoot = "${src.name}/python/CuTeDSL";
+
+    stdenv = torch.stdenv;
+    dependencies = [
+      torch
+      cudaPackages.cutlass
+    ];
+    build-system = [
+      setuptools-scm
+    ];
+  };
+
+in
+
+buildPythonPackage rec {
+  pname = "quack-kernels";
+  version = "0.3.4";
+  pyproject = true;
+
+  stdenv = torch.stdenv;
+  build-system = [
+    packaging
+    setuptools
+    setuptools-scm
+    torch
+  ];
+
+  dependencies = [
+    cudaPackages.cutlass
+    torch
+    #tvm-ffi comes with torch-c-dlpack_ext
+    apache-tvm-ffi
+    cutlass-dsl
+  ];
+
+  src = fetchFromGitHub {
+    # https://github.com/Dao-AILab/quack/tree/v0.3.4
+    name = "quack-kernels";
+    owner = "DAO-AILab";
+    repo = "quack";
+    rev = "v${version}";
+    hash = "sha256-CjaUe+gBHfWkdnxn+jiXq9s7ODjveptJXdjbiBywGjA=";
+  };
+
+  meta = {
+    description = "A Quirky Assortment of CuTe Kernels";
+    homepage = "https://github.com/Dao-AILab/quack";
+    license = lib.licenses.asl20;
+    maintainers = [
+    ];
+  };
+}

pkgs/development/python-modules/vllm/default.nix

@@ -93,6 +93,7 @@
   cupy,
   flashinfer,
   nvidia-ml-py,
+  quack-kernels,
 
   # optional-dependencies
   # audio
@@ -390,7 +391,9 @@ buildPythonPackage.override { stdenv = torch.stdenv; } (finalAttrs: {
   ]
   ++ lib.optionals cudaSupport [
     cudaPackages.cuda_nvcc
+    cudaPackages.cutlass
     autoAddDriverRunpath
+    quack-kernels
   ]
   ++ lib.optionals isCudaJetson [
     cudaPackages.autoAddCudaCompatRunpath

pkgs/top-level/python-packages.nix

@@ -16318,6 +16318,8 @@ self: super: with self; {
 
   qtpy = callPackage ../development/python-modules/qtpy { };
 
+  quack-kernels = callPackage ../development/python-modules/quack-kernels { };
+
   quadprog = callPackage ../development/python-modules/quadprog { };
 
   qualysclient = callPackage ../development/python-modules/qualysclient { };