xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-06-06 07:02:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix cache write

Browse Source
This commit is contained in:
Woosuk Kwon 2024-04-24 08:56:30 +00:00
parent d5fb1c20c1
commit 620e7646d3
1 changed files with 1 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
vllm/model_executor/models/jax/ops/write_to_cache.py
View File

@ -5,16 +5,6 @@ import jax.numpy as jnp
_PAD_SLOT_ID = -1 _PAD_SLOT_ID = -1
def write_to_kv_cache(
key: jax.Array, # [batch_size, seq_len, num_heads, head_size]
value: jax.Array, # [batch_size, seq_len, num_heads, head_size]
kv_cache: jax.Array, # [2, num_heads, num_blocks, block_size, head_size]
slot_mapping: jax.Array, # [batch_size, seq_len]
) -> jax.Array:
f = _write_to_kv_cache
return f(key, value, kv_cache, slot_mapping)
def _write_to_kv_cache( def _write_to_kv_cache(
key: jax.Array, # [batch_size, seq_len, num_heads, head_size] key: jax.Array, # [batch_size, seq_len, num_heads, head_size]
value: jax.Array, # [batch_size, seq_len, num_heads, head_size] value: jax.Array, # [batch_size, seq_len, num_heads, head_size]
@ -33,7 +23,7 @@ def _write_to_kv_cache(
return kv_cache return kv_cache
def _write_to_kv_cache_in_place( def write_to_kv_cache(
key: jax.Array, # [batch_size, seq_len, num_heads, head_size] key: jax.Array, # [batch_size, seq_len, num_heads, head_size]
value: jax.Array, # [batch_size, seq_len, num_heads, head_size] value: jax.Array, # [batch_size, seq_len, num_heads, head_size]
kv_cache: jax.Array, # [2, num_heads, num_blocks, block_size, head_size] kv_cache: jax.Array, # [2, num_heads, num_blocks, block_size, head_size]