Registry / ai-ml / vllm-flash-attn

vllm-flash-attn

library2.6.2pypypi✓ verified 85d ago

Forward-only flash-attention kernel optimized for vLLM inference. Version 2.6.2 is the latest, released as a lightweight wrapper around the Flash Attention CUDA kernel with a simplified forward-only API. Development is active alongside vLLM releases.

ai-ml llm-agents

pip install vllm-flash-attn

github ↗package ↗

INSTALL

IMPORT

SIG · VLLM-FLASH-ATTN

vllm-flash-attn

ai-mlpythonv2.6.2

Install

90.8s avg

Import

3800ms

Disk

5274MB

Pass rate

4/ 10

Env Coverage4 / 10

glibc

3.9–3.13

musl

3.9–3.13

Install & Compatibility

Where this runs

tested against v2.6.2 · pip install

no network on importno background threads

Install × environment matrix

Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux

musl

glibc

py 3.10

✕ build_error

✓ 93s

py 3.11

✕ build_error

✓ 88.5s

py 3.12

✕ build_error

✓ 79.8s

py 3.13

✕ build_error

py 3.9

✕ build_error

✓ 101.9s

5274MB installed

● package 5274MB

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

flash_attn_func

✓ from vllm_flash_attn import flash_attn_func

✗ from flash_attn import flash_attn_func

flash_attn provides a full trian/test API; vllm_flash_attn exports only forward-only functions.

flash_attn_with_kvcache

✓ from vllm_flash_attn import flash_attn_with_kvcache

✗ from flash_attn import flash_attn_with_kvcache

This function is specifically for inference with precomputed KV cache.

Basic usage of forward-only flash attention. Requires CUDA GPU.

import torch from vllm_flash_attn import flash_attn_func q = torch.randn(1, 1, 8, 64, device='cuda', dtype=torch.float16) k = torch.randn(1, 1, 8, 64, device='cuda', dtype=torch.float16) v = torch.randn(1, 1, 8, 64, device='cuda', dtype=torch.float16) out = flash_attn_func(q, k, v, softmax_scale=1.0, causal=False) print(out.shape)

Debug

Known issues

gotchaThis package is forward-only. It does not support backward pass gradients. Using it in training will silently produce wrong gradients or crash.

fix

Use the full flash-attn package (flash_attn) for training.

affects: all

deprecatedSupport for compute capability < 8.0 (e.g., V100) was dropped in v2.6.0. Older versions may still work but are unmaintained.

fix

Upgrade GPU to Volta or newer, or pin to vllm-flash-attn<2.6.0 if on older hardware.

affects: >=2.6.0

gotchaThe function signature for flash_attn_func changed in v2.6.0: the `softmax_scale` parameter is no longer optional and must be passed explicitly.

fix

Always pass softmax_scale=1.0 or your desired scaling factor.

affects: >=2.6.0

Errors

Common errors & fixes

ImportError: No module named 'vllm_flash_attn'

Package not installed or installed with an older name (flash_attn instead of vllm_flash_attn).

fix

pip install vllm-flash-attn

RuntimeError: FlashAttention only supports CUDA with compute capability >= 8.0

GPU older than Ampere (e.g., V100, GTX 1080).

fix

Update to a GPU with compute capability >= 8.0, or use a CPU implementation.

AssertionError: Input tensor must be contiguous in memory

Input tensors are not contiguous; flash attention requires contiguous memory layout.

fix

Call .contiguous() on tensors before passing: q.contiguous(), k.contiguous(), v.contiguous()

Upgrade

Version history

2.6.2latest on PyPI · released Sep 5, 2024

Audit

Dependencies

torchrequiredFlash attention kernels depend on PyTorch tensors and CUDA.

flash-attnrequiredThe underlying flash-attention library; vllm-flash-attn wraps it for forward-only use.

Agent activity

58 hits · last 30 days

node

Amazon

Perplexity

OpenAI (training)

Resources

githubhttps://github.com/vllm-project/flash-attention.git ↗

packagehttps://pypi.org/project/vllm-flash-attn/ ↗