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

Change version

Browse Source
This commit is contained in:
Woosuk Kwon 2024-04-26 05:00:26 +00:00
parent 2aa9831dd3
commit 21f35c2289
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -7,7 +7,7 @@ import jax.numpy as jnp
_PAD_SLOT_ID = -1 _PAD_SLOT_ID = -1
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]
k_cache: jax.Array, # [num_heads, num_blocks * block_size, head_size] k_cache: jax.Array, # [num_heads, num_blocks * block_size, head_size]
@ -27,7 +27,7 @@ def _write_to_kv_cache(
return k_cache, v_cache return k_cache, v_cache
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]
k_cache: jax.Array, # [num_heads, num_blocks * block_size, head_size] k_cache: jax.Array, # [num_heads, num_blocks * block_size, head_size]