| 1968 | } |
| 1969 | |
| 1970 | bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, const slot_info & sinfo) { |
| 1971 | auto & cells = v_cells[strm]; |
| 1972 | |
| 1973 | uint32_t v_trans; |
| 1974 | uint32_t n_layer; |
| 1975 | |
| 1976 | io.read_to(&v_trans, sizeof(v_trans)); |
| 1977 | io.read_to(&n_layer, sizeof(n_layer)); |
| 1978 | |
| 1979 | if (n_layer != layers.size()) { |
| 1980 | LLAMA_LOG_ERROR("%s: mismatched layer count (%u instead of %u)\n", __func__, n_layer, (uint32_t) layers.size()); |
| 1981 | return false; |
| 1982 | } |
| 1983 | |
| 1984 | if (cell_count > cells.size()) { |
| 1985 | LLAMA_LOG_ERROR("%s: not enough cells in kv cache to restore state (%u > %u)\n", __func__, cell_count, cells.size()); |
| 1986 | return false; |
| 1987 | } |
| 1988 | |
| 1989 | if (this->v_trans != (bool) v_trans) { |
| 1990 | LLAMA_LOG_ERROR("%s: incompatible V transposition\n", __func__); |
| 1991 | return false; |
| 1992 | } |
| 1993 | |
| 1994 | // For each layer, read the keys for each cell, one row is one cell, read as one contiguous block |
| 1995 | for (const auto & layer : layers) { |
| 1996 | const uint32_t il = layer.il; |
| 1997 | |
| 1998 | const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(il); |
| 1999 | |
| 2000 | auto * k = layer.k_stream[strm]; |
| 2001 | |
| 2002 | // Read type of key |
| 2003 | int32_t k_type_i_ref; |
| 2004 | io.read_to(&k_type_i_ref, sizeof(k_type_i_ref)); |
| 2005 | const int32_t k_type_i = (int32_t) k->type; |
| 2006 | if (k_type_i != k_type_i_ref) { |
| 2007 | LLAMA_LOG_ERROR("%s: mismatched key type (%d != %d, layer %d)\n", __func__, k_type_i, k_type_i_ref, il); |
| 2008 | return false; |
| 2009 | } |
| 2010 | |
| 2011 | // Read row size of key |
| 2012 | uint64_t k_size_row_ref; |
| 2013 | io.read_to(&k_size_row_ref, sizeof(k_size_row_ref)); |
| 2014 | const size_t k_size_row = ggml_row_size(k->type, n_embd_k_gqa); |
| 2015 | if (k_size_row != k_size_row_ref) { |
| 2016 | LLAMA_LOG_ERROR("%s: mismatched key row size (%zu != %zu, layer %d)\n", __func__, k_size_row, (size_t) k_size_row_ref, il); |
| 2017 | return false; |
| 2018 | } |
| 2019 | |
| 2020 | if (cell_count) { |
| 2021 | if (sinfo.is_contiguous()) { |
| 2022 | // Fast path: contiguous cells, single memcpy |
| 2023 | ggml_backend_tensor_set(k, io.read(cell_count * k_size_row), sinfo.head() * k_size_row, cell_count * k_size_row); |
| 2024 | } else { |
| 2025 | // Slow path: scatter to non-contiguous positions |
| 2026 | const void * src = io.read(cell_count * k_size_row); |
| 2027 | for (uint32_t i = 0; i < cell_count; ++i) { |
nothing calls this directly
no test coverage detected