Skip to content

Commit e0e806f

Browse files
authored
kv-cache : fix unified::seq_rm to work with seq_id < 0 (#13985)
ggml-ci
1 parent 7e00e60 commit e0e806f

File tree

1 file changed

+19
-4
lines changed

1 file changed

+19
-4
lines changed

src/llama-kv-cache-unified.cpp

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -149,12 +149,27 @@ bool llama_kv_cache_unified::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos
149149
p1 = std::numeric_limits<llama_pos>::max();
150150
}
151151

152-
for (uint32_t i = 0; i < cells.size(); ++i) {
153-
if (!cells.pos_in(i, p0, p1)) {
154-
continue;
152+
if (seq_id >= 0) {
153+
for (uint32_t i = 0; i < cells.size(); ++i) {
154+
if (!cells.pos_in(i, p0, p1)) {
155+
continue;
156+
}
157+
158+
if (cells.seq_has(i, seq_id) && cells.seq_rm(i, seq_id)) {
159+
if (new_head == cells.size()) {
160+
new_head = i;
161+
}
162+
}
155163
}
164+
} else {
165+
// match any sequence
166+
for (uint32_t i = 0; i < cells.size(); ++i) {
167+
if (!cells.pos_in(i, p0, p1)) {
168+
continue;
169+
}
170+
171+
cells.rm(i);
156172

157-
if (cells.seq_has(i, seq_id) && cells.seq_rm(i, seq_id)) {
158173
if (new_head == cells.size()) {
159174
new_head = i;
160175
}

0 commit comments

Comments
 (0)