Skip to content

Commit f181ce2

Browse files
authored
[Bugfix] fix _kv_cache_transpose_append buffer read region error (#2277)
* improve Install via environment variable * [HotFix] fix kv_cache_transpose_append buffer region
1 parent 0ca6b33 commit f181ce2

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

python/mlc_llm/nn/kv_cache.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -399,7 +399,7 @@ def tir_kv_cache_transpose_append(
399399
pages[T.floordiv(position, 16), 0, vh, T.floormod(position, 16), vf] = k_data[vgpos, vh, vf]
400400
with T.block("v_transpose_append"):
401401
vgpos, vh, vf = T.axis.remap("SSS", [global_pos, h, f])
402-
T.reads(position_map[vgpos], k_data[vgpos, vh, vf])
402+
T.reads(position_map[vgpos], v_data[vgpos, vh, vf])
403403
T.writes(pages[position_map[vgpos] // 16, 1, vh, position_map[vgpos] % 16, vf])
404404
position: T.int32 = position_map[vgpos] # type: ignore[name-defined,no-redef]
405405
pages[T.floordiv(position, 16), 1, vh, T.floormod(position, 16), vf] = v_data[vgpos, vh, vf]

0 commit comments

Comments
 (0)