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 14, 22:35 UTC
[BugFix] fix: aot passes kvcache dtype information (#19750)
Signed-off-by: Mickael Seznec <mickael@mistral.ai>
Mickaël Seznec committed
Aug 1, 2025 at 07:45 UTC
e1a7fe4af5e9c287501c648e64956a08705af86a
1 file changed
+21
-4
vllm/v1/attention/backends/flash_attn.py
+21
-4
@@ -99,6 +99,13 @@ class FlashAttentionBackend(AttentionBackend):
99
raise ValueError(f"Unknown cache layout format {cache_layout}.")
100
return stride_order
101
102
+ @staticmethod
103
+ def get_fp8_dtype_for_flashattn(kv_cache_dtype: str) -> torch.dtype:
104
+ if kv_cache_dtype in ("fp8", "fp8_e4m3"):
105
+ return torch.float8_e4m3fn
106
+ else:
107
+ raise ValueError(f"Unrecognized FP8 dtype: {kv_cache_dtype}")
108
+
109
110
@dataclass
111
class FlashAttentionMetadata:
@@ -161,6 +168,7 @@ class FlashAttentionMetadataBuilder(
168
self.parallel_config)
169
self.num_heads_kv = self.model_config.get_num_kv_heads(
170
self.parallel_config)
171
+ self.kv_cache_dtype = kv_cache_spec.dtype
172
self.headdim = self.model_config.get_head_size()
173
self.block_size = kv_cache_spec.block_size
174
@@ -239,17 +247,24 @@ class FlashAttentionMetadataBuilder(
247
248
def schedule(batch_size, cu_query_lens, max_query_len, seqlens,
249
max_seq_len, causal):
250
+ cache_dtype = self.cache_config.cache_dtype
251
+ if cache_dtype.startswith("fp8"):
252
+ qkv_dtype = FlashAttentionBackend.get_fp8_dtype_for_flashattn(
253
+ cache_dtype)
254
+ else:
255
+ qkv_dtype = self.kv_cache_dtype
256
if aot_schedule:
257
return get_scheduler_metadata(
258
batch_size=batch_size,
259
max_seqlen_q=max_query_len,
260
max_seqlen_k=max_seq_len,
247
- cache_seqlens=seqlens,
261
num_heads_q=self.num_heads_q,
262
num_heads_kv=self.num_heads_kv,
263
headdim=self.headdim,
251
- page_size=self.block_size,
264
+ cache_seqlens=seqlens,
265
+ qkv_dtype=qkv_dtype,
266
cu_seqlens_q=cu_query_lens,
267
+ page_size=self.block_size,
268
causal=causal,
269
window_size=self.aot_sliding_window,
270
num_splits=self.max_num_splits,
@@ -474,8 +489,10 @@ class FlashAttentionImpl(AttentionImpl):
489
)
490
491
if self.kv_cache_dtype.startswith("fp8"):
477
- key_cache = key_cache.view(torch.float8_e4m3fn)
478
- value_cache = value_cache.view(torch.float8_e4m3fn)
492
+ dtype = FlashAttentionBackend.get_fp8_dtype_for_flashattn(
493
+ self.kv_cache_dtype)
494
+ key_cache = key_cache.view(dtype)
495
+ value_cache = value_cache.view(dtype)
496
num_tokens, num_heads, head_size = query.shape
497
query, _ = ops.scaled_fp8_quant(
498
query.reshape(