diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index a7d3106c04..6515ec3d14 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -203,9 +203,11 @@ uint32_t llama_hparams::n_embd_r() const { // Corresponds to Mamba's conv_states size const uint32_t n_conv = (ssm_d_conv > 0 ? ssm_d_conv - 1 : 0) * (ssm_d_inner + 2*ssm_n_group*ssm_d_state); - // qwen4exp puts a PLE module on a delta-net layer, so the row holds a second dilated conv - // state; the rows are uniform, so every recurrent layer reserves it - return n_conv + ple_conv_state(); + // qwen4exp's PLE dilated conv history deliberately does not share this row: the Meta backend + // splits cache_r_l by head and cannot view one sub-range of a split axis, so a second history + // packed behind the first is unaddressable under -sm tensor. it lives in cache_ple_r_l instead, + // mirrored, because the whole PLE module is mirrored + return n_conv; } uint32_t llama_hparams::n_embd_s() const { diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index e2990972ef..4e77ad0459 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -51,7 +51,8 @@ llama_memory_recurrent::llama_memory_recurrent( auto it = ctx_map.find(buft); if (it == ctx_map.end()) { ggml_init_params params = { - /*.mem_size =*/ size_t(2u*n_layer*ggml_tensor_overhead()), + // r and s per layer, plus the separate PLE conv row where the model has one + /*.mem_size =*/ size_t((hparams.ple_conv_state() > 0 ? 3u : 2u)*n_layer*ggml_tensor_overhead()), /*.mem_buffer =*/ NULL, /*.no_alloc =*/ true, }; @@ -71,6 +72,7 @@ llama_memory_recurrent::llama_memory_recurrent( r_l.resize(n_layer); s_l.resize(n_layer); + p_l.resize(n_layer); for (int i = 0; i < n_layer; i++) { if (filter && !filter(i)) { @@ -103,6 +105,14 @@ llama_memory_recurrent::llama_memory_recurrent( ggml_format_name(s, "cache_s_l%d", i); r_l[i] = r; s_l[i] = s; + + // qwen4exp's PLE history needs a row of its own so that the Meta backend can mirror it while + // the delta-net conv state next door stays split across devices + if (hparams.ple_conv_state() > 0 && hparams.is_ple(i)) { + ggml_tensor * p = ggml_new_tensor_2d(ctx, type_r, hparams.ple_conv_state(), n_rows); + ggml_format_name(p, "cache_ple_r_l%d", i); + p_l[i] = p; + } } // allocate tensors and initialize the buffers to avoid NaNs in the padding @@ -119,11 +129,13 @@ llama_memory_recurrent::llama_memory_recurrent( { const size_t memory_size_r = size_r_bytes(); const size_t memory_size_s = size_s_bytes(); + const size_t memory_size_p = size_p_bytes(); - LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u seqs %2u rs_seq), R (%s): %7.2f MiB, S (%s): %7.2f MiB\n", __func__, - (float)(memory_size_r + memory_size_s) / (1024.0f * 1024.0f), mem_size, n_layer, n_seq_max, n_rs_seq, + LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u seqs %2u rs_seq), R (%s): %7.2f MiB, S (%s): %7.2f MiB, P (%s): %7.2f MiB\n", __func__, + (float)(memory_size_r + memory_size_s + memory_size_p) / (1024.0f * 1024.0f), mem_size, n_layer, n_seq_max, n_rs_seq, ggml_type_name(type_r), (float)memory_size_r / (1024.0f * 1024.0f), - ggml_type_name(type_s), (float)memory_size_s / (1024.0f * 1024.0f)); + ggml_type_name(type_s), (float)memory_size_s / (1024.0f * 1024.0f), + ggml_type_name(type_r), (float)memory_size_p / (1024.0f * 1024.0f)); } } @@ -740,6 +752,18 @@ size_t llama_memory_recurrent::size_s_bytes() const { return size_s_bytes; } +size_t llama_memory_recurrent::size_p_bytes() const { + size_t size_p_bytes = 0; + + for (const auto & p : p_l) { + if (p != nullptr) { + size_p_bytes += ggml_nbytes(p); + } + } + + return size_p_bytes; +} + void llama_memory_recurrent::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { GGML_UNUSED(flags); @@ -899,6 +923,17 @@ void llama_memory_recurrent::state_write_data(llama_io_write_i & io, const std:: const size_t buf_size = range_size * r_size_row; io.write_tensor(r_l[il], range.first * r_size_row, buf_size); } + + // the PLE conv history is a second recurrent row, so it has to travel with the first + if (p_l[il] != nullptr) { + const uint64_t p_size_row = ggml_row_size(p_l[il]->type, hparams.ple_conv_state()); + io.write(&p_size_row, sizeof(p_size_row)); + + for (const auto & range : cell_ranges) { + const size_t range_size = range.second - range.first; + io.write_tensor(p_l[il], range.first * p_size_row, range_size * p_size_row); + } + } } if (!s_trans) { @@ -1097,6 +1132,20 @@ bool llama_memory_recurrent::state_read_data(llama_io_read_i & io, uint32_t cell // Read and set the keys for the whole cell range io.read_tensor(r_l[il], head * r_size_row, cell_count * r_size_row); } + + if (p_l[il] != nullptr) { + uint64_t p_size_row_ref; + io.read(&p_size_row_ref, sizeof(p_size_row_ref)); + const size_t p_size_row = ggml_row_size(p_l[il]->type, hparams.ple_conv_state()); + if (p_size_row != p_size_row_ref) { + LLAMA_LOG_ERROR("%s: mismatched ple row size (%zu != %zu, layer %d)\n", __func__, p_size_row, (size_t) p_size_row_ref, il); + return false; + } + + if (cell_count) { + io.read_tensor(p_l[il], head * p_size_row, cell_count * p_size_row); + } + } } if (!s_trans) { @@ -1251,6 +1300,10 @@ ggml_tensor * llama_memory_recurrent_context::get_s_l(int32_t il) const { return mem->s_l[il]; } +ggml_tensor * llama_memory_recurrent_context::get_p_l(int32_t il) const { + return mem->p_l[il]; +} + int32_t llama_memory_recurrent_context::s_copy(int i) const { const uint32_t cell_idx = i + mem->head; const int32_t src0 = mem->cells[cell_idx].src0; diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index b13b7b748f..4abb3f5cf5 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -111,6 +111,8 @@ public: // per layer std::vector r_l; std::vector s_l; + // a second conv history that must stay replicated across devices, so it cannot share the r row + std::vector p_l; private: //const llama_model & model; @@ -125,6 +127,7 @@ private: size_t size_r_bytes() const; size_t size_s_bytes() const; + size_t size_p_bytes() const; void state_write_meta(llama_io_write_i & io, const std::vector> & cell_ranges, llama_seq_id seq_id = -1) const; void state_write_data(llama_io_write_i & io, const std::vector> & cell_ranges) const; @@ -170,6 +173,7 @@ public: ggml_tensor * get_r_l(int32_t il) const; ggml_tensor * get_s_l(int32_t il) const; + ggml_tensor * get_p_l(int32_t il) const; int32_t s_copy(int i) const; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index a926d32436..f8b4a30c97 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -395,6 +395,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_ssm_beta ("blk\\.\\d*\\.ssm_beta.weight"); static const std::regex pattern_ssm_beta_alpha ("blk\\.\\d*\\.ssm_ba.weight"); static const std::regex pattern_r_cache ("cache_r_l\\d*"); + static const std::regex pattern_ple_r_cache ("cache_ple_r_l\\d*"); static const std::regex pattern_s_cache ("cache_s_l\\d*"); static const std::regex pattern_ssm_conv1d ("blk\\.\\d*\\.ssm_conv1d.weight"); static const std::regex pattern_ssm_out_weight ("blk\\.\\d*\\.ssm_out.weight"); @@ -497,6 +498,12 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); } + // the PLE table is a model-level lookup and its conv kernel and norm are mirrored, so every + // device computes the whole dilated conv and needs the whole history + if (std::regex_match(tensor_name, pattern_ple_r_cache)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); + } + // standard attention if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_kv_weight)) { return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight", "ssm_out.weight"); diff --git a/src/models/models.h b/src/models/models.h index b22367fb0d..dacca61600 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2341,17 +2341,16 @@ struct llama_model_qwen4exp : public llama_model_base { ggml_tensor * gate, int layer); - // build_rs writes the state tensor in place, so both convolutions share one gather per layer - std::map rs_rows; + // build_rs writes the state tensor in place, so one gather per cache tensor is reused + std::map rs_rows; - // conv history at an explicit offset: delta-net and PLE share the row + // one conv history per cache tensor: delta-net and PLE each have their own ggml_tensor * build_conv_state_at( llm_graph_input_rs * inp, ggml_tensor * conv_states_all, ggml_tensor * x, int64_t state_cols, int64_t channels, - int64_t row_offset, int il); ggml_tensor * build_ple( diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index c438710995..b49cf591a2 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -730,9 +730,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear( // the channels must match how load_arch_tensors sizes wqkv, not ssm_d_inner const int64_t conv_channels = head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads; - // offset 0: delta-net history first, PLE history (if any) after it ggml_tensor * conv_input = build_conv_state_at(inp, conv_states_all, qkv_mixed, - conv_kernel_size - 1, conv_channels, 0, il); + conv_kernel_size - 1, conv_channels, il); ggml_tensor * state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs); @@ -985,36 +984,32 @@ void llm_graph_input_ple::set_input(const llama_ubatch * ubatch) { ggml_backend_tensor_set(rows, idx.data(), 0, idx.size()*ggml_element_size(rows)); } -// Read one conv history from the recurrent row at row_offset and write the new tail back. -// The shared build_conv_state cannot do this: the row holds the delta-net history and the PLE one. +// Read a conv history out of its own recurrent row and write the new tail back. +// The shared build_conv_state cannot do this: qwen4exp has two such rows per layer. ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at( llm_graph_input_rs * inp, ggml_tensor * conv_states_all, ggml_tensor * x, int64_t state_cols, int64_t channels, - int64_t row_offset, int il) { const auto * mctx_cur = inp->mctx; const auto kv_head = mctx_cur->get_head(); const int64_t n_seqs = ubatch.n_seqs; - const int64_t row_total = hparams.n_embd_r(); + const int64_t row_total = conv_states_all->ne[0]; - // the gather needs the whole row, then this convolution takes its slice - auto it = rs_rows.find(il); + // the row is exactly this convolution's state, so the gather is reused as a whole + GGML_ASSERT(state_cols * channels == row_total); + + auto it = rs_rows.find(conv_states_all); if (it == rs_rows.end()) { - it = rs_rows.emplace(il, build_rs(inp, conv_states_all, row_total, n_seqs)).first; + it = rs_rows.emplace(conv_states_all, build_rs(inp, conv_states_all, row_total, n_seqs)).first; } ggml_tensor * rows = it->second; - const size_t esz = ggml_element_size(rows); - - ggml_tensor * state = ggml_cont(ctx0, - ggml_view_2d(ctx0, rows, state_cols * channels, n_seqs, - rows->nb[1], row_offset * esz)); - state = ggml_reshape_3d(ctx0, state, state_cols, channels, n_seqs); + ggml_tensor * state = ggml_reshape_3d(ctx0, rows, state_cols, channels, n_seqs); cb(state, "conv_state_at", il); ggml_tensor * conv_input = ggml_concat(ctx0, state, ggml_transpose(ctx0, x), 0); @@ -1030,7 +1025,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at( ggml_tensor * dst = ggml_view_2d(ctx0, conv_states_all, state_cols * channels, n_seqs, conv_states_all->nb[1], - kv_head * row_size + row_offset * ggml_element_size(conv_states_all)); + kv_head * row_size); ggml_build_forward_expand(gf, ggml_cpy(ctx0, ggml_cont(ctx0, tail), dst)); @@ -1109,10 +1104,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_ple( const int64_t n_seq_tokens = ubatch.n_seq_tokens; // [hist + n_seq_tokens, hc_dim, n_seqs], tokens on ne[0] - ggml_tensor * padded = build_conv_state_at(inp, inp->mctx->get_r_l(il), + ggml_tensor * padded = build_conv_state_at(inp, inp->mctx->get_p_l(il), ggml_reshape_3d(ctx0, normalized, hc_dim, n_seq_tokens, n_seqs), - hist, hc_dim, - hparams.n_embd_r() - hparams.ple_conv_state(), il); + hist, hc_dim, il); ggml_tensor * conv_out = nullptr; for (int64_t k = 0; k < kern; ++k) {