MinivLLM:軽量vLLMクローン(paged/flash attention 実装)

AI/ML

概要

MinivLLM は Nano-vLLM を土台にした、vLLM 型推論エンジンの小型再現実装です。特徴は外部依存を最小化した自己完結型の paged attention(デコード向けのKVページング)と flash attention(高速な行列積近似/最適化)を Python 側で実装している点にあります。プロジェクトは主に二つの観点でのベンチマークを提供します:prefilling(入力の前詰め)における flash attention の性能と、デコード時における paged attention の挙動・スループット。サンプルの起動手順や依存同期に uv を利用する点も特徴で、研究や実験を手早く始めたい開発者向けのミニマル実装です。

GitHub

リポジトリの統計情報

  • スター数: 14
  • フォーク数: 0
  • ウォッチャー数: 14
  • コミット数: 12
  • ファイル数: 10
  • メインの言語: Python

主な特徴

  • Nano-vLLM を基にしたシンプルで理解しやすい vLLM 風推論エンジンの再現。
  • paged attention(KV をページ単位で扱う)と flash attention を自己完結的に実装。
  • prefilling と decoding のそれぞれでのベンチマークスクリプトを同梱。
  • uv パッケージマネージャを用いた依存同期と手軽な起動フロー。

技術的なポイント

MinivLLM の主要な技術的関心は「メモリ効率」と「速度最適化」のトレードオフを、簡潔な実装で示す点にあります。まず paged attention はデコード時に発生する KV キャッシュの肥大化問題に対処する手法で、全 KV を一度にメモリ上に展開する代わりに、必要な領域(ページ)だけを読み書きすることでメモリ使用量を抑え、長文生成時のスケーラビリティを改善します。MinivLLM ではこのページングの概念を模した実装を行い、デコードでのレイテンシやスループットに与える影響を測定できるようにしています。

一方 flash attention は Attention 計算の高速化技術で、従来のトークン毎の QK^T 計算を最適化し、メモリアクセスや数値安定性の観点で有利になるアルゴリズムや低レベル最適化(積和演算のブロッキング、精度管理など)を取り入れます。本リポジトリは C++/CUDA の大規模実装に頼らず、Python レイヤーで可能な限りの最適化を再現し、prefilling(長いコンテキストを最初にバッファする段階)における効果を評価できるようにした点が特徴です。

さらに、ベンチマーク構成は prefilling と decoding を分けて評価する点が実用的です。モデルの初期コンテキスト読み込みでの高速化(flash attention)と、生成ループでのメモリ管理(paged attention)は要求される最適化が異なるため、個別に測定することでどのアプローチがどの場面で有効かを明確化します。依存管理に uv を使うことで環境再現性を高め、main.py やベンチマークスクリプトをそのまま動かして比較実験が行えるようになっています。

実験的・教育的な趣旨のため、コードは複雑なフレームワーク依存を避け、アルゴリズムの本質が理解できる構造になっています。これにより、研究者やエンジニアは実装を読み解きながら、独自の最適化(例えば異なるページサイズやブロックサイズ、数値フォーマット)を試すことが容易です。将来的には GPU ネイティブの高速化(CUDA 実装や Triton など)と組み合わせることで、より実運用に近い速度比較も可能になります。

プロジェクトの構成

主要なファイルとディレクトリ:

  • .gitignore: file
  • LICENSE: file
  • README.md: file
  • benchmark_decoding.py: file
  • benchmark_prefilling.py: file

…他 5 ファイル

まとめ

ミニマルかつ教育的な vLLM の再現実装で、paged/flash attention の比較実験に最適。

リポジトリ情報:

READMEの抜粋:

MyVLLM

A custom implementation of vLLM inference engine with attention mechanism benchmarks, based on Nano-vLLM but with self-contained paged attention and flash attention implementation.

Benchmarking on flash attention in prefilling time and paged attention in decoding time are provided.

Quickstart

# Install uv package manager
curl -LsSf https://astral.sh/uv/install.sh | sh

# Sync dependencies
uv sync

# Run the main inference engine
uv run python main.py

# Run prefilling be...