🐛 remove kv_scale from calls in _custom_ops

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

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
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,

diff --git a/vllm/attention/ops/paged_attn.py b/vllm/attention/ops/paged_attn.py
@@ -76,7 +76,7 @@ def write_to_paged_cache(
             value_cache,
             slot_mapping.flatten(),
             kv_cache_dtype,
-            # kv_scale,
+            kv_scale,
         )
 
     @staticmethod
@@ -123,7 +123,7 @@ def forward_decode(
                 max_context_len,
                 alibi_slopes,
                 kv_cache_dtype,
-                # kv_scale,
+                kv_scale,
             )
         else:
             # Run PagedAttention V2.
@@ -155,7 +155,7 @@ def forward_decode(
                 max_context_len,
                 alibi_slopes,
                 kv_cache_dtype,
-                # kv_scale,
+                kv_scale,
             )
         return output