flashinfer
v0.2.0.post1
| بلوق | الوثائق | الركود | منتدى المناقشة |
FlashInfer هو مولد مكتبة و kernel لنماذج اللغة الكبيرة التي توفر تنفيذًا عالي الأداء لنواة GPU LLM مثل Flashattention و SPARSETTENT و PAGEATTENTINE وأخذ العينات والمزيد. يركز FlashInfer على خدمة LLM واستدلالها ، ويقدم أداءً أحدث عبر سيناريوهات متنوعة.
تحقق من مدونة الإصدار V0.2 للحصول على ميزات جديدة!
تشمل الميزات الأساسية لـ FlashInfer:
plan FlashInfer / run مرحلة حساب الانتباه حيث نقوم بجدولة حساب المدخلات المتغيرة الطول في مرحلة plan لتخفيف مشكلة توازن الحمل.تدعم FlashInfer Pytorch و TVM و C ++ (رأس فقط) ، ويمكن دمجها بسهولة في المشاريع الحالية.
إن استخدام API Pytorch الخاص بنا هو أسهل طريقة للبدء:
نحن نقدم عجلات مسبقة الصنع للينكس. يمكنك تثبيت FlashInfer مع الأمر التالي:
# For CUDA 12.4 & torch 2.4
pip install flashinfer -i https://flashinfer.ai/whl/cu124/torch2.4
# For other CUDA & torch versions, please check https://docs.flashinfer.ai/installation.htmlنقدم أيضًا عجلات مبنية ليلا لتجربة أحدث الميزات من الفرع الرئيسي:
pip install flashinfer -i https://flashinfer.ai/whl/nightly/cu124/torch2.4بدلاً من ذلك ، يمكنك بناء FlashInfer من المصدر:
git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
cd flashinfer
pip install -e . -v بشكل افتراضي ، يستخدم FlashInfer مجموعة للتو في الوقت المناسب (JIT) لنواةه. لتوصيل النواة الأساسية مسبقًا ، قم بتعيين متغير البيئة FLASHINFER_ENABLE_AOT=1 قبل تشغيل أمر التثبيت:
FLASHINFER_ENABLE_AOT=1 pip install -e . -vلمزيد من التفاصيل ، راجع التثبيت من وثائق المصدر.
فيما يلي مثال أدنى على استخدام حبات الانتباه/الاهتمام المسبق لـ FlashInfer:
import torch
import flashinfer
kv_len = 2048
num_kv_heads = 32
head_dim = 128
k = torch . randn ( kv_len , num_kv_heads , head_dim ). half (). to ( 0 )
v = torch . randn ( kv_len , num_kv_heads , head_dim ). half (). to ( 0 )
# decode attention
num_qo_heads = 32
q = torch . randn ( num_qo_heads , head_dim ). half (). to ( 0 )
o = flashinfer . single_decode_with_kv_cache ( q , k , v ) # decode attention without RoPE on-the-fly
o_rope_on_the_fly = flashinfer . single_decode_with_kv_cache ( q , k , v , pos_encoding_mode = "ROPE_LLAMA" ) # decode with LLaMA style RoPE on-the-fly
# append attention
append_qo_len = 128
q = torch . randn ( append_qo_len , num_qo_heads , head_dim ). half (). to ( 0 ) # append attention, the last 128 tokens in the KV-Cache are the new tokens
o = flashinfer . single_prefill_with_kv_cache ( q , k , v , causal = True ) # append attention without RoPE on-the-fly, apply causal mask
o_rope_on_the_fly = flashinfer . single_prefill_with_kv_cache ( q , k , v , causal = True , pos_encoding_mode = "ROPE_LLAMA" ) # append attention with LLaMA style RoPE on-the-fly, apply causal mask
# prefill attention
qo_len = 2048
q = torch . randn ( qo_len , num_qo_heads , head_dim ). half (). to ( 0 ) # prefill attention
o = flashinfer . single_prefill_with_kv_cache ( q , k , v , causal = False ) # prefill attention without RoPE on-the-fly, do not apply causal maskتحقق من وثائق لاستخدام حلقات Decode/Aspend/Prefill و kernels المتتالية المشتركة.
نحن نتعرف على أداء flashinfer kernel مع NVBench ويمكنك تجميع وتشغيل المعايير مع الأوامر التالية:
mkdir build
cp cmake/config.cmake build # you can modify the config.cmake to enable/disable benchmarks and change CUDA architectures
cd build
cmake ..
make -j12 يمكنك تشغيل ./bench_{single/batch}_{prefill/decode} single/batch )_ ./bench_single_prefill prefill/decode} لقياس الأداء (على سبيل المثال ./bench_{single/batch}_{prefill/decode} --help الخيارات المتاحة.
يوفر FlashInfer أيضًا ربطات C ++ API و TVM ، يرجى الرجوع إلى الوثائق لمزيد من التفاصيل.
يسرنا أن نشارك أن FlashInfer يتم تبنيه من قبل العديد من المشاريع المتطورة ، بما في ذلك على سبيل المثال لا الحصر:
يستوحى FlashInfer من مشاريع Flashattention 1 & 2 و VLLM و Stream-K و Cutlass و Aitemplate.
