[RFC]: Kernel Library Restructure / Packaging Split (addressing long build times)

[seemethere]

Motivation.

Motivation

vLLM local builds take too long.

On CI the wall clock time for a build (with no cache) exceeds 5 hours. For the repository it’s noted that most of the commits to the vllm directory did not modify the csrc directory, effectively making them python only changes and not requiring extra build time (95% within the last year).

It’s also acknowledged in the install from source docs about how long the build time is with most people being pointed towards python only builds as a means to get a reasonable build time to do most of their work. (link)

As well, I think it’s worth it to consider how we structure the code to make it more clear over what is kernel specific vs. what is engine specific, which could benefit the project by clearly identifying when we would need to run things like performance testing vs. unit testing (i.e. kernel vs. engine).

Proposed Change.

The vLLM package is currently made up of two pretty distinct paths:

• csrc: Kernels / PyTorch Custom Operators
  • To also include vllm-flash-attn source builds
• vllm: Engine code that sometimes utilizes things from csrc

My proposal is split up csrc / vllm packages into two separate python packages with matching directory structures (names are flexible here):

• vllm: The current package that users install with a dependency on vllm-ops
• vllm-kernels: Made up of vLLM Kernels + Custom Ops + dependencies like vllm-flash-attn

Potential directory structure:

vllm/
├─ vllm-engine/            <-- Current Package
│  ├─ vllm/
│  ├─ pyproject.toml
├─ vllm-kernels/           <-- New package 
│  ├─ vllm_kernels/        <-- Python frontend
│  ├─ csrc/                <-- Kernel code
│  ├─ pyproject.toml

The changes will follow into two steps:

• Step 1: refactor as two packages,
• Step 2: polish the kernel APIs, and separate the libraries based on the backend type.

Along with this change I recommend that we restructure kernels and the main engine code into clear separate directories for better code organization over which directories build which packages.

Feedback Period.

No formal feedback period! I'd like to start this work ASAP while my time isn't allocated to another project so ideally within the next month or so!

CC List.

cc @simon-mo, @tlrmchlsmth, @youkaichao, @comaniac

Any Other Things.

Other potential benefits

Package size could be spread across 2 packages

We all understand how much it sucks dealing with space limits on PyPI, by effectively splitting our package size across two different packages on PyPI we can increase the number of releases we can do in the future.

Decoupling release cycles

*Potentially *splitting out the release cadence of the kernels and the engine, rolling out kernel updates without having to do a full engine upgrade and vice versa. This can make it so that you won’t even need to do an update to the engine and still roll out things like performance upgrades.

Backends could be more swappable

Another potential idea for this is that we can make the vllm-kernels package swappable between different accelerator types (cpu, rocm, cuda, etc.) and make it so that we can actually publish accelerator backends (like rocm) onto pypi and have everything installable from there. This also opens up the potential for vendors to manage their own backend packages themselves, freeing up vllm core maintainers to focus on just the core set of backends.

Some caveats

Note

I’m not necessarily proposing that we have other people depend on vllm-kernels but just that we split the packaging up for easier organization / caching for both our local developers as well as for CI

Ensuring compatibility between the two packages could be a challenge

Having two separate packages does open us up to issues where users could install incompatible versions of the packages together.

Example:

Imagine a user has vllm-kernels==1.0.0 while also having vllm==2.0.0 installed, these could be wildly different and should be incompatible.

Potential solution(s):

Hard-lock dependencies

We could also hard lock the engine package to the kernels as well by just explicitly making each version of vllm depend on the requisite vllm-kernels package (i.e. vllm==2.0.0 depend on vllm-kernels==2.0.0)

Soft-lock dependencies

We could just make the vLLM package depend on a specific minor version of the vllm-kernels package. (i.e. have vllm==2.0.0 depend on the vllm-kernels>=2.0.0,<=2.1.0)

Open Discussion Points

• Directory / project structure: How should we represent the two packages within the repository?
  • Does vllm-kernels potentially move to a separate repository?
  • Do we feel interested in maintaining a top level repository structure to represent the two packages (i.e. having top-level vllm-engine, vllm-kernels directories?)?
• Compatibility: Do we always want to release these packages together? Do we think that we should de-couple them for some of the benefits listed above?

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[hmellor]

Instead of creating a new Python package for C++ kernels might I suggest the Hugging Face 🤗 solution?

It would looks something like:

• Creating a vllm-kernels directory at the root of the repo

• Create a directory structure that looks something like:

  vllm-kernels
  ├── build.toml
  ├── relu_kernel
  │   └── relu.cu
  └── torch-ext
      └── torch_bindings.cpp
      └── torch_bindings.h
      └── relu
          └── __init__.py
  

• This can then be built in the Docker container provided by https://github.com/huggingface/kernel-builder with a simple command like docker run --rm -v vllm-kernels:/kernelcode ghcr.io/huggingface/kernel-builder:latest

