Instructions to use replicate/flashinfer-draft with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Kernels
How to use replicate/flashinfer-draft with Kernels:
# !pip install kernels from kernels import get_kernel kernel = get_kernel("replicate/flashinfer-draft") - Notebooks
- Google Colab
- Kaggle
| # Generate FlashInfer Source Files | |
| ```bash | |
| # 1. Clone FlashInfer (pin version with --branch) | |
| git clone --depth 1 --branch v0.2.0 https://github.com/flashinfer-ai/flashinfer flashinfer | |
| # 2. Apply patch (if needed) | |
| cd flashinfer && git apply ../generated.patch | |
| # 3. Generate AOT files | |
| export CUDA_HOME=/usr/local/cuda-12.6 && export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.7 8.9 9.0" && uv run --with torch python -m flashinfer.aot | |
| # 4. Copy generated files to csrc directory | |
| cp -r build/aot/generated ../csrc/generated | |
| ``` | |
| now manually comment out the `TORCH_LIBRARY_FRAGMENT` in the generated files | |
| specifically in the activation `csrc/generated/gelu_and_mul.cu` and etc since those are included in the current build. | |
| ### build and test | |
| ```bash | |
| nix develop -L .#test --command python tests/simple_test.py | |
| ``` | |
| ```txt | |
| GELU and multiply operation completed. Output shape: torch.Size([128, 2048]) | |
| Output tensor sample: tensor([[ 7.8613e-02, -4.7656e-01, -4.8637e-03, 6.2073e-02, -3.8745e-01], | |
| [ 1.5686e-01, -5.0964e-03, -3.8981e-04, 1.8945e+00, -7.3792e-02], | |
| ... | |
| ``` |