fix cache_ops.reshape_and_cache param type compatible issue

aac64dda · huangwb · d6af14b2 · aac64dda
Commit aac64dda authored May 22, 2024 by huangwb
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

server/text_generation_server/utils/paged_attention.py server/text_generation_server/utils/paged_attention.py +2 -1

No files found.
--- a/server/text_generation_server/utils/paged_attention.py
+++ b/server/text_generation_server/utils/paged_attention.py
@@ -27,7 +27,8 @@ def reshape_and_cache(
    elif IS_ROCM_SYSTEM:
        from vllm import cache_ops

-        cache_ops.reshape_and_cache(key, value, key_cache, value_cache, slots)
+        # cache_ops.reshape_and_cache(key, value, key_cache, value_cache, slots)
+        cache_ops.reshape_and_cache(key, value, key_cache, value_cache, slots.int())
    elif IS_XPU_SYSTEM:
        ipex.llm.modules.PagedAttention.reshape_and_cache(
            key, value, key_cache, value_cache, slots