116f4be405 · tduncan/Nvidia-Optimization-Stack-VLLM

vllm/config/cache.py

+2 -1

@@ -8,6 +8,7 @@ from pydantic import Field, SkipValidation, field_validator, model_validator

8

9 from vllm.config.utils import config

10 from vllm.logger import init_logger

11 +from vllm.utils.torch_utils import is_quantized_kv_cache

12

13 logger = init_logger(__name__)

14

@@ -236,7 +237,7 @@ class CacheConfig:

237 @field_validator("cache_dtype", mode="after")

238 @classmethod

239 def _validate_cache_dtype(cls, cache_dtype: CacheDType) -> CacheDType:

239 - if cache_dtype.startswith("fp8"):

240 + if is_quantized_kv_cache(cache_dtype):

241 logger.info(

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

243 "memory footprint and boosts the performance. "

vllm/model_executor/layers/attention/mla_attention.py

+5 -4

@@ -241,6 +241,7 @@ from vllm.utils.flashinfer import has_flashinfer, has_nvidia_artifactory

241 from vllm.utils.math_utils import cdiv, round_down

242 from vllm.utils.torch_utils import (

243 direct_register_custom_op,

244 + is_quantized_kv_cache,

245 kv_cache_dtype_str_to_dtype,

246 )

247 from vllm.v1.attention.backend import (

@@ -342,7 +343,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

343 # Automatically convert fp8 kv-cache format to "fp8_ds_mla"

344 if (

345 self.attn_backend.get_name() == "FLASHMLA_SPARSE"

345 - and kv_cache_dtype.startswith("fp8")

346 + and is_quantized_kv_cache(kv_cache_dtype)

347 and kv_cache_dtype != "fp8_ds_mla"

348 ):

349 assert cache_config is not None

@@ -356,7 +357,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

357

358 if (

359 self.attn_backend.get_name() == "FLASHINFER_MLA_SPARSE"

359 - and kv_cache_dtype.startswith("fp8")

360 + and is_quantized_kv_cache(kv_cache_dtype)

361 ):

362 logger.info_once(

363 "Using standard fp8 KV cache format. To use DeepSeek's fp8_ds_mla "

@@ -571,7 +572,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

572 if self.impl.dcp_world_size == -1:

573 self.impl.dcp_world_size = get_dcp_group().world_size

574

574 - fp8_attention = self.kv_cache_dtype.startswith("fp8")

575 + fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype)

576

577 num_actual_toks = attn_metadata.num_actual_tokens

578

@@ -1434,7 +1435,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1435 is enabled, else model dtype.

1436 """

1437 use_fp8 = (

1437 - vllm_config.cache_config.cache_dtype.startswith("fp8")

1438 + is_quantized_kv_cache(vllm_config.cache_config.cache_dtype)

1439 and vllm_config.attention_config.use_prefill_query_quantization

1440 and backend_supports_prefill_query_quantization()

1441 )

vllm/model_executor/layers/quantization/kv_cache.py

+1 -1

@@ -9,7 +9,7 @@ from vllm.model_executor.layers.quantization.base_config import (

9 QuantizeMethodBase,

10 )

11 from vllm.platforms import current_platform

12 -from vllm.v1.attention.backend import is_quantized_kv_cache

12 +from vllm.utils.torch_utils import is_quantized_kv_cache

13

14 logger = init_logger(__name__)

15

vllm/model_executor/models/extract_hidden_states.py

+1 -2

@@ -23,14 +23,13 @@ from vllm.model_executor.layers.attention.kv_transfer_utils import (

23 )

24 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

25 from vllm.model_executor.models.utils import maybe_prefix

26 -from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype

26 +from vllm.utils.torch_utils import is_quantized_kv_cache, kv_cache_dtype_str_to_dtype

27 from vllm.v1.attention.backend import (

28 AttentionBackend,

29 AttentionImpl,

30 AttentionMetadataBuilder,

31 AttentionType,

32 CommonAttentionMetadata,

33 - is_quantized_kv_cache,

33 )

34 from vllm.v1.kv_cache_interface import (

35 AttentionSpec,

vllm/platforms/cpu.py

+2 -2

@@ -16,7 +16,7 @@ import torch

16

17 from vllm import envs

18 from vllm.logger import init_logger

19 -from vllm.v1.attention.backend import is_quantized_kv_cache

19 +from vllm.utils.torch_utils import is_quantized_kv_cache

20 from vllm.v1.attention.backends.registry import AttentionBackendEnum

21

22 from .interface import CpuArchEnum, Platform, PlatformEnum

@@ -183,7 +183,7 @@ class CpuPlatform(Platform):

183 "backend is not compatible with FP8 KV cache."

184 )

185

186 - if cache_config.cache_dtype.startswith("fp8"):

186 + if is_quantized_kv_cache(cache_config.cache_dtype):

187 logger.warning(

188 "CPU backend doesn't support KV cache quantization fallback to auto."

189 )

vllm/platforms/cuda.py

+2 -1

@@ -23,6 +23,7 @@ import vllm._C_stable_libtorch # noqa

23 import vllm.envs as envs

24 from vllm.logger import init_logger

25 from vllm.utils.import_utils import import_pynvml

26 +from vllm.utils.torch_utils import is_quantized_kv_cache

27 from vllm.v1.attention.backends.registry import AttentionBackendEnum

28

29 from .interface import DeviceCapability, Platform, PlatformEnum

@@ -87,7 +88,7 @@ def _get_backend_priorities(

88 # Sparse MLA backend priorities

89 # See https://github.com/vllm-project/vllm/issues/35807 for

90 # benchmark results

90 - if kv_cache_dtype is not None and kv_cache_dtype.startswith("fp8"):

91 + if kv_cache_dtype is not None and is_quantized_kv_cache(kv_cache_dtype):

92 # Prefer FlashInfer for fp8 kv cache

93 sparse_backends = [

94 AttentionBackendEnum.FLASHINFER_MLA_SPARSE,

vllm/utils/torch_utils.py

+4

@@ -61,6 +61,10 @@ MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP = {

61 T = TypeVar("T")

62

63

64 +def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:

65 + return kv_cache_dtype.startswith("fp8")

66 +

67 +

68 def is_strictly_contiguous(t: torch.Tensor) -> bool:

69 """

70 Check if tensor is contiguous AND has no degenerate strides.

vllm/v1/attention/backend.py

-4

@@ -954,10 +954,6 @@ class SparseMLAAttentionImpl(AttentionImplBase[T], Generic[T]):

954 )

955

956

957 -def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:

958 - return kv_cache_dtype.startswith("fp8")

959 -

960 -

957 def subclass_attention_backend(

958 name_prefix: str,

959 attention_backend_cls: type[AttentionBackend],

vllm/v1/attention/backends/cpu_attn.py

+1 -1

@@ -9,6 +9,7 @@ from vllm import _custom_ops as ops

9 from vllm.config import VllmConfig

10 from vllm.logger import init_logger

11 from vllm.platforms import CpuArchEnum, current_platform

12 +from vllm.utils.torch_utils import is_quantized_kv_cache

13 from vllm.v1.attention.backend import (

14 AttentionBackend,

15 AttentionImpl,

@@ -16,7 +17,6 @@ from vllm.v1.attention.backend import (

17 AttentionMetadataBuilder,

18 AttentionType,

19 CommonAttentionMetadata,

19 - is_quantized_kv_cache,

20 )

21 from vllm.v1.attention.backends.utils import (

22 split_decodes_and_prefills,

vllm/v1/attention/backends/flash_attn.py

+5 -5

@@ -10,12 +10,12 @@ import numpy as np

10 import torch

11

12 from vllm.model_executor.layers.attention import Attention

13 +from vllm.utils.torch_utils import is_quantized_kv_cache

14 from vllm.v1.attention.backend import (

15 AttentionBackend,

16 AttentionImpl,

17 AttentionType,

18 MultipleOf,

18 - is_quantized_kv_cache,

19 )

20 from vllm.v1.attention.backends.fa_utils import (

21 flash_attn_supports_fp8,

@@ -177,7 +177,7 @@ class FlashAttentionBackend(AttentionBackend):

177 def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:

178 if kv_cache_dtype is None:

179 return True

180 - if kv_cache_dtype.startswith("fp8"):

180 + if is_quantized_kv_cache(kv_cache_dtype):

181 return flash_attn_supports_fp8()

182 return kv_cache_dtype in ["auto", "float16", "bfloat16"]

183

@@ -430,7 +430,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

430 batch_size, cu_query_lens, max_query_len, seqlens, max_seq_len, causal

431 ):

432 cache_dtype = self.cache_config.cache_dtype

433 - if cache_dtype.startswith("fp8"):

433 + if is_quantized_kv_cache(cache_dtype):

434 qkv_dtype = FlashAttentionBackend.get_fp8_dtype_for_flashattn(

435 cache_dtype

436 )

@@ -726,7 +726,7 @@ class FlashAttentionImpl(AttentionImpl):

726 # For decoder and cross-attention, use KV cache as before

727 key_cache, value_cache = kv_cache.unbind(0)

728

729 - if self.kv_cache_dtype.startswith("fp8"):

729 + if is_quantized_kv_cache(self.kv_cache_dtype):

730 # queries are quantized in the attention layer

731 dtype = FlashAttentionBackend.get_fp8_dtype_for_flashattn(

732 self.kv_cache_dtype

@@ -978,7 +978,7 @@ class FlashAttentionImpl(AttentionImpl):

978 )

979

980 # For encoder attention, process FP8 quantization if needed

981 - if self.kv_cache_dtype.startswith("fp8"):

981 + if is_quantized_kv_cache(self.kv_cache_dtype):

982 raise NotImplementedError(

983 "quantization is not supported for encoder attention"

984 )

vllm/v1/attention/backends/flash_attn_diffkv.py

+2 -1

@@ -4,6 +4,7 @@

4

5 import torch

6

7 +from vllm.utils.torch_utils import is_quantized_kv_cache

8 from vllm.v1.attention.backend import AttentionType

9 from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available

10 from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (

@@ -191,7 +192,7 @@ class FlashAttentionDiffKVImpl(FlashAttentionImpl):

192 key_cache = kv_cache[..., : self.head_size]

193 value_cache = kv_cache[..., self.head_size :]

194

194 - if self.kv_cache_dtype.startswith("fp8"):

195 + if is_quantized_kv_cache(self.kv_cache_dtype):

196 # queries are quantized in the attention layer

197 dtype = FlashAttentionBackend.get_fp8_dtype_for_flashattn(

198 self.kv_cache_dtype

vllm/v1/attention/backends/flashinfer.py

+9 -10

@@ -42,7 +42,7 @@ from vllm.utils.flashinfer import (

42 )

43 from vllm.utils.math_utils import cdiv

44 from vllm.utils.platform_utils import is_pin_memory_available

45 -from vllm.utils.torch_utils import is_strictly_contiguous

45 +from vllm.utils.torch_utils import is_quantized_kv_cache, is_strictly_contiguous

46 from vllm.v1.attention.backend import (

47 AttentionBackend,

48 AttentionCGSupport,

@@ -602,7 +602,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

602 self.page_size = self.kv_cache_spec.block_size

603

604 self.cache_dtype = self.cache_config.cache_dtype

605 - if self.cache_dtype.startswith("fp8"):

605 + if is_quantized_kv_cache(self.cache_dtype):

606 self.kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(

607 self.cache_dtype

608 )

@@ -1269,7 +1269,7 @@ class FlashInferImpl(AttentionImpl):

1269 def fused_output_quant_supported(self, quant_key: QuantKey):

1270 return (

1271 self.support_trtllm_attn

1272 - and self.kv_cache_dtype.startswith("fp8")

1272 + and is_quantized_kv_cache(self.kv_cache_dtype)

1273 and quant_key in (kFp8StaticTensorSym, kNvfp4Dynamic)

1274 )

1275

@@ -1317,12 +1317,12 @@ class FlashInferImpl(AttentionImpl):

1317

1318 if self.bmm1_scale is None:

1319 self.bmm1_scale = self.scale

1320 - if self.kv_cache_dtype.startswith("fp8"):

1320 + if is_quantized_kv_cache(self.kv_cache_dtype):

1321 self.bmm1_scale *= layer._q_scale_float * layer._k_scale_float

1322

1323 if self.bmm2_scale is None:

1324 self.bmm2_scale = 1.0

1325 - if self.kv_cache_dtype.startswith("fp8"):

1325 + if is_quantized_kv_cache(self.kv_cache_dtype):

1326 self.bmm2_scale *= layer._v_scale_float

1327

1328 prefill_use_trtllm = isinstance(attn_metadata.prefill, TRTLLMPrefill)

@@ -1375,8 +1375,8 @@ class FlashInferImpl(AttentionImpl):

1375

1376 # The FlashInfer api requires data to be in fp8_e4m3 or fp8_e5m2

1377 # to process the cache when the kv_cache_dtype is fp8

1378 - if self.kv_sharing_target_layer_name is None and self.kv_cache_dtype.startswith(

1379 - "fp8"

1378 + if self.kv_sharing_target_layer_name is None and is_quantized_kv_cache(

1379 + self.kv_cache_dtype

1380 ):

1381 torch_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(

1382 self.kv_cache_dtype

@@ -1486,9 +1486,8 @@ class FlashInferImpl(AttentionImpl):

1486 assert self.o_sf_scale is None

1487 out = output[num_decode_tokens:]

1488

1489 - if (

1490 - attn_metadata.q_data_type != FP8_DTYPE

1491 - and self.kv_cache_dtype.startswith("fp8")

1489 + if attn_metadata.q_data_type != FP8_DTYPE and is_quantized_kv_cache(

1490 + self.kv_cache_dtype

1491 ):

1492 # TRTLLM prefill attention does not support BF16 Q

1493 # and fp8 kv cache. So to enable prefill attention

vllm/v1/attention/backends/flex_attention.py

+1 -2

@@ -27,14 +27,13 @@ from vllm.config.cache import CacheDType

27 from vllm.logger import init_logger

28 from vllm.platforms import current_platform

29 from vllm.utils.math_utils import cdiv

30 -from vllm.utils.torch_utils import is_torch_equal_or_newer

30 +from vllm.utils.torch_utils import is_quantized_kv_cache, is_torch_equal_or_newer

31 from vllm.v1.attention.backend import (

32 AttentionBackend,

33 AttentionImpl,

34 AttentionMetadataBuilder,

35 AttentionType,

36 CommonAttentionMetadata,

37 - is_quantized_kv_cache,

37 )

38 from vllm.v1.kv_cache_interface import AttentionSpec

39

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

+1 -1

@@ -17,12 +17,12 @@ from vllm.model_executor.layers.attention.mla_attention import (

17 )

18 from vllm.platforms.interface import DeviceCapability

19 from vllm.utils.platform_utils import num_compute_units

20 +from vllm.utils.torch_utils import is_quantized_kv_cache

21 from vllm.v1.attention.backend import (

22 AttentionCGSupport,

23 AttentionLayer,

24 AttentionType,

25 MultipleOf,

25 - is_quantized_kv_cache,

26 )

27

28 logger = init_logger(__name__)

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

+2 -2

@@ -20,12 +20,12 @@ from vllm.model_executor.layers.attention.mla_attention import (

20 )

21 from vllm.platforms.interface import DeviceCapability

22 from vllm.utils.math_utils import round_up

23 +from vllm.utils.torch_utils import is_quantized_kv_cache

24 from vllm.v1.attention.backend import (

25 AttentionCGSupport,

26 AttentionLayer,

27 AttentionType,

28 MultipleOf,

28 - is_quantized_kv_cache,

29 )

30 from vllm.v1.attention.backends.fa_utils import (

31 flash_attn_supports_mla,

@@ -319,7 +319,7 @@ class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):

319 q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1

320 )

321

322 - if self.kv_cache_dtype.startswith("fp8"):

322 + if is_quantized_kv_cache(self.kv_cache_dtype):

323 raise NotImplementedError("FP8 FlashAttention MLA not yet supported")

324

325 kv_c_cache = kv_c_and_k_pe_cache[..., : self.kv_lora_rank]

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

+3 -3

@@ -16,12 +16,12 @@ from vllm.model_executor.layers.attention.mla_attention import (

16 QueryLenSupport,

17 )

18 from vllm.platforms.interface import DeviceCapability

19 +from vllm.utils.torch_utils import is_quantized_kv_cache

20 from vllm.v1.attention.backend import (

21 AttentionCGSupport,

22 AttentionLayer,

23 AttentionType,

24 MultipleOf,

24 - is_quantized_kv_cache,

25 )

26 from vllm.v1.attention.backends.utils import KVCacheLayoutType

27

@@ -184,12 +184,12 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):

184

185 if self.bmm1_scale is None:

186 self.bmm1_scale = self.scale

187 - if self.kv_cache_dtype.startswith("fp8"):

187 + if is_quantized_kv_cache(self.kv_cache_dtype):

188 self.bmm1_scale *= layer._q_scale_float * layer._k_scale_float

189

190 if self.bmm2_scale is None:

191 self.bmm2_scale = 1.0

192 - if self.kv_cache_dtype.startswith("fp8"):

192 + if is_quantized_kv_cache(self.kv_cache_dtype):

193 self.bmm2_scale *= layer._k_scale_float

194

195 # Reuse pre-allocated zero-init output buffer to avoid a memset

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

+3 -2

@@ -26,6 +26,7 @@ from vllm.model_executor.layers.attention.mla_attention import (

26 get_mla_dims,

27 )

28 from vllm.platforms.interface import DeviceCapability

29 +from vllm.utils.torch_utils import is_quantized_kv_cache

30 from vllm.v1.attention.backend import (

31 AttentionBackend,

32 AttentionCGSupport,

@@ -341,11 +342,11 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata

342

343 if self.bmm1_scale is None:

344 self.bmm1_scale = self.scale

344 - if self.kv_cache_dtype.startswith("fp8"):

345 + if is_quantized_kv_cache(self.kv_cache_dtype):

346 self.bmm1_scale *= layer._q_scale_float * layer._k_scale_float

347 if self.bmm2_scale is None:

348 self.bmm2_scale = 1.0

348 - if self.kv_cache_dtype.startswith("fp8"):

349 + if is_quantized_kv_cache(self.kv_cache_dtype):

350 self.bmm2_scale *= layer._k_scale_float

351

352 o = trtllm_batch_decode_with_kv_cache_mla(

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

+6 -3

@@ -20,6 +20,7 @@ from vllm.model_executor.layers.attention.mla_attention import (

20 )

21 from vllm.platforms.interface import DeviceCapability

22 from vllm.utils.platform_utils import num_compute_units

23 +from vllm.utils.torch_utils import is_quantized_kv_cache

24 from vllm.v1.attention.backend import (

25 AttentionCGSupport,

26 AttentionLayer,

@@ -128,7 +129,9 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

129

130 self.cg_buf_tile_scheduler_metadata = None

131 self.cg_buf_num_splits = None

131 - self.is_fp8_kvcache = vllm_config.cache_config.cache_dtype.startswith("fp8")

132 + self.is_fp8_kvcache = is_quantized_kv_cache(

133 + vllm_config.cache_config.cache_dtype

134 + )

135

136 num_sms = num_compute_units(self.device.index)

137

@@ -269,7 +272,7 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

272 q = reshape_query_for_spec_decode(q, num_decodes)

273

274 scheduler_metadata = attn_metadata.decode.scheduler_metadata

272 - if envs.VLLM_BATCH_INVARIANT and not self.kv_cache_dtype.startswith("fp8"):

275 + if envs.VLLM_BATCH_INVARIANT and not is_quantized_kv_cache(self.kv_cache_dtype):

276 device = q.device

277 dtype = torch.int32

278

@@ -299,7 +302,7 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

302 scheduler_metadata.tile_scheduler_metadata = tile_scheduler_metadata

303 scheduler_metadata.num_splits = num_splits

304

302 - if self.kv_cache_dtype.startswith("fp8"):

305 + if is_quantized_kv_cache(self.kv_cache_dtype):

306 o, lse = flash_mla_with_kvcache_fp8(

307 q=q,

308 k_cache=kv_c_and_k_pe_cache.unsqueeze(-2), # Add head dim of 1

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

+2 -1

@@ -16,6 +16,7 @@ from vllm.model_executor.layers.attention.mla_attention import (

16 from vllm.platforms import current_platform

17 from vllm.platforms.interface import DeviceCapability

18 from vllm.utils.platform_utils import num_compute_units

19 +from vllm.utils.torch_utils import is_quantized_kv_cache

20 from vllm.v1.attention.backend import (

21 AttentionBackend,

22 AttentionCGSupport,

@@ -571,7 +572,7 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):

572 vllm_config = get_current_vllm_config()

573 max_tokens = vllm_config.scheduler_config.max_num_batched_tokens

574 q_concat_shape = (max_tokens, num_heads, head_size)

574 - if kv_cache_dtype.startswith("fp8"):

575 + if is_quantized_kv_cache(kv_cache_dtype):

576 assert kv_cache_dtype == "fp8_ds_mla", (

577 "FlashMLA Sparse Attention backend fp8 only supports "

578 "fp8_ds_mla kv-cache dtype"

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

+1 -1

@@ -14,11 +14,11 @@ from vllm.model_executor.layers.attention.mla_attention import (

14 MLACommonMetadata,

15 )

16 from vllm.platforms.interface import DeviceCapability

17 +from vllm.utils.torch_utils import is_quantized_kv_cache

18 from vllm.v1.attention.backend import (

19 AttentionLayer,

20 AttentionType,

21 MultipleOf,

21 - is_quantized_kv_cache,

22 )

23 from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd

24

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

+2 -1

@@ -13,6 +13,7 @@ from vllm.logger import init_logger

13 from vllm.model_executor.layers.attention.mla_attention import (

14 get_mla_dims,

15 )

16 +from vllm.utils.torch_utils import is_quantized_kv_cache

17 from vllm.v1.attention.backend import (

18 AttentionBackend,

19 AttentionCGSupport,

@@ -231,7 +232,7 @@ class XPUMLASparseImpl(SparseMLAAttentionImpl[XPUMLASparseMetadata]):

232 # NOTE(lucas): for the sparse FlashMLA kernels the kernels want to use

233 # MQA 576/512 approach for both prefill and decode

234

234 - if self.kv_cache_dtype.startswith("fp8"):

235 + if is_quantized_kv_cache(self.kv_cache_dtype):

236 raise NotImplementedError("FP8 kv is not supported with XPU MLA Sparse yet")

237

238 # Concatenate q if it's a tuple (ql_nope, q_pe)

vllm/v1/attention/backends/rocm_aiter_fa.py

+8 -7

@@ -16,6 +16,7 @@ from vllm.platforms import current_platform

16 from vllm.platforms.interface import DeviceCapability

17 from vllm.utils.math_utils import cdiv

18 from vllm.utils.platform_utils import num_compute_units

19 +from vllm.utils.torch_utils import is_quantized_kv_cache

20 from vllm.v1.attention.backend import (

21 AttentionBackend,

22 AttentionCGSupport,

@@ -291,7 +292,7 @@ if current_platform.is_rocm():

292 new_key_cache = key_cache.view_as(k_cache_template)

293 new_value_cache = value_cache.view_as(v_cache_template)

294 QUANT = False

294 - if kv_cache_dtype.startswith("fp8"):

295 + if is_quantized_kv_cache(kv_cache_dtype):

296 QUANT = True

297 grid = (

298 num_tokens,

@@ -494,7 +495,7 @@ class AiterFlashAttentionMetadataBuilder(

495 if (

496 rocm_aiter_ops.is_shuffle_kv_cache_enabled()

497 and self.scale.numel() == 1

497 - and self.vllm_config.cache_config.cache_dtype.startswith("fp8")

498 + and is_quantized_kv_cache(self.vllm_config.cache_config.cache_dtype)

499 ):

500 layers = get_layers_from_vllm_config(self.vllm_config, Attention)

501 first_layer_name = [k for k in layers][0]

@@ -887,7 +888,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

888 cu_seqlens_kv=swa_cu_seqlens,

889 token_to_batch=swa_token_to_batch,

890 seq_starts=swa_seq_starts,

890 - dequant=self.kv_cache_dtype.startswith("fp8"),

891 + dequant=is_quantized_kv_cache(self.kv_cache_dtype),

892 kv_cache_layout="NHD",

893 total_tokens=swa_total_tokens,

894 )

@@ -982,7 +983,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

983 cu_seqlens_kv=cu_seqlens_kv[chunk_idx],

984 token_to_batch=token_to_batch[chunk_idx],

985 seq_starts=chunk_starts[chunk_idx],

985 - dequant=self.kv_cache_dtype.startswith("fp8"),

986 + dequant=is_quantized_kv_cache(self.kv_cache_dtype),

987 kv_cache_layout="SHUFFLE"

988 if rocm_aiter_ops.is_shuffle_kv_cache_enabled()

989 else "NHD",

@@ -1081,7 +1082,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

1082 num_actual_tokens = attn_metadata.num_actual_tokens

1083 key_cache, value_cache = kv_cache.unbind(0)

1084

1084 - if self.kv_cache_dtype.startswith("fp8"):

1085 + if is_quantized_kv_cache(self.kv_cache_dtype):

1086 key_cache = key_cache.view(current_platform.fp8_dtype())

1087 value_cache = value_cache.view(current_platform.fp8_dtype())

1088

@@ -1370,7 +1371,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

1371 # key and value may be None in the case of cross attention. They are

1372 # calculated once based on the output from the encoder and then cached

1373 # in KV cache.

1373 - if self.kv_cache_dtype.startswith("fp8"):

1374 + if is_quantized_kv_cache(self.kv_cache_dtype):

1375 key_cache = key_cache.view(current_platform.fp8_dtype())

1376 value_cache = value_cache.view(current_platform.fp8_dtype())

1377 # Reshape the input keys and values and store them in the cache.

@@ -1436,7 +1437,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

1437 key_cache, value_cache = kv_cache.unbind(0)

1438 flash_layout = True

1439

1439 - is_fp8_kv_cache = self.kv_cache_dtype.startswith("fp8")

1440 + is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype)

1441 if is_fp8_kv_cache:

1442 key_cache = key_cache.view(current_platform.fp8_dtype())

1443 value_cache = value_cache.view(current_platform.fp8_dtype())

vllm/v1/attention/backends/rocm_aiter_unified_attn.py

+3 -2

@@ -11,6 +11,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

11 QuantKey,

12 kFp8StaticTensorSym,

13 )

14 +from vllm.utils.torch_utils import is_quantized_kv_cache

15 from vllm.v1.attention.backend import AttentionLayer, AttentionType, MultipleOf

16 from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata

17 from vllm.v1.attention.backends.rocm_attn import (

@@ -200,7 +201,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):

201

202 softmax_scale = self.scale

203 fp8_post_attn_v_rescale = False

203 - if self.kv_cache_dtype.startswith("fp8"):

204 + if is_quantized_kv_cache(self.kv_cache_dtype):

205 key_cache = key_cache.view(self.fp8_dtype)

206 value_cache = value_cache.view(self.fp8_dtype)

207 # When Q is FP8, triton kernel skips K/V dequant (for fp8xfp8 matmul).

@@ -299,7 +300,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):

300 key_cache, value_cache = kv_cache.unbind(0)

301 flash_layout = True

302

302 - is_fp8_kv_cache = self.kv_cache_dtype.startswith("fp8")

303 + is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype)

304 if is_fp8_kv_cache:

305 key_cache = key_cache.view(self.fp8_dtype)

306 value_cache = value_cache.view(self.fp8_dtype)

vllm/v1/attention/backends/rocm_attn.py

+4 -3

@@ -16,6 +16,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

16 kFp8StaticTensorSym,

17 )

18 from vllm.platforms import current_platform

19 +from vllm.utils.torch_utils import is_quantized_kv_cache

20 from vllm.v1.attention.backend import (

21 AttentionBackend,

22 AttentionCGSupport,

@@ -315,7 +316,7 @@ class RocmAttentionImpl(AttentionImpl):

316 layer: The attention layer

317 """

318 # For encoder attention, process FP8 quantization if needed

318 - if self.kv_cache_dtype.startswith("fp8"):

319 + if is_quantized_kv_cache(self.kv_cache_dtype):

320 raise NotImplementedError(

321 "quantization is not supported for encoder attention"

322 )

@@ -406,7 +407,7 @@ class RocmAttentionImpl(AttentionImpl):

407 kv_cache, self.num_kv_heads, self.head_size

408 )

409

409 - if self.kv_cache_dtype.startswith("fp8"):

410 + if is_quantized_kv_cache(self.kv_cache_dtype):

411 key_cache = key_cache.view(self.fp8_dtype)

412 value_cache = value_cache.view(self.fp8_dtype)

413 assert layer._q_scale_float == 1.0, (

@@ -513,7 +514,7 @@ class RocmAttentionImpl(AttentionImpl):

514 )

515 flash_layout = False

516

516 - is_fp8_kv_cache = self.kv_cache_dtype.startswith("fp8")

517 + is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype)

518 if is_fp8_kv_cache:

519 key_cache = key_cache.view(self.fp8_dtype)

520 value_cache = value_cache.view(self.fp8_dtype)

vllm/v1/attention/backends/triton_attn.py

+5 -4

@@ -18,6 +18,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

18 from vllm.platforms import current_platform

19 from vllm.platforms.interface import DeviceCapability

20 from vllm.utils.math_utils import next_power_of_2

21 +from vllm.utils.torch_utils import is_quantized_kv_cache

22 from vllm.v1.attention.backend import (

23 AttentionBackend,

24 AttentionCGSupport,

@@ -472,7 +473,7 @@ class TritonAttentionImpl(AttentionImpl):

473

474 # For decoder and cross-attention, use KV cache as before

475 key_cache, value_cache = kv_cache.unbind(1)

475 - if self.kv_cache_dtype.startswith("fp8"):

476 + if is_quantized_kv_cache(self.kv_cache_dtype):

477 if key_cache.dtype != self.fp8_dtype:

478 key_cache = key_cache.view(self.fp8_dtype)

479 value_cache = value_cache.view(self.fp8_dtype)

@@ -546,7 +547,7 @@ class TritonAttentionImpl(AttentionImpl):

547 layer: The attention layer

548 """

549 # For encoder attention, process FP8 quantization if needed

549 - if self.kv_cache_dtype.startswith("fp8"):

550 + if is_quantized_kv_cache(self.kv_cache_dtype):

551 raise NotImplementedError(

552 "quantization is not supported for encoder attention"

553 )

@@ -588,7 +589,7 @@ class TritonAttentionImpl(AttentionImpl):

589 key_cache, value_cache = kv_cache.unbind(1)

590

591 # Reshape the input keys and values and store them in the cache.

591 - if self.kv_cache_dtype.startswith("fp8"):

592 + if is_quantized_kv_cache(self.kv_cache_dtype):

593 key_cache = key_cache.view(self.fp8_dtype)

594 value_cache = value_cache.view(self.fp8_dtype)

595 # triton kernel does not support uint8 kv_cache

@@ -623,7 +624,7 @@ class TritonAttentionImpl(AttentionImpl):

624 key_cache, value_cache = kv_cache.unbind(1)

625 flash_layout = True

626

626 - is_fp8_kv_cache = self.kv_cache_dtype.startswith("fp8")

627 + is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype)

628 if is_fp8_kv_cache:

629 key_cache = key_cache.view(self.fp8_dtype)

630 value_cache = value_cache.view(self.fp8_dtype)

vllm/v1/attention/ops/triton_reshape_and_cache_flash.py

+11 -8

@@ -5,6 +5,7 @@ import torch

5

6 from vllm.platforms import current_platform

7 from vllm.triton_utils import tl, triton

8 +from vllm.utils.torch_utils import is_quantized_kv_cache

9

10

11 @triton.jit

@@ -145,16 +146,18 @@ def triton_reshape_and_cache_flash(

146 block_stride = key_cache.stride()[0]

147 page_stride = key_cache.stride()[1]

148

148 - assert kv_cache_dtype == "auto" or kv_cache_dtype.startswith("fp8"), (

149 + assert kv_cache_dtype == "auto" or is_quantized_kv_cache(kv_cache_dtype), (

150 f"unsupported kv_cache_dtype (str), got {kv_cache_dtype}."

151 )

152 kv_cache_torch_dtype = (

153 current_platform.fp8_dtype()

153 - if kv_cache_dtype.startswith("fp8")

154 + if is_quantized_kv_cache(kv_cache_dtype)

155 else key_cache.dtype

156 )

157

157 - if key_cache.dtype != kv_cache_torch_dtype and kv_cache_dtype.startswith("fp8"):

158 + if key_cache.dtype != kv_cache_torch_dtype and is_quantized_kv_cache(

159 + kv_cache_dtype

160 + ):

161 # to avoid erounous implicit cast in triton kernel (tl.store to uint8)

162 # (e.g. explicit cast to fp8e4m3fnuz is not supported in triton 3.4)

163 key_cache = key_cache.view(kv_cache_torch_dtype)

@@ -164,7 +167,7 @@ def triton_reshape_and_cache_flash(

167 "uint8 is not supported by triton reshape_and_cache_flash"

168 )

169

167 - FP8_KV_CACHE = kv_cache_dtype.startswith("fp8")

170 + FP8_KV_CACHE = is_quantized_kv_cache(kv_cache_dtype)

171 assert (not FP8_KV_CACHE) or kv_cache_torch_dtype in [

172 torch.float8_e4m3fn,

173 torch.float8_e5m2,

@@ -323,16 +326,16 @@ def triton_reshape_and_cache_flash_diffkv(

326 block_stride = kv_cache.stride()[0]

327 page_stride = kv_cache.stride()[1]

328

326 - assert kv_cache_dtype == "auto" or kv_cache_dtype.startswith("fp8"), (

329 + assert kv_cache_dtype == "auto" or is_quantized_kv_cache(kv_cache_dtype), (

330 f"unsupported kv_cache_dtype (str), got {kv_cache_dtype}."

331 )

332 kv_cache_torch_dtype = (

333 current_platform.fp8_dtype()

331 - if kv_cache_dtype.startswith("fp8")

334 + if is_quantized_kv_cache(kv_cache_dtype)

335 else kv_cache.dtype

336 )

337

335 - if kv_cache.dtype != kv_cache_torch_dtype and kv_cache_dtype.startswith("fp8"):

338 + if kv_cache.dtype != kv_cache_torch_dtype and is_quantized_kv_cache(kv_cache_dtype):

339 # to avoid erounous implicit cast in triton kernel (tl.store to uint8)

340 # (e.g. explicit cast to fp8e4m3fnuz is not supported in triton 3.4)

341 kv_cache = kv_cache.view(kv_cache_torch_dtype)

@@ -341,7 +344,7 @@ def triton_reshape_and_cache_flash_diffkv(

344 "uint8 is not supported by triton reshape_and_cache_flash_diffkv"

345 )

346

344 - FP8_KV_CACHE = kv_cache_dtype.startswith("fp8")

347 + FP8_KV_CACHE = is_quantized_kv_cache(kv_cache_dtype)

348 assert (not FP8_KV_CACHE) or kv_cache_torch_dtype in [

349 torch.float8_e4m3fn,

350 torch.float8_e5m2,

vllm/v1/worker/gpu_model_runner.py

+2 -1

@@ -109,6 +109,7 @@ from vllm.utils.nvtx_pytorch_hooks import PytHooks

109 from vllm.utils.platform_utils import is_pin_memory_available, num_compute_units

110 from vllm.utils.torch_utils import (

111 get_dtype_size,

112 + is_quantized_kv_cache,

113 kv_cache_dtype_str_to_dtype,

114 )

115 from vllm.v1.attention.backend import (

@@ -896,7 +897,7 @@ class GPUModelRunner(

897 If these are left at 0.0 (default after wake_up), all KV cache values

898 become effectively zero, causing gibberish output.

899 """

899 - if not self.cache_config.cache_dtype.startswith("fp8"):

900 + if not is_quantized_kv_cache(self.cache_config.cache_dtype):

901 return

902

903 kv_caches = getattr(self, "kv_caches", [])

vllm/v1/worker/gpu_worker.py

+2 -2

@@ -46,7 +46,7 @@ from vllm.tasks import SupportedTask

46 from vllm.tracing import instrument

47 from vllm.utils.mem_constants import GiB_bytes

48 from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling

49 -from vllm.utils.torch_utils import set_random_seed

49 +from vllm.utils.torch_utils import is_quantized_kv_cache, set_random_seed

50 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput

51 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec

52 from vllm.v1.outputs import (

@@ -197,7 +197,7 @@ class Worker(WorkerBase):

197 # especially the FP8 scaling factor.

198 if (

199 (tags is None or "kv_cache" in tags)

200 - and self.cache_config.cache_dtype.startswith("fp8")

200 + and is_quantized_kv_cache(self.cache_config.cache_dtype)

201 and hasattr(self.model_runner, "init_fp8_kv_scales")

202 ):

203 self.model_runner.init_fp8_kv_scales()

[1/N][Cleanup] Standardize on use of `is_quantized_kv_cache` (#38659)