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

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

[Attention] Flash Attention 3 - fp8 (#14570)

Signed-off-by: Mickael Seznec <mickael@mistral.ai>

Mickaël Seznec committed Mar 20, 2025 at 06:14 UTC a597a57595b5b6bb5eca63ed510468a21e683cbd

15 files changed +272 -76

cmake/external_projects/vllm_flash_attn.cmake

+1 -1

@@ -38,7 +38,7 @@ else()

38 FetchContent_Declare(

39 vllm-flash-attn

40 GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git

41 - GIT_TAG 9bfa9869829d8c593527eb34c5271d0090f7ccc9

41 + GIT_TAG dc9d410b3e2d6534a4c70724c2515f4def670a22

42 GIT_PROGRESS TRUE

43 # Don't share the vllm-flash-attn build between build types

44 BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

tests/kernels/test_flash_attn.py

+68 -8

@@ -15,6 +15,7 @@ NUM_HEADS = [(4, 4), (8, 2), (16, 2)]

15 HEAD_SIZES = [128, 256]

16 BLOCK_SIZES = [16, 32]

17 DTYPES = [torch.float16, torch.bfloat16]

18 +QDTYPES = [None, torch.float8_e4m3fn]

19 # one value large enough to test overflow in index calculation.

20 # one value small enough to test the schema op check

21 NUM_BLOCKS = [32768, 2048]

@@ -85,6 +86,7 @@ def ref_paged_attn(

86 @pytest.mark.parametrize("num_blocks", NUM_BLOCKS)

87 @pytest.mark.parametrize("sliding_window", [None, 256])

88 @pytest.mark.parametrize("fa_version", [2, 3])

89 +@pytest.mark.parametrize("q_dtype", QDTYPES)

90 @torch.inference_mode()

91 def test_flash_attn_with_paged_kv(

92 use_out: bool,

@@ -97,11 +99,15 @@ def test_flash_attn_with_paged_kv(

99 num_blocks: int,

100 sliding_window: Optional[int],

101 fa_version: int,

102 + q_dtype: Optional[torch.dtype],

103 ) -> None:

104 torch.set_default_device("cuda")

105 if not is_fa_version_supported(fa_version):

106 pytest.skip(f"Flash attention version {fa_version} not supported due "

107 f"to: \"{fa_version_unsupported_reason(fa_version)}\"")

108 + if q_dtype is not None and (dtype != torch.bfloat16 or fa_version == 2):

109 + pytest.skip("Flash attention with quantized inputs is only "

110 + "supported on version 3 with bfloat16 base type")

111

112 current_platform.seed_everything(0)

113 num_seqs = len(kv_lens)

@@ -130,10 +136,28 @@ def test_flash_attn_with_paged_kv(

136

137 q = query.unsqueeze(1)

138 out = torch.empty_like(q) if use_out else None

139 +

140 + maybe_quantized_query = q

141 + maybe_quantized_key_cache = key_cache

142 + maybe_quantized_value_cache = value_cache

143 + q_descale = None

144 + k_descale = None

145 + v_descale = None

146 + if q_dtype is not None:

147 + # QKV are drawn from N(0, 1): no need for a fp8 scaling factor

148 + maybe_quantized_query = query.to(q_dtype)

149 + maybe_quantized_key_cache = key_cache.to(q_dtype)

150 + maybe_quantized_value_cache = value_cache.to(q_dtype)

151 +

152 + scale_shape = (num_seqs, num_kv_heads)

153 + q_descale = torch.ones(scale_shape, dtype=torch.float32)

154 + k_descale = torch.ones(scale_shape, dtype=torch.float32)

155 + v_descale = torch.ones(scale_shape, dtype=torch.float32)

156 +

157 output = flash_attn_with_kvcache(

134 - q=q,

135 - k_cache=key_cache,

136 - v_cache=value_cache,

158 + q=maybe_quantized_query,

159 + k_cache=maybe_quantized_key_cache,

160 + v_cache=maybe_quantized_value_cache,

161 out=out,

162 softmax_scale=scale,

163 causal=True,

@@ -142,10 +166,17 @@ def test_flash_attn_with_paged_kv(

166 softcap=soft_cap if soft_cap is not None else 0,

167 window_size=window_size,

168 fa_version=fa_version,

169 + q_descale=q_descale,

170 + k_descale=k_descale,

171 + v_descale=v_descale,

172 )

173 output = output if not use_out else out

174 output = output.squeeze(1)

175

176 + atol, rtol = 1.5e-2, 1e-2

177 + if q_dtype is not None:

178 + atol, rtol = 1.5e-1, 1.5e-1

179 +

180 ref_output = ref_paged_attn(query=query,

181 key_cache=key_cache,

182 value_cache=value_cache,

@@ -155,7 +186,7 @@ def test_flash_attn_with_paged_kv(

186 scale=scale,

187 soft_cap=soft_cap,

188 sliding_window=sliding_window)

158 - torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=1e-2), \

189 + torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol), \

190 f"{torch.max(torch.abs(output - ref_output))}"

191

192

@@ -171,6 +202,7 @@ def test_flash_attn_with_paged_kv(

202 @pytest.mark.parametrize("soft_cap", [None, 10.0, 50.0])

203 @pytest.mark.parametrize("num_blocks", NUM_BLOCKS)

204 @pytest.mark.parametrize("fa_version", [2, 3])

205 +@pytest.mark.parametrize("q_dtype", QDTYPES)

206 @torch.inference_mode()

207 def test_varlen_with_paged_kv(

208 use_out: bool,

@@ -183,11 +215,15 @@ def test_varlen_with_paged_kv(

215 soft_cap: Optional[float],

216 num_blocks: int,

217 fa_version: int,

218 + q_dtype: Optional[torch.dtype],

219 ) -> None:

220 torch.set_default_device("cuda")

221 if not is_fa_version_supported(fa_version):

222 pytest.skip(f"Flash attention version {fa_version} not supported due "

223 f"to: \"{fa_version_unsupported_reason(fa_version)}\"")

224 + if q_dtype is not None and (dtype != torch.bfloat16 or fa_version == 2):

225 + pytest.skip("Flash attention with quantized inputs is only "

226 + "supported on version 3 with bfloat16 base type")

227 current_platform.seed_everything(0)

228 num_seqs = len(seq_lens)

229 query_lens = [x[0] for x in seq_lens]

@@ -223,10 +259,28 @@ def test_varlen_with_paged_kv(

259 dtype=torch.int32)

260

261 out = torch.empty_like(query) if use_out else None

262 +

263 + maybe_quantized_query = query

264 + maybe_quantized_key_cache = key_cache

265 + maybe_quantized_value_cache = value_cache

266 + q_descale = None

267 + k_descale = None

268 + v_descale = None

269 + if q_dtype is not None:

270 + # QKV are drawn from N(0, 1): no need for a fp8 scaling factor

271 + maybe_quantized_query = query.to(q_dtype)

272 + maybe_quantized_key_cache = key_cache.to(q_dtype)

273 + maybe_quantized_value_cache = value_cache.to(q_dtype)

274 +

275 + scale_shape = (num_seqs, num_kv_heads)

276 + q_descale = torch.ones(scale_shape, dtype=torch.float32)

277 + k_descale = torch.ones(scale_shape, dtype=torch.float32)

278 + v_descale = torch.ones(scale_shape, dtype=torch.float32)

279 +

280 output = flash_attn_varlen_func(

227 - q=query,

228 - k=key_cache,

229 - v=value_cache,

281 + q=maybe_quantized_query,

282 + k=maybe_quantized_key_cache,

283 + v=maybe_quantized_value_cache,

284 out=out,

285 cu_seqlens_q=cu_query_lens,

286 seqused_k=kv_lens,

@@ -238,6 +292,9 @@ def test_varlen_with_paged_kv(

292 block_table=block_tables,

293 softcap=soft_cap if soft_cap is not None else 0,

294 fa_version=fa_version,

295 + q_descale=q_descale,

296 + k_descale=k_descale,

297 + v_descale=v_descale,

298 )

299 output = output if not use_out else out

300

@@ -252,5 +309,8 @@ def test_varlen_with_paged_kv(

309 sliding_window=sliding_window,

310 soft_cap=soft_cap,

311 )

255 - torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=1e-2), \

312 + atol, rtol = 1.5e-2, 1e-2

313 + if q_dtype is not None:

314 + atol, rtol = 1.5e-1, 1.5e-1

315 + torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol), \

316 f"{torch.max(torch.abs(output - ref_output))}"

vllm/attention/__init__.py

+8 -4

@@ -4,12 +4,16 @@ from vllm.attention.backends.abstract import (AttentionBackend,

4 AttentionMetadata,

5 AttentionMetadataBuilder,

6 AttentionState, AttentionType)

7 -from vllm.attention.backends.utils import get_flash_attn_version

7 from vllm.attention.layer import Attention

8 from vllm.attention.selector import get_attn_backend

10 __all__ = [

12 - "Attention", "AttentionBackend", "AttentionMetadata", "AttentionType",

13 - "AttentionMetadataBuilder", "Attention", "AttentionState",

14 - "get_attn_backend", "get_flash_attn_version"

11 + "Attention",

12 + "AttentionBackend",

13 + "AttentionMetadata",

14 + "AttentionType",

15 + "AttentionMetadataBuilder",

16 + "Attention",

17 + "AttentionState",

18 + "get_attn_backend",

19 ]

vllm/attention/backends/abstract.py

@@ -232,6 +232,7 @@ class AttentionMetadataBuilder(ABC, Generic[T]):

232

233 class AttentionLayer(Protocol):

234

235 + _q_scale: torch.Tensor

236 _k_scale: torch.Tensor

237 _v_scale: torch.Tensor

238 _k_scale_float: float

vllm/attention/backends/flash_attn.py

+68 -11

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

19 # yapf: enable

20 from vllm.attention.backends.utils import (

21 PAD_SLOT_ID, CommonAttentionState, compute_slot_mapping,

22 - compute_slot_mapping_start_idx, get_flash_attn_version,

23 - get_num_prefill_decode_query_kv_tokens, get_seq_len_block_table_args,

24 - is_all_cross_attn_metadata_set, is_all_encoder_attn_metadata_set,

25 - is_block_tables_empty)

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

26 from vllm.logger import init_logger

27 from vllm.multimodal import MultiModalPlaceholderMap

28 from vllm.utils import async_tensor_h2d, make_tensor_with_pad

@@ -630,9 +630,11 @@ class FlashAttentionImpl(AttentionImpl):

630 self.sliding_window = ((sliding_window - 1,

631 0) if sliding_window is not None else (-1, -1))

632 self.kv_cache_dtype = kv_cache_dtype

633 - if is_quantized_kv_cache(self.kv_cache_dtype):

633 + self.vllm_flash_attn_version = get_flash_attn_version()

634 + if (is_quantized_kv_cache(self.kv_cache_dtype)

635 + and self.vllm_flash_attn_version != 3):

636 raise NotImplementedError(

635 - "FlashAttention with FP8 KV cache not yet supported")

637 + "Only FlashAttention3 supports FP8 KV cache")

638 if logits_soft_cap is None:

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

640 logits_soft_cap = 0

@@ -647,7 +649,6 @@ class FlashAttentionImpl(AttentionImpl):

649 f"Head size {head_size} is not supported by FlashAttention. "

650 f"Supported head sizes are: {support_head_sizes}.")

651 self.attn_type = attn_type

650 - self.vllm_flash_attn_version = get_flash_attn_version()

652

653 def forward(

654 self,

@@ -671,13 +672,19 @@ class FlashAttentionImpl(AttentionImpl):

672 for profiling run.

673 attn_metadata: Metadata for attention.

674 NOTE: It in-place updates the output tensor.

675 + NOTE: FP8 quantization, flash-attn expect the size of

676 + {q,k,v}_descale to be (num_sequences, num_kv_heads).

677 + We use torch's .expand() to avoid duplicating values

678 """

675 - # NOTE(woosuk): FlashAttention does not support FP8 KV cache.

676 - assert layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0, (

677 - "key/v_scale is not supported in FlashAttention.")

678 -

679 assert output is not None, "Output tensor must be provided."

680

681 + # NOTE(woosuk): FlashAttention2 does not support FP8 KV cache.

682 + if self.vllm_flash_attn_version < 3 or output.dtype != torch.bfloat16:

683 + assert (

684 + layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0), (

685 + "key/v_scale is only supported in FlashAttention 3 with "

686 + "base dtype bfloat16")

687 +

688 attn_type = self.attn_type

689 if (attn_type == AttentionType.ENCODER

690 and (not attn_metadata.is_all_encoder_attn_metadata_set)):

@@ -694,6 +701,7 @@ class FlashAttentionImpl(AttentionImpl):

701 window_size = self.sliding_window

702 alibi_slopes: Optional[torch.Tensor] = self.alibi_slopes

703 logits_soft_cap: Optional[float] = self.logits_soft_cap

704 + fp8_attention = kv_cache_dtype.startswith("fp8")

705

706 if kv_cache.numel() > 0:

707 key_cache = kv_cache[0]

@@ -729,6 +737,19 @@ class FlashAttentionImpl(AttentionImpl):

737 layer._v_scale,

738 )

739

740 + if fp8_attention:

741 + kv_cache = kv_cache.view(torch.float8_e4m3fn)

742 + key_cache = key_cache.view(torch.float8_e4m3fn)

743 + value_cache = value_cache.view(torch.float8_e4m3fn)

744 +

745 + if fp8_attention:

746 + num_tokens, num_heads, head_size = query.shape

747 + query, _ = ops.scaled_fp8_quant(

748 + query.reshape(

749 + (num_tokens, num_heads * head_size)).contiguous(),

750 + layer._q_scale)

751 + query = query.reshape((num_tokens, num_heads, head_size))

752 +

753 (num_prefill_query_tokens, num_prefill_kv_tokens,

754 num_decode_query_tokens) = \

755 get_num_prefill_decode_query_kv_tokens(attn_metadata, attn_type)

@@ -753,6 +774,23 @@ class FlashAttentionImpl(AttentionImpl):

774 key = key[:num_prefill_kv_tokens]

775 value = value[:num_prefill_kv_tokens]

776

777 + if fp8_attention:

778 + num_kv_tokens, num_kv_heads, head_size = key.shape

779 +

780 + key, _ = ops.scaled_fp8_quant(

781 + key.reshape((num_kv_tokens,

782 + num_kv_heads * head_size)).contiguous(),

783 + layer._k_scale)

784 + key = key.reshape((num_kv_tokens, num_kv_heads, head_size))

785 +

786 + value, _ = ops.scaled_fp8_quant(

787 + value.reshape((num_kv_tokens,

788 + num_kv_heads * head_size)).contiguous(),

789 + layer._v_scale)

790 + value = value.reshape(

791 + (num_kv_tokens, num_kv_heads, head_size))

792 +

793 + descale_shape = (q_seq_start_loc.shape[0] - 1, key.shape[1])

794 flash_attn_varlen_func(

795 q=query,

796 k=key,

@@ -768,13 +806,19 @@ class FlashAttentionImpl(AttentionImpl):

806 softcap=logits_soft_cap,

807 out=prefill_output,

808 fa_version=self.vllm_flash_attn_version,

809 + q_descale=layer._q_scale.expand(descale_shape),

810 + k_descale=layer._k_scale.expand(descale_shape),

811 + v_descale=layer._v_scale.expand(descale_shape),

812 )

813 else:

814 # prefix-enabled attention

815 assert attn_type == AttentionType.DECODER, (

816 "Only decoder-only models support prefix caching")

817 assert prefill_meta.seq_lens is not None

818 + assert prefill_meta.query_start_loc is not None

819 max_seq_len = max(prefill_meta.seq_lens)

820 + descale_shape = (prefill_meta.query_start_loc.shape[0] - 1,

821 + key.shape[1])

822 flash_attn_varlen_func( # noqa

823 q=query,

824 k=key_cache,

@@ -791,6 +835,9 @@ class FlashAttentionImpl(AttentionImpl):

835 softcap=logits_soft_cap,

836 out=prefill_output,

837 fa_version=self.vllm_flash_attn_version,

838 + q_descale=layer._q_scale.expand(descale_shape),

839 + k_descale=layer._k_scale.expand(descale_shape),

840 + v_descale=layer._v_scale.expand(descale_shape),

841 )

842

843 if decode_meta := attn_metadata.decode_metadata:

@@ -804,6 +851,9 @@ class FlashAttentionImpl(AttentionImpl):

851 assert attn_type == AttentionType.DECODER, (

852 "Only decoder-only models support max_decode_query_len > 1"

853 )

854 + assert decode_meta.query_start_loc is not None

855 + descale_shape = (decode_meta.query_start_loc.shape[0] - 1,

856 + key.shape[1])

857 flash_attn_varlen_func(

858 q=decode_query,

859 k=key_cache,

@@ -820,6 +870,9 @@ class FlashAttentionImpl(AttentionImpl):

870 block_table=decode_meta.block_tables,

871 out=decode_output,

872 fa_version=self.vllm_flash_attn_version,

873 + q_descale=layer._q_scale.expand(descale_shape),

874 + k_descale=layer._k_scale.expand(descale_shape),

875 + v_descale=layer._v_scale.expand(descale_shape),

876 )

877 else:

878 # Use flash_attn_with_kvcache for normal decoding.

@@ -828,6 +881,7 @@ class FlashAttentionImpl(AttentionImpl):

881 _,

882 block_tables_arg,

883 ) = get_seq_len_block_table_args(decode_meta, False, attn_type)

884 + descale_shape = (seq_lens_arg.shape[0], key_cache.shape[-2])

885 flash_attn_with_kvcache(

886 q=decode_query.unsqueeze(1),

887 k_cache=key_cache,

@@ -841,6 +895,9 @@ class FlashAttentionImpl(AttentionImpl):

895 softcap=logits_soft_cap,

896 out=decode_output.unsqueeze(1),

897 fa_version=self.vllm_flash_attn_version,

898 + q_descale=layer._q_scale.expand(descale_shape),

899 + k_descale=layer._k_scale.expand(descale_shape),

900 + v_descale=layer._v_scale.expand(descale_shape),

901 )

902 return output

903

vllm/attention/backends/mla/common.py

+1 -1

@@ -203,9 +203,9 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionLayer,

203 AttentionState, MLAAttentionImpl)

204 from vllm.attention.backends.utils import (PAD_SLOT_ID, compute_slot_mapping,

205 compute_slot_mapping_start_idx,

206 - get_flash_attn_version,

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

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

210 LinearBase, RowParallelLinear,

211 UnquantizedLinearMethod)

vllm/attention/backends/utils.py

-34

@@ -8,13 +8,11 @@ from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Type, TypeVar, Union

8 import numpy as np

9 import torch

11 -from vllm import envs

11 from vllm.attention import (AttentionMetadata, AttentionMetadataBuilder,

12 AttentionState)

13 from vllm.attention.backends.abstract import AttentionType

14 from vllm.logger import init_logger

15 from vllm.multimodal import MultiModalPlaceholderMap

17 -from vllm.platforms import current_platform

16 from vllm.utils import async_tensor_h2d, make_tensor_with_pad

18 logger = init_logger(__name__)

@@ -585,35 +583,3 @@ def get_num_prefill_decode_query_kv_tokens(

583

584 return (num_prefill_query_tokens, num_prefill_kv_tokens,

585 num_decode_query_tokens)

588 -

589 -

590 -def get_flash_attn_version():

591 - try:

592 - from vllm.vllm_flash_attn.flash_attn_interface import (

593 - fa_version_unsupported_reason, is_fa_version_supported)

594 -

595 - # if hopper default to FA3, otherwise stick to FA2 for now

596 - # TODO(lucas): profile FA3 on ampere to see if it makes sense to

597 - # use FA3 as default for both

598 - if current_platform.get_device_capability()[0] == 9:

599 - fa_version = 3 if is_fa_version_supported(3) else 2

600 - else:

601 - fa_version = 2

602 -

603 - if envs.VLLM_FLASH_ATTN_VERSION is not None:

604 - assert envs.VLLM_FLASH_ATTN_VERSION in [2, 3]

605 - fa_version = envs.VLLM_FLASH_ATTN_VERSION

606 - if (current_platform.get_device_capability()[0] == 10

607 - and envs.VLLM_FLASH_ATTN_VERSION == 3):

608 - logger.warning("Cannot use FA version 3 on Blackwell platform",

609 - "defaulting to FA version 2.")

610 - fa_version = 2

611 -

612 - if not is_fa_version_supported(fa_version):

613 - logger.error("Cannot use FA version %d is not supported due to %s",

614 - fa_version, fa_version_unsupported_reason(fa_version))

615 -

616 - assert is_fa_version_supported(fa_version)

617 - return fa_version

618 - except (ImportError, AssertionError):

619 - return None

vllm/attention/layer.py

+7 -2

@@ -84,6 +84,9 @@ class Attention(nn.Module):

84 self.calculate_kv_scales = calculate_kv_scales

85 self._k_scale = torch.tensor(1.0, dtype=torch.float32)

86 self._v_scale = torch.tensor(1.0, dtype=torch.float32)

87 + # FlashAttn doesn't support quantizing the kv-cache only

88 + # but requires q to be quantized as well.

89 + self._q_scale = torch.tensor(1.0, dtype=torch.float32)

91 # We also keep the float32 versions of k/v_scale for attention

92 # backends that don't support tensors (Flashinfer)

@@ -153,6 +156,7 @@ class Attention(nn.Module):

156 ).parallel_config.pipeline_parallel_size)

157 ]

158

159 + self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)

160 self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)

161 self.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32)

162

@@ -178,7 +182,7 @@ class Attention(nn.Module):

182 if self.calculate_kv_scales:

183 attn_metadata = get_forward_context().attn_metadata

184 if attn_metadata.enable_kv_scales_calculation:

181 - self.calc_kv_scales(key, value)

185 + self.calc_kv_scales(query, key, value)

186 if self.use_output:

187 output_shape = (output_shape

188 if output_shape is not None else query.shape)

@@ -225,7 +229,8 @@ class Attention(nn.Module):

229 return torch.ops.vllm.unified_attention(

230 query, key, value, self.layer_name)

231

228 - def calc_kv_scales(self, key, value):

232 + def calc_kv_scales(self, query, key, value):

233 + self._q_scale.copy_(torch.abs(query).max() / self.q_range)

234 self._k_scale.copy_(torch.abs(key).max() / self.k_range)

235 self._v_scale.copy_(torch.abs(value).max() / self.v_range)

236 self._k_scale_float = self._k_scale.item()

vllm/envs.py

+6 -1

@@ -78,6 +78,7 @@ if TYPE_CHECKING:

78 VLLM_ENABLE_V1_MULTIPROCESSING: bool = True

79 VLLM_LOG_BATCHSIZE_INTERVAL: float = -1

80 VLLM_DISABLE_COMPILE_CACHE: bool = False

81 + Q_SCALE_CONSTANT: int = 200

82 K_SCALE_CONSTANT: int = 200

83 V_SCALE_CONSTANT: int = 100

84 VLLM_SERVER_DEV_MODE: bool = False

@@ -524,13 +525,17 @@ environment_variables: dict[str, Callable[[], Any]] = {

525 # Pad the fp8 weights to 256 bytes for ROCm

526 "VLLM_ROCM_FP8_PADDING":

527 lambda: bool(int(os.getenv("VLLM_ROCM_FP8_PADDING", "1"))),

528 +

529 + # Divisor for dynamic query scale factor calculation for FP8 KV Cache

530 + "Q_SCALE_CONSTANT":

531 + lambda: int(os.getenv("Q_SCALE_CONSTANT", "200")),

532 # Divisor for dynamic key scale factor calculation for FP8 KV Cache

533 "K_SCALE_CONSTANT":

534 lambda: int(os.getenv("K_SCALE_CONSTANT", "200")),

530 -

535 # Divisor for dynamic value scale factor calculation for FP8 KV Cache

536 "V_SCALE_CONSTANT":

537 lambda: int(os.getenv("V_SCALE_CONSTANT", "100")),

538 +

539 # If set, enable multiprocessing in LLM for the V1 code path.

540 "VLLM_ENABLE_V1_MULTIPROCESSING":

541 lambda: bool(int(os.getenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1"))),

vllm/fa_utils.py new

+42

@@ -0,0 +1,42 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +from typing import Optional

3 +

4 +from vllm import envs

5 +from vllm.logger import init_logger

6 +

7 +logger = init_logger(__name__)

8 +

9 +

10 +def get_flash_attn_version() -> Optional[int]:

11 + # import here to avoid circular dependencies

12 + from vllm.platforms import current_platform

13 + try:

14 + from vllm.vllm_flash_attn.flash_attn_interface import (

15 + fa_version_unsupported_reason, is_fa_version_supported)

16 + device_capability = current_platform.get_device_capability()

17 +

18 + assert device_capability is not None

19 +

20 + # 1. default version depending on platform

21 + fa_version = 3 if (device_capability.major == 9

22 + and is_fa_version_supported(3)) else 2

23 +

24 + # 2. override if passed by environment

25 + if envs.VLLM_FLASH_ATTN_VERSION is not None:

26 + assert envs.VLLM_FLASH_ATTN_VERSION in [2, 3]

27 + fa_version = envs.VLLM_FLASH_ATTN_VERSION

28 +

29 + # 3. fallback for unsupported combinations

30 + if device_capability.major == 10 and fa_version == 3:

31 + logger.warning("Cannot use FA version 3 on Blackwell platform",

32 + "defaulting to FA version 2.")

33 + fa_version = 2

34 +

35 + if not is_fa_version_supported(fa_version):

36 + logger.error("Cannot use FA version %d is not supported due to %s",

37 + fa_version, fa_version_unsupported_reason(fa_version))

38 +

39 + assert is_fa_version_supported(fa_version)

40 + return fa_version

41 + except (ImportError, AssertionError):

42 + return None

vllm/model_executor/layers/quantization/kv_cache.py

+13 -3

@@ -26,11 +26,14 @@ class BaseKVCacheMethod(QuantizeMethodBase):

27 def create_weights(self, layer: torch.nn.Module):

28 """

29 - Create "weight" (aka k_scale and v_scale) for an attention layer.

29 + Create "weight" (aka q_scale, k_scale and v_scale)

30 + for an attention layer.

31 """

31 - # Initialize the KV cache scales to -1.0, which is an invalid value.

32 - # If the k/v_scale appears in the checkpoint, it will be

32 + # Initialize the Q and KV cache scales to -1.0, an invalid value.

33 + # If the q and k/v_scales appear in the checkpoint, it will be

34 # overwritten when loading weights.

35 + layer.q_scale = torch.nn.Parameter(torch.tensor(-1.0),

36 + requires_grad=False)

37 layer.k_scale = torch.nn.Parameter(torch.tensor(-1.0),

38 requires_grad=False)

39 layer.v_scale = torch.nn.Parameter(torch.tensor(-1.0),

@@ -75,6 +78,13 @@ class BaseKVCacheMethod(QuantizeMethodBase):

78 raise ValueError("Only support per-tensor scaling factor "

79 "for fp8 KV cache")

81 + if layer.q_scale < 0.0:

82 + logger.warning_once(

83 + "Checkpoint does not provide a q scaling factor. "

84 + "Setting it to k_scale. This only matters for "

85 + "the flash-attn backend.")

86 + layer._q_scale.copy_(k_scale)

87 +

88 # These are used in the final Attention.forward()

89 layer._k_scale.copy_(k_scale)

90 layer._v_scale.copy_(v_scale)

vllm/platforms/cuda.py

+12 -9

@@ -14,6 +14,7 @@ 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

18 from vllm.logger import init_logger

19 from vllm.utils import import_pynvml

@@ -240,15 +241,6 @@ class CudaPlatformBase(Platform):

241 "Cannot use FlashAttention-2 backend for dtype other than "

242 "torch.float16 or torch.bfloat16.")

243 target_backend = _Backend.XFORMERS

243 - elif kv_cache_dtype is not None and \

244 - kv_cache_dtype.startswith("fp8"):

245 - logger.info(

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

247 - logger.warning(

248 - "Please use FlashInfer backend with FP8 KV Cache for "

249 - "better performance by setting environment variable "

250 - "VLLM_ATTENTION_BACKEND=FLASHINFER")

251 - target_backend = _Backend.XFORMERS

244 elif block_size % 16 != 0:

245 logger.info(

246 "Cannot use FlashAttention-2 backend for block size not "

@@ -270,6 +262,17 @@ class CudaPlatformBase(Platform):

262 "Cannot use FlashAttention-2 backend for head size %d.",

263 head_size)

264 target_backend = _Backend.XFORMERS

265 + fp8_kv_cache = (kv_cache_dtype is not None

266 + and kv_cache_dtype.startswith("fp8"))

267 + if (fp8_kv_cache and get_flash_attn_version() != 3):

268 + logger.info(

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

270 + )

271 + logger.warning(

272 + "Please use FlashInfer backend with FP8 KV Cache for "

273 + "better performance by setting environment variable "

274 + "VLLM_ATTENTION_BACKEND=FLASHINFER")

275 + target_backend = _Backend.XFORMERS

276 except ImportError:

277 logger.info(

278 "Cannot use FlashAttention-2 backend because the "

vllm/v1/attention/backends/flash_attn.py

+40 -1

@@ -6,11 +6,12 @@ from typing import TYPE_CHECKING, Any, Optional

6 import numpy as np

7 import torch

9 +from vllm import _custom_ops as ops

10 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,

11 AttentionMetadata, AttentionType,

12 is_quantized_kv_cache)

12 -from vllm.attention.backends.utils import get_flash_attn_version

13 from vllm.attention.ops.triton_merge_attn_states import merge_attn_states

14 +from vllm.fa_utils import get_flash_attn_version

15 from vllm.logger import init_logger

16 from vllm.platforms import current_platform

17 from vllm.utils import cdiv

@@ -226,6 +227,9 @@ class FlashAttentionImpl(AttentionImpl):

227 attn_metadata: Metadata for attention.

228 Returns:

229 shape = [num_tokens, num_heads * head_size]

230 + NOTE: FP8 quantization, flash-attn expect the size of

231 + {q,k,v}_descale to be (num_sequences, num_kv_heads).

232 + We use torch's .expand() to avoid duplicating values

233 """

234 assert output is not None, "Output tensor must be provided."

235

@@ -259,6 +263,17 @@ class FlashAttentionImpl(AttentionImpl):

263 layer._k_scale,

264 layer._v_scale,

265 )

266 + descale_shape = (attn_metadata.query_start_loc.shape[0] - 1,

267 + key.shape[1])

268 + if self.kv_cache_dtype.startswith("fp8"):

269 + key_cache = key_cache.view(torch.float8_e4m3fn)

270 + value_cache = value_cache.view(torch.float8_e4m3fn)

271 + num_tokens, num_heads, head_size = query.shape

272 + query, _ = ops.scaled_fp8_quant(

273 + query.reshape(

274 + (num_tokens, num_heads * head_size)).contiguous(),

275 + layer._q_scale)

276 + query = query.reshape((num_tokens, num_heads, head_size))

277

278 # Compute attention and update output up to `num_actual_tokens`.

279 if not attn_metadata.use_cascade:

@@ -279,6 +294,9 @@ class FlashAttentionImpl(AttentionImpl):

294 block_table=attn_metadata.block_table,

295 softcap=self.logits_soft_cap,

296 fa_version=self.vllm_flash_attn_version,

297 + q_descale=layer._q_scale.expand(descale_shape),

298 + k_descale=layer._k_scale.expand(descale_shape),

299 + v_descale=layer._v_scale.expand(descale_shape),

300 )

301 return output

302

@@ -301,6 +319,9 @@ class FlashAttentionImpl(AttentionImpl):

319 block_table=attn_metadata.block_table,

320 common_prefix_len=attn_metadata.common_prefix_len,

321 fa_version=self.vllm_flash_attn_version,

322 + q_descale=layer._q_scale,

323 + k_descale=layer._k_scale,

324 + v_descale=layer._v_scale,

325 )

326 return output

327

@@ -391,6 +412,9 @@ def cascade_attention(

412 block_table: torch.Tensor,

413 common_prefix_len: int,

414 fa_version: int,

415 + q_descale: Optional[torch.Tensor] = None,

416 + k_descale: Optional[torch.Tensor] = None,

417 + v_descale: Optional[torch.Tensor] = None,

418 ) -> torch.Tensor:

419 assert alibi_slopes is None, ("Cascade attention does not support ALiBi.")

420 # TODO: Support sliding window.

@@ -402,6 +426,7 @@ def cascade_attention(

426 assert common_prefix_len % block_size == 0

427 num_common_kv_blocks = common_prefix_len // block_size

428 assert num_common_kv_blocks > 0

429 + descale_shape = (cu_prefix_query_lens.shape[0] - 1, key_cache.shape[-2])

430

431 # Process shared prefix.

432 prefix_output, prefix_lse = flash_attn_varlen_func(

@@ -419,8 +444,16 @@ def cascade_attention(

444 softcap=logits_soft_cap,

445 return_softmax_lse=True,

446 fa_version=fa_version,

447 + q_descale=q_descale.expand(descale_shape)

448 + if q_descale is not None else None,

449 + k_descale=k_descale.expand(descale_shape)

450 + if k_descale is not None else None,

451 + v_descale=v_descale.expand(descale_shape)

452 + if v_descale is not None else None,

453 )

454

455 + descale_shape = (cu_query_lens.shape[0] - 1, key_cache.shape[-2])

456 +

457 # Process suffix per query.

458 suffix_output, suffix_lse = flash_attn_varlen_func(

459 q=query,

@@ -437,6 +470,12 @@ def cascade_attention(

470 softcap=logits_soft_cap,

471 return_softmax_lse=True,

472 fa_version=fa_version,

473 + q_descale=q_descale.expand(descale_shape)

474 + if q_descale is not None else None,

475 + k_descale=k_descale.expand(descale_shape)

476 + if k_descale is not None else None,

477 + v_descale=v_descale.expand(descale_shape)

478 + if v_descale is not None else None,

479 )

480

481 # Merge prefix and suffix outputs, and store the result in output.

vllm/v1/executor/multiproc_executor.py

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

5 import signal

6 import sys

7 import time

8 +import traceback

9 import weakref

10 from dataclasses import dataclass

11 from enum import Enum, auto

@@ -370,6 +371,9 @@ class WorkerProc:

371 func = partial(cloudpickle.loads(method), self.worker)

372 output = func(*args, **kwargs)

373 except Exception as e:

374 + # Notes have been introduced in python 3.11

375 + if hasattr(e, "add_note"):

376 + e.add_note(traceback.format_exc())

377 self.worker_response_mq.enqueue(

378 (WorkerProc.ResponseStatus.FAILURE, e))

379 logger.exception("WorkerProc hit an exception: %s", exc_info=e)

vllm/v1/worker/gpu_model_runner.py

+1 -1

@@ -1558,7 +1558,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1558 block_size=block_size,

1559 num_kv_heads=attn_module.num_kv_heads,

1560 head_size=attn_module.head_size,

1561 - dtype=attn_module.dtype,

1561 + dtype=self.kv_cache_dtype,

1562 use_mla=use_mla)

1563 elif attn_module.attn_type in (AttentionType.ENCODER,

1564 AttentionType.ENCODER_ONLY):