• The results could then be uploaded to the Hugging Face Hub (like this one for example which contains the vLLM quantization kernels https://huggingface.co/kernels-community/quantization)

• vLLM can then consume the kernels using https://github.com/huggingface/kernels

  activation = get_kernel("kernels-community/activation")
  x = torch.randn((10, 10), dtype=torch.float16, device="cuda")
  y = torch.empty_like(x)
  activation.gelu_fast(y, x)
  

As you can see by taking a look around https://huggingface.co/kernels-community, we have already done this for many families of kernels from vLLM so I'm confident we could make it work for an official vllm-kernels!

cc @LucasWilkinson

[hmellor]

Oh and as a bonus, we would get:

• download statistics for bragging rights
• traceability when other packages use our kernels (better than copy/paste with a comment saying where it came from)

[seemethere]

activation = get_kernel("kernels-community/activation")
x = torch.randn((10, 10), dtype=torch.float16, device="cuda")
y = torch.empty_like(x)
activation.gelu_fast(y, x)

This seems interesting and I can see benefits from both an ease of implementation / binary size side but do we think that doing this might introduce some potential flakiness if we rely on network calls to huggingface.co for grabbing kernels?

I'd ideally like to make it so that we reduce the amount of network calls we have to make (since most of the models are actually downloaded already).

But I do see the value in this approach and am curious if others also see this too.

[hmellor]

potential flakiness if we rely on network calls to huggingface.co for grabbing kernels?

• For Docker images - we can download and bundle the binaries at build time (just as we do for other dependencies).
• For the Python install use case - these kernel binaries are downloaded and cached, just like models.
• Also, vLLM already makes many calls to huggingface.co to fetch configs and model weights, so we're not introducing an unknown.

Either way, a user still has to download the kernel binaries from somewhere; whether it be bundled in a wheel from PyPI or via get_kernel from huggingface.co.

[LucasWilkinson]

Thanks for the RFC, @seemethere. @hmellor and I discussed leveraging Hugging Face’s kernel builder, but that unfortunately fell on the back burner. As a result I do have some thoughts, though.

Does vllm-kernels potentially move to a separate repository?

Based on my experience with vllm-flash-attention, I’m fairly opposed. Kernel changes are frequently paired with vLLM-side changes, so you end up coordinating a vLLM PR and a vllm-kernels PR. If you have change A (vLLM-A + vllm-kernels-A) and change B (vLLM-B + vllm-kernels-B), you can easily end up getting blocked or end up in weird situations. For example say vllm-kernels-A lands but vLLM-A is held up in review. Then vllm-kernels-B lands, this results in vLLM-B being blocked on vLLM-A since there is no way of picking up vllm-kernels-B independent of vllm-kernels-A. Not to mention its much nicer for the reviewers to be able tor review the kernel changes and vLLM side changes in single PR with a unified UI/discussion. For vllm-flash-attention managing PRs in this way is somewhat manageable (albiet annoying) due the low number of PRs generally, however I fear this would become much harder to maintain for a general kernel library with a higher volume of PRs. HF kernels could help a bit here but making it possible to version smaller sub-groups of kernels but I think the general problem still exists. For this reason Im very pro sticking with a monolithic

Do we always want to release these packages together? Should we decouple them for the benefits above?

I think we should decouple them to make the Python-only install cleaner. Ideally, we’d publish release-only wheels paired with each vLLM PyPI release (to avoid wasted storage). It’d also be nice to have a private PyPI—like our nightly index—that automatically pushes a new kernels wheel whenever vllm-kernels is updated, to be used by the python only install, and automatically hard-locking to last commit that updated vllm-kernels. I’m not sure how that would work in CI; maybe @khluu has thoughts.

Instead of creating a new Python package for C++ kernels, might we use the Hugging Face solution?

Re: HF kernels (repeating some things from @hmellor and I's chat) I strongly believe we should still ship standalone wheels that don’t require network calls. Currenlty this lets users run local models on air-gapped systems and simplifies standalone containers. I think we’d need broad buy-in befor such a change. I suspect we could still use the HF kernels API and still do this. Under the hood of get_kernel we could load a pre-compiled shared lib in the wheel or call out to a JIT. I don't understand HF kernels well enough yet to know how possible this is (@hmellor). We should also make sure that adopting HF kernels doesn't hurt the kernel developer experience, i.e. we can still have fast incremental builds using ccache (investigating this is on my TODO list unless someone else is willing to pick it up). In any case, whether or not we adopt HF kernels, breaking things into separate folders is good first step in my opinion. I am very for just getting started with that while we resolve some of these open questions.

** tl dr;

Overall, I think separating the kernels into the structure you mentioned is a great first step (regardless of where we eventually end up):

vllm/
├─ vllm-engine/            ← Current package
│  ├─ vllm/
│  ├─ pyproject.toml
├─ vllm-kernels/           ← New package
│  ├─ vllm_kernels/        ← Python frontend
│  ├─ csrc/                ← Kernel code
│  ├─ pyproject.toml

—or, alternatively (personally id prefer this), similar to what @hmellor suggested:

vllm/                     ← Current package
│─ vllm/
│─ pyproject.toml
vllm-kernels/             ← New package
│─ vllm_kernels/          ← Python frontend
│─ csrc/                  ← Kernel code
│─ pyproject.toml

[LucasWilkinson]

Also, vLLM already makes many calls to huggingface.co to fetch configs and model weights, so we're not introducing an unknown.

This does not happen if you pass a model as a local path though right? Or am I mistaken and this results in network calls too.

[LucasWilkinson]

For Docker images - we can download and bundle the binaries at build time (just as we do for other dependencies).

How would this work from an implementation perspective? would we have to maintain a list of possible get_kernel kernel calls? may _custom_ops.py could contain all the get_kernel calls? structured such that simply importing that file would be sufficient to trigger a download of all the kernels

[tlrmchlsmth]

@hmellor I see the value of having something like huggingface kernels, but a few questions/concerns:

• What would the kernel developer experience be like? Right now it's convenient to build kernels in-tree. One of the goals with this project is to keep development easy. So what does the kernel developer workflow look like? How does it work with ccache or sccache?
• Target management. How do we manage packaging kernels with different CUDA versions? And what about RoCM?

[hmellor]

This does not happen if you pass a model as a local path though right? Or am I mistaken and this results in network calls too.

@LucasWilkinson

I think unless you but the HF Hub library in offline mode (HF_HUB_OFFLINE=1) it will check the Hub to see what it needs and then just not download it again if it's in the cache.

[hmellor]

we should still ship standalone wheels that don’t require network calls

How would this work from an implementation perspective? would we have to maintain a list of possible get_kernel kernel calls? may _custom_ops.py could contain all the get_kernel calls? structured such that simply importing that file would be sufficient to trigger a download of all the kernels

@LucasWilkinson

get_kernel is a combination of install_kernel and import_from_path. So I could see a vLLM (not kernels) build process where (for Python or Docker) we download all the kernels we need (method to be decided) and bundle them in the build artefact. Then at runtime, we can use import_from_path (or whichever method is most appropriate) so that air-gapped machines are unaffected.

def get_kernel(repo_id: str, revision: str = "main") -> ModuleType:
    package_name, package_path = install_kernel(repo_id, revision=revision)
    return import_from_path(package_name, package_path / package_name / "__init__.py")

You can take a look in https://github.com/huggingface/kernels/blob/main/src/kernels/utils.py for all the available functions for managing and using kernels.

[hmellor]

In general, it's worth having a browse of the docs for kernel-builder as the team has written down a lot of useful information. For your specific questions:

We should also make sure that adopting HF kernels doesn't hurt the kernel developer experience, i.e. we can still have fast incremental builds using ccache

@LucasWilkinson

What would the kernel developer experience be like?

@tlrmchlsmth

We have this guide for local development of kernels.

On the topic of incremental build caching, the kernel-builder uses Nix which does perform efficient caching of build artefacts.

Target management. How do we manage packaging kernels with different CUDA versions? And what about RoCM?

@tlrmchlsmth

The kernel-builder builds binaries that are compatible with everything PyTorch is compatible with. They're listed in this doc. As for which versions of PyTorch does the kernel-builder support, it will support the current and previous version of PyTorch. So the kernels-builder would be updated in lock-step with PyTorch in vLLM.

[tlrmchlsmth]

@hmellor is there something that tracks the PyTorch nightly? One downside of this would be that it rules out upgrading to CUDA 12.9 in order to pick up performance improvements

[hmellor]

@hmellor is there something that tracks the PyTorch nightly? One downside of this would be that it rules out upgrading to CUDA 12.9 in order to pick up performance improvements

These are both good points, I'll speak to the team and get back to you.

[LucasWilkinson]

Looks like: https://github.com/huggingface/kernels/blob/main/docs/locking.md can help with downloading assuming we list all of the dependencies

I guess im mostly just interested in what the end2end flow of HF kernels would look like; for example if we had a PR that updated:

vllm/attention/backends/flash_attn.py

// with other associated backend changes

- ops.reshape_and_cache_flash(...)
+ ops.reshape_and_cache_flash(..., stride_order=layer.stride_order)

vllm-kernels/cache_ops/build.toml

// specifies version v0.2.1, not sure where a version is specified

vllm-kernels/cache_ops/cache_kernel/reshape_and_cache.cu

- void reshape_and_cache_flash(...)
+ void reshape_and_cache_flash(..., std::vector<int> stride_order)

// With other kernel changes

pyproject.toml

[tool.kernels.dependencies]
- "vllm-kernels/cache_ops" == v0.2.0
+ "vllm-kernels/cache_ops" == v0.2.1

Then once this lands on main could we have a CI pass that recognizes that vllm-kernels/cache_ops's version changed and pushes a new one to the HF hub? I think we'd ideally want to do this once it lands on main that way we don't run into issues where two PRs could have the same version number for a kernel library.