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, 11:00 UTC
[FP8][ROCm][Attention] Enable FP8 KV cache on ROCm for V1 (#17870)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Gregory Shtrasberg committed
May 11, 2025 at 03:58 UTC
06c0922a69c16761949b5a9c4cc1723de1fc25f4
3 files changed
+17
-8
vllm/attention/ops/chunked_prefill_paged_decode.py
+2
-1
@@ -9,6 +9,7 @@
9
import torch
10
11
from vllm import _custom_ops as ops
12
+from vllm.platforms import current_platform
13
from vllm.platforms.rocm import use_rocm_custom_paged_attention
14
from vllm.triton_utils import tl, triton
15
@@ -267,7 +268,7 @@ def chunked_prefill_paged_decode(
268
assert value_cache.dtype == torch.uint8
269
270
if kv_cache_dtype in ("fp8", "fp8_e4m3"):
270
- target_dtype = torch.float8_e4m3fn
271
+ target_dtype = current_platform.fp8_dtype()
272
elif kv_cache_dtype == "fp8_e5m2":
273
target_dtype = torch.float8_e5m2
274
else:
vllm/engine/arg_utils.py
+3
-1
@@ -1205,7 +1205,9 @@ class EngineArgs:
1205
and not envs.is_set("VLLM_ATTENTION_BACKEND")
1206
) or envs.VLLM_ATTENTION_BACKEND == "FLASH_ATTN_VLLM_V1"
1207
supported = False
1208
- if fp8_attention and will_use_fa:
1208
+ if current_platform.is_rocm():
1209
+ supported = True
1210
+ elif fp8_attention and will_use_fa:
1211
from vllm.attention.utils.fa_utils import (
1212
flash_attn_supports_fp8)
1213
supported = flash_attn_supports_fp8()
vllm/v1/attention/backends/triton_attn.py
+12
-6
@@ -9,6 +9,7 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
9
AttentionMetadata, AttentionType)
10
from vllm.attention.ops.triton_unified_attention import unified_attention
11
from vllm.logger import init_logger
12
+from vllm.platforms import current_platform
13
from vllm.v1.attention.backends.flash_attn import (
14
FlashAttentionMetadata, FlashAttentionMetadataBuilder)
15
@@ -108,6 +109,8 @@ class TritonAttentionImpl(AttentionImpl):
109
"are not implemented for "
110
"TritonAttentionImpl")
111
112
+ self.fp8_dtype = current_platform.fp8_dtype()
113
+
114
def forward(
115
self,
116
layer: torch.nn.Module,
@@ -161,15 +164,18 @@ class TritonAttentionImpl(AttentionImpl):
164
)
165
166
if self.kv_cache_dtype.startswith("fp8"):
164
- key_cache = key_cache.view(torch.float8_e4m3fn)
165
- value_cache = value_cache.view(torch.float8_e4m3fn)
167
+ key_cache = key_cache.view(self.fp8_dtype)
168
+ value_cache = value_cache.view(self.fp8_dtype)
169
num_tokens, num_heads, head_size = query.shape
170
assert layer._q_scale == 1.0, \
171
"A non 1.0 q_scale is not currently supported."
169
- query, _ = ops.scaled_fp8_quant(
170
- query.reshape(
171
- (num_tokens, num_heads * head_size)).contiguous(),
172
- layer._q_scale)
172
+ if not current_platform.is_rocm():
173
+ # Skip Q quantization on ROCm, since dequantizing back to
174
+ # f32 in the attention kernel is not supported.
175
+ query, _ = ops.scaled_fp8_quant(
176
+ query.reshape(
177
+ (num_tokens, num_heads * head_size)).contiguous(),
178
+ layer._q_scale)
179
query = query.reshape((num_tokens, num_heads, head_size))
180
181
use_local_attn = \