Skip to content

GitLab

Commit b550cf96 authored by zhuwenwen's avatar zhuwenwen
Browse files

update reshape_and_cache_cuda import

parent 7ff04b72
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
vllm/v1/attention/backends/flash_attn.py
View file @ b550cf96
...@@ -566,8 +566,8 @@ class FlashAttentionImpl(AttentionImpl): ...@@ -566,8 +566,8 @@ class FlashAttentionImpl(AttentionImpl):
layer._v_scale, layer._v_scale,
) )
else: else:
from lightop import reshape_and_cache_cuda
if envs.VLLM_USE_OPT_RESHAPE_AND_CACHE and key.dtype == value.dtype == "fp16": if envs.VLLM_USE_OPT_RESHAPE_AND_CACHE and key.dtype == value.dtype == "fp16":
from lightop import reshape_and_cache_cuda
reshape_and_cache_cuda( reshape_and_cache_cuda(
key, value, key, value,
key_cache, value_cache, key_cache, value_cache,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment