diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 7a0c4bd721..0da038f4e5 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1794,80 +1794,98 @@ void llama_kv_cache::set_input_qsa( ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio) const { - GGML_ASSERT(n_stream == 1 && "TODO: support multiple streams"); GGML_ASSERT(ratio > 0); - const auto & cells = v_cells[0]; - GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer)); const int64_t n_kv = cell_blk->ne[0]; - const int64_t n_blocks = blk_pos->ne[0]/4; + const int64_t n_ns = cell_blk->ne[1]; // streams in this ubatch + const int64_t n_blocks = blk_pos->ne[0]/(4*n_ns); const int64_t n_tokens = ubatch->n_tokens; const int64_t r = ratio; + GGML_ASSERT(n_tokens % n_ns == 0); + const int64_t n_tps = n_tokens/n_ns; // tokens per stream + int32_t * dst_cell_blk = (int32_t *) cell_blk->data; int32_t * dst_blk_cells = (int32_t *) blk_cells->data; int32_t * dst_blk_pos = (int32_t *) blk_pos->data; float * dst_bias = (float *) bias->data; // block b covers [b*ratio, (b+1)*ratio), so its first token is at b*ratio. All three - // mrope sections carry that position: exact for text, approximate for image tokens. - for (int64_t s = 0; s < 4; ++s) { - for (int64_t b = 0; b < n_blocks; ++b) { - dst_blk_pos[s*n_blocks + b] = (int32_t) (b*r); + // mrope sections carry it: exact for text, approximate for images. Positions repeat per stream. + for (int64_t sec = 0; sec < 4; ++sec) { + for (int64_t s = 0; s < n_ns; ++s) { + for (int64_t b = 0; b < n_blocks; ++b) { + dst_blk_pos[sec*(n_blocks*n_ns) + s*n_blocks + b] = (int32_t) (b*r); + } } } - // an incompletely populated block cannot be pooled. Those cells are the tail, which - // the bias below forces in anyway, so they point at block 0 just to keep the gather - // in range. -1 marks a cell with no usable block. - std::vector blk_of(n_kv, -1); - std::vector filled(n_blocks, 0); + // One pass per stream: cell j is a different token in each, so no mapping is shared. + // n_ns == 1 is the single-stream behaviour this replaced. + std::vector blk_of(n_kv); + std::vector filled(n_blocks); - std::fill(dst_blk_cells, dst_blk_cells + r*n_blocks, 0); + for (int64_t s = 0; s < n_ns; ++s) { + // ubatch index s*n_tps belongs to this stream; ask which cells array it uses + const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0]; + const auto & cells = v_cells[seq_to_stream[seq_of_stream]]; - for (int64_t j = 0; j < n_kv; ++j) { - if (cells.is_empty(j)) { - continue; - } + int32_t * cur_cell_blk = dst_cell_blk + s*n_kv; + int32_t * cur_blk_cells = dst_blk_cells + s*(r*n_blocks); - const llama_pos p = cells.pos_get(j); - const int64_t b = p/r; - - if (b >= n_blocks) { - continue; - } - - blk_of[j] = (int32_t) b; - dst_blk_cells[b*r + (p%r)] = (int32_t) j; - filled[b]++; - } - - for (int64_t j = 0; j < n_kv; ++j) { - if (blk_of[j] >= 0 && filled[blk_of[j]] < r) { - blk_of[j] = -1; - } - dst_cell_blk[j] = blk_of[j] < 0 ? 0 : blk_of[j]; - } - - for (int64_t i = 0; i < n_tokens; ++i) { - const llama_seq_id seq_id = ubatch->seq_id[i][0]; - const llama_pos q = ubatch->pos[i]; - - // from here on we are inside an incomplete block, always attended to, which is what - // lands the selection on block boundaries as the reference does - const llama_pos tail_start = (q + 1)/r*r; + // an incomplete block cannot be pooled: those tail cells are forced in by the bias + // below, so block 0 only keeps the gather in range. -1 = no usable block. + std::fill(blk_of.begin(), blk_of.end(), -1); + std::fill(filled.begin(), filled.end(), 0); + std::fill(cur_blk_cells, cur_blk_cells + r*n_blocks, 0); for (int64_t j = 0; j < n_kv; ++j) { - float v = -INFINITY; - - if (!cells.is_empty(j) && cells.seq_has(j, seq_id) && cells.pos_get(j) <= q) { - // finite, so it can never meet a -inf and produce a nan - v = cells.pos_get(j) >= tail_start ? 1e9f : (blk_of[j] < 0 ? -INFINITY : 0.0f); + if (cells.is_empty(j)) { + continue; } - dst_bias[i*n_kv + j] = v; + const llama_pos p = cells.pos_get(j); + const int64_t b = p/r; + + if (b >= n_blocks) { + continue; + } + + blk_of[j] = (int32_t) b; + cur_blk_cells[b*r + (p%r)] = (int32_t) j; + filled[b]++; + } + + for (int64_t j = 0; j < n_kv; ++j) { + if (blk_of[j] >= 0 && filled[blk_of[j]] < r) { + blk_of[j] = -1; + } + cur_cell_blk[j] = blk_of[j] < 0 ? 0 : blk_of[j]; + } + + for (int64_t ii = 0; ii < n_tps; ++ii) { + const int64_t i = s*n_tps + ii; + const llama_seq_id seq_id = ubatch->seq_id[i][0]; + const llama_pos q = ubatch->pos[i]; + + // the rest is an incomplete block, always attended to, which is what lands the + // selection on block boundaries like the reference + const llama_pos tail_start = (q + 1)/r*r; + + float * cur_bias = dst_bias + i*n_kv; + + for (int64_t j = 0; j < n_kv; ++j) { + float v = -INFINITY; + + if (!cells.is_empty(j) && cells.seq_has(j, seq_id) && cells.pos_get(j) <= q) { + // finite, so it can never meet a -inf and produce a nan + v = cells.pos_get(j) >= tail_start ? 1e9f : (blk_of[j] < 0 ? -INFINITY : 0.0f); + } + + cur_bias[j] = v; + } } } } @@ -2670,6 +2688,11 @@ uint32_t llama_kv_cache_context::get_n_kv() const { return n_kv; } +uint32_t llama_kv_cache_context::get_n_stream() const { + // streams in the current slot info, matching get_k/get_v's `ns`. 1 if unified. + return sinfos[i_cur].s1 - sinfos[i_cur].s0 + 1; +} + ggml_type llama_kv_cache_context::type_k() const { return kv->type_k(); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 37616e7b35..29391724ea 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -375,6 +375,9 @@ public: uint32_t get_n_kv() const; + // streams in the current slot info; 1 for a unified cache + uint32_t get_n_stream() const; + ggml_type type_k() const; ggml_type type_v() const; diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 686bae2d05..275b1a3922 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -9,6 +9,7 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, true); + ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); ml.get_key(LLM_KV_SSM_INNER_SIZE, hparams.ssm_d_inner); ml.get_key(LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state); @@ -22,6 +23,7 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { GGML_ASSERT(hparams.hc_low_rank > 0 && "qwen4exp needs a hyper-connection low rank"); hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd; + ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head); ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); @@ -136,6 +138,7 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0); layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0); + const int64_t idx_dim = hparams.indexer_head_size; layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", il), { n_embd, hparams.indexer_n_head * idx_dim }, 0); layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", il), { n_embd, idx_dim }, 0); @@ -392,11 +395,12 @@ public: mctx->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio); } + // Per-stream: a cell index means a different token in each stream. n_stream 1 = the old shapes. ggml_tensor * k_idxs = nullptr; // I32 [n_tokens] - ggml_tensor * cell_blk = nullptr; // I32 [n_kv] - ggml_tensor * blk_cells = nullptr; // I32 [ratio*n_blocks] - ggml_tensor * blk_pos = nullptr; // I32 [4*n_blocks] - ggml_tensor * bias = nullptr; // F32 [n_kv, n_tokens] + ggml_tensor * cell_blk = nullptr; // I32 [n_kv, n_stream] + ggml_tensor * blk_cells = nullptr; // I32 [ratio*n_blocks, n_stream] + ggml_tensor * blk_pos = nullptr; // I32 [4*n_blocks*n_stream] + ggml_tensor * bias = nullptr; // F32 [n_kv, n_tokens/n_stream, n_stream] const llama_kv_cache_context * mctx; const uint32_t ratio; @@ -417,13 +421,18 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( const int64_t n_blocks = (n_kv + r - 1)/r; + // n_tps: tokens divide evenly across streams, as build_attn_mask_top_k and the KQ mask assume. + const int64_t n_stream = mctx_idx->get_n_stream(); + GGML_ASSERT(n_tokens % n_stream == 0); + const int64_t n_tps = n_tokens/n_stream; + auto qsa = std::make_unique(mctx_idx, (uint32_t) r); qsa->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); - qsa->cell_blk = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_kv); - qsa->blk_cells = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, r*n_blocks); - qsa->blk_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks); - qsa->bias = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv, n_tokens); + qsa->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_stream); + qsa->blk_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, r*n_blocks, n_stream); + qsa->blk_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks*n_stream); + qsa->bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_kv, n_tps, n_stream); ggml_set_input(qsa->cell_blk); ggml_set_input(qsa->blk_cells); @@ -440,30 +449,33 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, k_raw, inp->k_idxs, il)); - // one key head, so the per-cell rows are contiguous and this view is dense + // one key head, so rows are contiguous. get_k gives [idx_dim, n_head_kv, n_kv, n_stream]. ggml_tensor * k_all = mctx_idx->get_k(ctx0, il); - k_all = ggml_view_2d(ctx0, k_all, idx_dim, n_kv, k_all->nb[2], 0); + k_all = ggml_view_3d(ctx0, k_all, idx_dim, n_kv, n_stream, k_all->nb[2], k_all->nb[3], 0); + // gathers per stream: blk_cells row s indexes stream s's own cells ggml_tensor * members = ggml_get_rows(ctx0, k_all, inp->blk_cells); - members = ggml_reshape_3d(ctx0, members, idx_dim, r, n_blocks); + members = ggml_reshape_4d(ctx0, members, idx_dim, r, n_blocks, n_stream); // mean over the block's members; compress_ratio is small, so summing slices beats // transposing to reach ggml_sum_rows ggml_tensor * pooled = nullptr; for (int64_t i = 0; i < r; ++i) { ggml_tensor * slice = ggml_cont(ctx0, - ggml_view_2d(ctx0, members, idx_dim, n_blocks, members->nb[2], i*members->nb[1])); + ggml_view_3d(ctx0, members, idx_dim, n_blocks, n_stream, + members->nb[2], members->nb[3], i*members->nb[1])); pooled = pooled ? ggml_add(ctx0, pooled, slice) : slice; } pooled = ggml_scale(ctx0, pooled, 1.0f/(float) r); cb(pooled, "indexer_k_pooled", il); - pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_blocks); + // rope wants [n_dims, n_head, n_tokens]: lay every stream's blocks flat, split after. + pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_blocks*n_stream); pooled = build_norm(pooled, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il); pooled = ggml_rope_multi(ctx0, pooled, inp->blk_pos, nullptr, n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - pooled = ggml_reshape_2d(ctx0, pooled, idx_dim, n_blocks); + pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, n_blocks, n_stream); cb(pooled, "indexer_k", il); ggml_tensor * q = build_lora_mm(model.layers[il].index_q_proj, cur); @@ -475,15 +487,15 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( cb(q, "indexer_q", il); // Each head's dot product is rectified before summing, as in DeepSeek's lightning - // indexer but with no learned per-head weight. The reference's constant divisor - // cannot reorder anything, so it is left out. + // no per-head weight; the constant divisor cannot reorder. mul_mat matches ne[2], so + // stream s's queries only meet stream s's blocks. ggml_tensor * score = ggml_mul_mat(ctx0, pooled, - ggml_reshape_2d(ctx0, ggml_cont(ctx0, q), idx_dim, n_idx_h*n_tokens)); - score = ggml_reshape_3d(ctx0, score, n_blocks, n_idx_h, n_tokens); + ggml_reshape_3d(ctx0, ggml_cont(ctx0, q), idx_dim, n_idx_h*n_tps, n_stream)); + score = ggml_reshape_4d(ctx0, score, n_blocks, n_idx_h, n_tps, n_stream); score = ggml_relu(ctx0, score); score = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)); score = ggml_sum_rows(ctx0, score); - score = ggml_reshape_2d(ctx0, score, n_blocks, n_tokens); + score = ggml_reshape_3d(ctx0, score, n_blocks, n_tps, n_stream); cb(score, "indexer_score", il); // Give every token of a block its block's score rather than expanding the block @@ -491,8 +503,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( // budget is a whole number of blocks and members tie, so the cut still lands on // a block boundary. get_rows gathers rows, so scores are transposed first. ggml_tensor * expanded = ggml_get_rows(ctx0, - ggml_cont(ctx0, ggml_transpose(ctx0, score)), inp->cell_blk); - expanded = ggml_cont(ctx0, ggml_transpose(ctx0, expanded)); + ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), inp->cell_blk); + expanded = ggml_cont(ctx0, ggml_permute(ctx0, expanded, 1, 0, 2, 3)); expanded = ggml_add(ctx0, expanded, inp->bias); cb(expanded, "indexer_score_tokens", il); @@ -501,6 +513,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( const int64_t width = std::min(n_kv, (int64_t) hparams.indexer_top_k + r - 1); ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, width)); + + // build_attn_mask_top_k reads [n_top_k, n_batch, 1, n_stream], matching the KQ mask. + top_k = ggml_reshape_4d(ctx0, top_k, width, n_tps, 1, n_stream); cb(top_k, "indexer_top_k", il); return top_k; @@ -701,6 +716,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear( k_conv = ggml_l2_norm(ctx0, k_conv, eps_norm); + // repeat to match shapes when head keys != value keys; unneeded with the fused GDN if (num_k_heads != num_v_heads && (!cparams.fused_gdn_ar || !cparams.fused_gdn_ch)) { GGML_ASSERT(num_v_heads % num_k_heads == 0);