@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / dccf535f8e

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 18:30 UTC

[V1] Enable V1 Fp8 cache for FA3 in the oracle (#15191)

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Lucas Wilkinson committed Mar 23, 2025 at 18:07 UTC dccf535f8edb861e95cf2f0b3512e1fd737265c2

9 files changed +45 -23

.gitignore

+2 -1

@@ -2,7 +2,8 @@

2 /vllm/_version.py

4 # vllm-flash-attn built from source

5 -vllm/vllm_flash_attn/

5 +vllm/vllm_flash_attn/*

6 +!vllm/vllm_flash_attn/fa_utils.py

8 # Byte-compiled / optimized / DLL files

9 __pycache__/

vllm/attention/backends/flash_attn.py

+12 -4

@@ -22,12 +22,13 @@ from vllm.attention.backends.utils import (

22 compute_slot_mapping_start_idx, get_num_prefill_decode_query_kv_tokens,

23 get_seq_len_block_table_args, is_all_cross_attn_metadata_set,

24 is_all_encoder_attn_metadata_set, is_block_tables_empty)

25 -from vllm.fa_utils import get_flash_attn_version

25 from vllm.logger import init_logger

26 from vllm.multimodal import MultiModalPlaceholderMap

27 from vllm.utils import async_tensor_h2d, make_tensor_with_pad

28 from vllm.vllm_flash_attn import (flash_attn_varlen_func,

29 flash_attn_with_kvcache)

30 +from vllm.vllm_flash_attn.fa_utils import (flash_attn_supports_fp8,

31 + get_flash_attn_version)

33 if TYPE_CHECKING:

34 from vllm.worker.model_runner import (ModelInputForGPUBuilder,

@@ -632,10 +633,13 @@ class FlashAttentionImpl(AttentionImpl):

633 self.kv_cache_dtype = kv_cache_dtype

634 self.vllm_flash_attn_version = get_flash_attn_version(

635 requires_alibi=self.alibi_slopes is not None)

635 - if (is_quantized_kv_cache(self.kv_cache_dtype)

636 - and self.vllm_flash_attn_version != 3):

636 + if is_quantized_kv_cache(self.kv_cache_dtype) and (

637 + not self.kv_cache_dtype.startswith("fp8")

638 + or not flash_attn_supports_fp8()):

639 raise NotImplementedError(

638 - "Only FlashAttention3 supports FP8 KV cache")

640 + f"FlashAttention does not support {self.kv_cache_dtype} "

641 + "kv-cache on this device "

642 + f"(FA supports fp8 = {flash_attn_supports_fp8()}).")

643 if logits_soft_cap is None:

644 # In flash-attn, setting logits_soft_cap as 0 means no soft cap.

645 logits_soft_cap = 0

@@ -704,6 +708,10 @@ class FlashAttentionImpl(AttentionImpl):

708 logits_soft_cap: Optional[float] = self.logits_soft_cap

709 fp8_attention = kv_cache_dtype.startswith("fp8")

710

711 + if fp8_attention and not flash_attn_supports_fp8():

712 + raise NotImplementedError(

713 + "FlashAttention does not support FP8 kv-cache on this device.")

714 +

715 if kv_cache.numel() > 0:

716 key_cache = kv_cache[0]

717 value_cache = kv_cache[1]

vllm/attention/backends/mla/common.py

+1 -1

@@ -205,7 +205,6 @@ from vllm.attention.backends.utils import (PAD_SLOT_ID, compute_slot_mapping,

205 compute_slot_mapping_start_idx,

206 is_block_tables_empty)

207 from vllm.attention.ops.triton_merge_attn_states import merge_attn_states

208 -from vllm.fa_utils import get_flash_attn_version

208 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

209 LinearBase, RowParallelLinear,

210 UnquantizedLinearMethod)

@@ -214,6 +213,7 @@ from vllm.model_executor.layers.rotary_embedding import (

213 from vllm.multimodal import MultiModalPlaceholderMap

214 from vllm.platforms import current_platform

215 from vllm.utils import async_tensor_h2d, cdiv, make_tensor_with_pad, round_down

216 +from vllm.vllm_flash_attn.fa_utils import get_flash_attn_version

217

218 try:

219 from vllm.vllm_flash_attn import flash_attn_varlen_func

vllm/config.py

-4

@@ -1157,10 +1157,6 @@ class CacheConfig:

1157 if self.cache_dtype == "auto":

1158 pass

1159 elif self.cache_dtype in ("fp8", "fp8_e4m3", "fp8_e5m2"):

1160 - if envs.VLLM_USE_V1:

1161 - raise NotImplementedError(

1162 - "V1 does not yet support fp8 KV cache. "

1163 - "Set VLLM_USE_V1=0 to enable fp8 kv cache.")

1160 logger.info(

1161 "Using fp8 data type to store kv cache. It reduces the GPU "

1162 "memory footprint and boosts the performance. "

vllm/engine/arg_utils.py

+14 -3

@@ -1562,9 +1562,20 @@ class EngineArgs:

1562

1563 # No Fp8 KV cache so far.

1564 if self.kv_cache_dtype != "auto":

1565 - _raise_or_fallback(feature_name="--kv-cache-dtype",

1566 - recommend_to_remove=False)

1567 - return False

1565 + fp8_attention = self.kv_cache_dtype.startswith("fp8")

1566 + will_use_fa = (

1567 + current_platform.is_cuda()

1568 + and not envs.is_set("VLLM_ATTENTION_BACKEND")

1569 + ) or envs.VLLM_ATTENTION_BACKEND == "FLASH_ATTN_VLLM_V1"

1570 + supported = False

1571 + if fp8_attention and will_use_fa:

1572 + from vllm.vllm_flash_attn.fa_utils import (

1573 + flash_attn_supports_fp8)

1574 + supported = flash_attn_supports_fp8()

1575 + if not supported:

1576 + _raise_or_fallback(feature_name="--kv-cache-dtype",

1577 + recommend_to_remove=False)

1578 + return False

1579

1580 # No Prompt Adapter so far.

1581 if self.enable_prompt_adapter:

vllm/platforms/cuda.py

+3 -5

@@ -14,7 +14,6 @@ from typing_extensions import ParamSpec

14 # import custom ops, trigger op registration

15 import vllm._C # noqa

16 import vllm.envs as envs

17 -from vllm.fa_utils import get_flash_attn_version

17 from vllm.logger import init_logger

18 from vllm.utils import import_pynvml

@@ -258,7 +257,7 @@ class CudaPlatformBase(Platform):

257 try:

258 import vllm.vllm_flash_attn # noqa: F401

259 from vllm.attention.backends.flash_attn import ( # noqa: F401

261 - FlashAttentionBackend)

260 + FlashAttentionBackend, flash_attn_supports_fp8)

261

262 supported_sizes = \

263 FlashAttentionBackend.get_supported_head_sizes()

@@ -269,10 +268,9 @@ class CudaPlatformBase(Platform):

268 target_backend = _Backend.XFORMERS

269 fp8_kv_cache = (kv_cache_dtype is not None

270 and kv_cache_dtype.startswith("fp8"))

272 - if (fp8_kv_cache and get_flash_attn_version() != 3):

271 + if (fp8_kv_cache and not flash_attn_supports_fp8()):

272 logger.info(

274 - "Cannot use FlashAttention-2 backend for FP8 KV cache."

275 - )

273 + "Cannot use FlashAttention backend for FP8 KV cache.")

274 logger.warning(

275 "Please use FlashInfer backend with FP8 KV Cache for "

276 "better performance by setting environment variable "

vllm/v1/attention/backends/flash_attn.py

+6 -4

@@ -11,10 +11,11 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,

11 AttentionMetadata, AttentionType,

12 is_quantized_kv_cache)

13 from vllm.attention.ops.triton_merge_attn_states import merge_attn_states

14 -from vllm.fa_utils import get_flash_attn_version

14 from vllm.logger import init_logger

15 from vllm.platforms import current_platform

16 from vllm.utils import cdiv

17 +from vllm.vllm_flash_attn.fa_utils import (flash_attn_supports_fp8,

18 + get_flash_attn_version)

20 if TYPE_CHECKING:

21 from vllm.v1.core.sched.output import SchedulerOutput

@@ -182,9 +183,6 @@ class FlashAttentionImpl(AttentionImpl):

183 else:

184 self.sliding_window = (sliding_window - 1, 0)

185 self.kv_cache_dtype = kv_cache_dtype

185 - if is_quantized_kv_cache(self.kv_cache_dtype):

186 - raise NotImplementedError(

187 - "FlashAttention V1 with FP8 KV cache not yet supported")

186 if logits_soft_cap is None:

187 # In flash-attn, setting logits_soft_cap as 0 means no soft cap.

188 logits_soft_cap = 0

@@ -206,6 +204,10 @@ class FlashAttentionImpl(AttentionImpl):

204 "are not implemented for "

205 "FlashAttentionImpl")

206 self.vllm_flash_attn_version = get_flash_attn_version()

207 + if is_quantized_kv_cache(self.kv_cache_dtype) \

208 + and not flash_attn_supports_fp8():

209 + raise NotImplementedError(

210 + "FlashAttention does not support fp8 kv-cache on this device.")

211

212 def forward(

213 self,

vllm/v1/attention/backends/mla/common.py

+1 -1

@@ -196,7 +196,6 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionLayer,

196 AttentionMetadata,

197 MLAAttentionImpl)

198 from vllm.attention.ops.triton_merge_attn_states import merge_attn_states

199 -from vllm.fa_utils import get_flash_attn_version

199 from vllm.logger import init_logger

200 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

201 LinearBase, RowParallelLinear,

@@ -204,6 +203,7 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,

203 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding

204 from vllm.platforms import current_platform

205 from vllm.utils import cdiv, round_down

206 +from vllm.vllm_flash_attn.fa_utils import get_flash_attn_version

207

208 try:

209 from vllm.vllm_flash_attn import flash_attn_varlen_func

vllm/vllm_flash_attn/fa_utils.py renamed

@@ -46,3 +46,9 @@ def get_flash_attn_version(requires_alibi: bool = False) -> Optional[int]:

46 return fa_version

47 except (ImportError, AssertionError):

48 return None

49 +

50 +

51 +def flash_attn_supports_fp8() -> bool:

52 + from vllm.platforms import current_platform

53 + return get_flash_attn_version() == 3 and \

54 + current_platform.get_device_capability().major == 9