Top AI Repos — open-source AI, indexed and scored
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
FlexAttention based, minimal vllm-style inference engine for fast Gemma 2 inference.
| Date | Stars |
|---|---|
| 2026-07-31 | 356 |
| 2026-08-06 | 357 |
Today
+1 stars today
This week
— stars this week
This month
— stars this month
Momentum
0.0
growth rate 0.00%/day
# flex-nano-vllm
FlexAttention based, minimal vllm-style inference engine for fast Gemma 2 inference.
## Introduction
This project has no flash-attn dependency, no custom triton kernel. Everything is implemented with FlexAttention. The code is commented, the structure is flat. Read the accompanying write-up: [vLLM flex attention from scratch](https://jonathanc.net/blog/vllm-flex-attention-from-scratch).
## Code Structure
```
flex-nano-vllm/
├── benchmark.py # Testing and benchmarking script.
├── benchmark_vllm.py # vLLM comparison benchmark (uses uv inline dependency to run vLLM).
├── visualize.py # Performance visualization script.
└── flex_nano_vllm/
├── inference.py # Main inference engine, uses paged attention.
├── modeling_gemma2.py # Gemma2 model implementation, copied from transformers.
└── paged_attention.py # Paged attention implementation, including page table and paged kv cache. Based on attention-gym.
```
## Quick Start
```
uv sync
# run test and benchmark
uv run benchmark.py
# compare with vllm
uv run benchmark_vllm.py
# enable profiling to save more metrics to a csv file
# ENABLE_PROFILING=1 uv run benchmark_vllm.py
```
## Results
Test configuration:
- PyTorch version: 2.7.1+cu128
- GPU: RTX 3090 x 1 (24GB)
- Model: google/gemma-2-2b
- Workload: 512 requests, max 512 input tokens, variable output tokens (128-512)
- Configs tested: vLLM at 50% & 90% GPU memory, flex-nano-vllm with same page allocation as vLLM
| Implementation | Output Tokens/s | Request/s | Total Throughput* |
|---------------|----------------|-----------|------------------|
| vLLM v1, 90% GPU memory, high batch size† | 3,840 | 17.67 | 7,234 |
| vLLM v1, 90% GPU memory | 3,772 | 15.26 | 6,401 |
| flex-nano-vllm, 90% GPU memory, high batch size† | 3,440 | 14.30 | 5,817 |
| flex-nano-vllm, 90% GPU memory | 3,076 | 13.06 | 5,382 |
| vLLM v1, 50% GPU memory | 3,020 | 13.74 | 5,448 |
| flex-nano-vllm, 50% GPU memory | 2,313 | 9.96 | 4,068 |
*Total throughput includes both input and output tokens
† High batch size means max_num_seqs=512 in vllm (maximum allowed concurrency)
## License
This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
Third-party code incorporated in this project retains its original licenses. See [THIRD_PARTY_LICENSES.md](THIRD_PARTY_LICENSES.md) for details.
## Acknowledgments
- [GeeeekExplorer/nano-vllm](https://github.com/GeeeekExplorer/nano-vllm): this project is inspired by nano-vllm.
- [pytorch-labs/attention-gym](https://github.com/pytorch-labs/attention-gym): The paged attention implementation is based on attention-gym.
- [huggingface/transformers](https://github.com/huggingface/transformers): I copied the gemma2 model from transformers and modified it to use flex attention / paged attention.
- [vllm-project/vllm](https://github.com/vllm-project/vllm): vLLM has support for flex attention backend, which helped me find a useful flag in flex_attention.
Excerpt of 3,121 characters
Read on GitHubWould you bet a product on this? Bounded 0–100 and slow moving.
matched fp:56121ace12b4f692, desc:inference engine