🐛 revert kv_scale hacks, build c code ourselves

Signed-off-by: Joe Runde <Joseph.Runde@ibm.com>
IBM · Apr 11, 2024 · 3b914ed · 3b914ed
1 parent de60593
commit 3b914ed
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/Dockerfile.ubi b/Dockerfile.ubi
@@ -257,7 +257,7 @@ FROM base AS vllm
 WORKDIR /vllm-staging
 # COPY files from various places into a staging directory
 COPY --link vllm vllm
-COPY --from=prebuilt-wheel --link /workspace/vllm/*.so vllm/
+COPY --from=build --link /workspace/vllm/*.so vllm/
 COPY --from=gen-protos --link /workspace/vllm/entrypoints/grpc/pb vllm/entrypoints/grpc/pb
 
 # custom COPY command to use umask to control permissions and grant permissions

diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py
@@ -49,7 +49,7 @@ def paged_attention_v1(
     vllm_ops.paged_attention_v1(out, query, key_cache, value_cache,
                                 num_kv_heads, scale, block_tables,
                                 context_lens, block_size, max_context_len,
-                                alibi_slopes, kv_cache_dtype)#, kv_scale)
+                                alibi_slopes, kv_cache_dtype, kv_scale)
 
 
 def paged_attention_v2(
@@ -73,8 +73,8 @@ def paged_attention_v2(
     vllm_ops.paged_attention_v2(out, exp_sum, max_logits, tmp_out, query,
                                 key_cache, value_cache, num_kv_heads, scale,
                                 block_tables, context_lens, block_size,
-                                max_context_len, alibi_slopes, kv_cache_dtype)
-                                # kv_scale)
+                                max_context_len, alibi_slopes, kv_cache_dtype,
+                                kv_scale)
 
 
 # pos encoding ops
@@ -173,7 +173,7 @@ def reshape_and_cache(
     kv_scale: float,
 ) -> None:
     vllm_cache_ops.reshape_and_cache(key, value, key_cache, value_cache,
-                                     slot_mapping, kv_cache_dtype)#, kv_scale)
+                                     slot_mapping, kv_cache_dtype, kv_scale)
 
 
 def copy_blocks(key_caches: torch.Tensor, value_caches: torch.Tensor,