Repository navigation
llama: support both embd + raw tokens in batch - #29622
Conversation
|
I tested this on an H200 with the CI build flag The server vision tests ( From prints in This seems to come from the |
I already started working on this, so I would suggest discussing this carefully to avoid duplicated works. The main complication is that for models that doesn't support mixed batch, simply feeding everything into cgraph will make How do you plan to resolve that? |
|
I'm having a look on the issue with graph realloc, no the problem was not about I'm experimenting with different approaches, will push when I find one that is optimal |
|
51e4ca9 should fix it: no more realloc after switching between mixed batch and normal batch trade-off: graph now uses more memory, the trick I used is to have |
|
I re-ran the same setup on the newest changes. The vision tests still abort (8 of 48, master 0). The failing path is the plain per-chunk mtmd path, so not the mixed batch path. What I think is the cause now: On the complication, perhaps it can be handled in the batch allocator. When the batch is mixed and the arch doesn't support it, split it at every change of entry kind, so every ubatch is token-only or embd-only, like the per-chunk path produces today. I could look into it if that helps. |
|
OK it's quite a bummer, turns out My fix was to simply have a different set of input nodes for mixed embd+text case, but trade-off is that memory is increased CC @ggerganov if you have any other ideas |
|
Re-ran the same setup (da0fdfb), now it works. For the models that don't support mixed batches, I have a working version where the allocator splits the batch at each change of entry kind, so every ubatch is only either tok or embd instead of returning an error, without need to change the code in |
I think this should work for now. If I understand correctly, the overhead is relatively small - just ~1.5x more input buffers? |
| llama_seq_id * seq_id_unq; // [n_seqs_unq] | s | seq_id | ||
| int32_t * seq_idx; // [LLAMA_MAX_SEQ] | - | seq_idx | ||
| int8_t * output; // [n_tokens] | i | - | ||
| int8_t * is_embd; // [n_tokens] | i | - (mixed ubatch only) |
There was a problem hiding this comment.
I think calling this member type would make it more clear:
| int8_t * is_embd; // [n_tokens] | i | - (mixed ubatch only) | |
| int8_t * type; // [n_tokens] | i | - (mixed ubatch only, 0 - token, 1 - embd) |
But feel free to keep it like this if you prefer.
| const int64_t n_tokens = ubatch->n_tokens; | ||
|
|
||
| if (ubatch->token && n_pos_per_embd == 4) { | ||
| if (ubatch->token && !ubatch->is_mixed() && n_pos_per_embd == 4) { |
There was a problem hiding this comment.
How do we handle the case of ubatch->is_mixed() && n_pos_per_embd == 4? As it is, I think it would go in the else branch which seems incorrect?
There was a problem hiding this comment.
yeah right, that was a bit messy: there are 2 places handling the same thing (one in batch allocator and one in graph input). I consolidated this logic into batch allocator: e7e6738
| // build flat pos array | ||
| // token batch: pos[i] = tokens[i].pos[0] | ||
| // embedding batch: pos[j*n_tok + i] = tokens[i].pos[j] (section-major) | ||
| // mixed batch: same as embedding batch, token pos is broadcast |
There was a problem hiding this comment.
It is not completely clear what we mean by "token pos is broadcast" here - try to clarify
|
@sihanyu03 sorry for the delay. I'm working on the mtmd side and it's attached to a model (unfortunately I cannot share more info here, but I'm currently on a rush). So it would be better if I implement the mtmd changes on my side. Nevertheless, my implementation will only cover the case for non-causal models (as it is currently impossible to support this case on master branch). Please feel free to push your PR for causal one. (Edit: ok so fortunately we can already demo this on Clef, see #29969 - still, I need the non-causal case for an upcoming model)
Yes this sounds good. Please proceed with it. |
I'll look into it! |
Upstream ggml-org#29622 adds a mixed token/embd branch to every input embedding graph through ggml_build_forward_select(). Its nodes are not flagged for compute, but the backend translated them anyway, and the DUP in that branch was unsupported, so the scheduler split the graph and passed the embeddings across the split with a fixed token count. The first single-token decode then failed (test-thread-safety on CPU and GPU). Build the OV model from the compute nodes only, and translate a same-type contiguous DUP like CONT so the graph stays on one backend.
ggml-org#29622 also moves the per-token embedding scale (gemma3, gemma3n, gemma4) into a new [1, n_tokens] input. Give it a dynamic token dim and pad it per chunk on the static (NPU) path.
|
I'm getting a performance regression from this PR and every subsequent commit. My setup:
Model: Settings (llama-server via llama-swap):
test prompt: I gave it a small text file and asked it to summarize it. PR #29612 (one previous to this): pp: ~420t/s tg: ~32t/s I've observed the same performance regression up to the current PR as of writing (PR #29971). |
Upstream ggml-org#29622 adds a mixed token/embd branch to every input embedding graph through ggml_build_forward_select(). Its nodes are not flagged for compute, but the backend translated them anyway, and the DUP in that branch was unsupported, so the scheduler split the graph and passed the embeddings across the split with a fixed token count. The first single-token decode then failed (test-thread-safety on CPU and GPU). Build the OV model from the compute nodes only, and translate a same-type contiguous DUP like CONT so the graph stays on one backend.
ggml-org#29622 also moves the per-token embedding scale (gemma3, gemma3n, gemma4) into a new [1, n_tokens] input. Give it a dynamic token dim and pad it per chunk on the static (NPU) path.
* ggml-openvino: skip unselected graph branches and support DUP Upstream #29622 adds a mixed token/embd branch to every input embedding graph through ggml_build_forward_select(). Its nodes are not flagged for compute, but the backend translated them anyway, and the DUP in that branch was unsupported, so the scheduler split the graph and passed the embeddings across the split with a fixed token count. The first single-token decode then failed (test-thread-safety on CPU and GPU). Build the OV model from the compute nodes only, and translate a same-type contiguous DUP like CONT so the graph stays on one backend. * ggml-openvino: make inp_scale_rows token dim dynamic #29622 also moves the per-token embedding scale (gemma3, gemma3n, gemma4) into a new [1, n_tokens] input. Give it a dynamic token dim and pad it per chunk on the static (NPU) path. * ggml-openvino: skip GPU MUL_MAT op tests with unbound Q4_1/Q4_K weights Op tests build Q4_1/Q4_K weights as u4 with an f16 zero point. The GPU plugin fails to compile that form for some row counts with "clFinish, error code: -5 CL_OUT_OF_RESOURCES", which aborts test-backend-ops on the MUL_MAT cases added in #29869 (e.g. m=1000, n=2, k=1024). Model weights use a u4 zero point and are not affected. Report these cases as unsupported on GPU until the plugin is fixed. Op tests check support before allocating, so the check matches unbound weights only; model loading probes with a dummy buffer and keeps its weights on the GPU. * ggml-openvino: create FILL in the output type translate_fill always built an f32 constant, so an f16 FILL produced f32 data and the copy back overran the f16 output buffer. Use the output type for the constant. * ggml-openvino: reject CONCAT with a quantized type Quantized inputs are dequantized when translated, so the backend cannot write a quantized CONCAT output. Report it as unsupported, as for CPY to a quantized type. * ggml-openvino: handle the single recurrent state gather of build_rs #29856 changed build_rs to gather all recurrent states with one GET_ROWS on the s_copy leaf and take the ubatch and extra states as views of it. The stateful path matched only the previous form, a GET_ROWS per view of s_copy, so Qwen3.5 failed with stateful execution on CPU and GPU ("is_axis_valid(axis, r)" in a Concat). For a single-slot cache, treat the GET_ROWS on the s_copy leaf as the active-state gather, keep the rank-4 layout of reshapes that read a view of it, and map the copy of the empty extra-state view to the single-slot remainder writeback. Do not warn about the dynamic dim of empty views. * openvino: align eltwise operand ranks to work around a GPU-plugin defect * openvino: match the MoE fusion on the rank-3 stateful graph * ggml-openvino: do not unsqueeze an RMS norm output in AlignEltwiseOperandRanks The pass unsqueezes the lower-rank operand of an Add/Multiply/Subtract whose operand ranks differ. In gemma-3 the lower-rank operand of the post-attention residual add is the norm output, and unsqueezing it makes the GPU plugin compute the layer wrongly: gemma-3 returns empty answers on GPU with stateful execution. Skip the rewrite when the lower-rank operand is an RMS norm output. * docs : update OpenVINO validated models --------- Co-authored-by: Mustafa Cavus <mustafa.cavus@intel.com>
[qvac-b11259: scale tokens after build_inp_embd, this base has no tok_scale parameter (ggml-org#29622)] Assisted-by: Claude Code (cherry picked from commit 4fbc76d)
|
I discovered this PR while working on benchmarks for our yzma project. After a bunch of automated runs, here is an edited version of the result. Turns out that it slows down decode for text-only batches. This was also reported in #30018 (Gemma 4, Vulkan) and #30033 (Intel), but not sure those are being looked at or fully make sense. So I measured about 10 percent on CUDA. Eventually I found the cause, and a small patch that brings back the old speed. Cause:
The cost grows with the thread count. Decode tokens per second on an RTX 4070 Laptop, CUDA, Q4_K_M,
A side effect on CUDA: with the new layout the allocator places Patch: graph reuse already compares - const bool has_mixed = llm_arch_supports_mixed_batch(arch) && cparams.ctx_type == LLAMA_CONTEXT_TYPE_DEFAULT;
+ const bool has_mixed = ubatch.is_mixed() && llm_arch_supports_mixed_batch(arch) && cparams.ctx_type == LLAMA_CONTEXT_TYPE_DEFAULT; inp->scale_tok = tok_scale*(scale_tok_only ? hparams.f_embedding_scale : 1.0f);
- if (inp->scale_tok != 1.0f) {
+ if (inp->scale_tok != 1.0f && !ubatch.is_mixed()) {
+ // one kind of row only, so a constant scale is enough
+ if (ubatch.token) {
+ cur = ggml_scale(ctx0, cur, inp->scale_tok);
+ }
+ } else if (inp->scale_tok != 1.0f) {With the patch on b11400, 8 threads:
The greedy output text is identical to b11399 on both models, and no GLU fusion is blocked. I have not tested a mixed batch with the patch. The comment above the scale says it comes after the select "so that the graph is the same for any batch contents". If one graph for all batches is a goal, the patch works against it, and another fix may be better. I can open a PR or you all can decide on what you think best to do about it. Thank you for looking! UPDATE: actually I think that patch just does not keep the correct design on what the idea in this PR is trying to accomplish. Currently working on a better idea. |
* ggml-openvino: skip unselected graph branches and support DUP Upstream ggml-org#29622 adds a mixed token/embd branch to every input embedding graph through ggml_build_forward_select(). Its nodes are not flagged for compute, but the backend translated them anyway, and the DUP in that branch was unsupported, so the scheduler split the graph and passed the embeddings across the split with a fixed token count. The first single-token decode then failed (test-thread-safety on CPU and GPU). Build the OV model from the compute nodes only, and translate a same-type contiguous DUP like CONT so the graph stays on one backend. * ggml-openvino: make inp_scale_rows token dim dynamic ggml-org#29622 also moves the per-token embedding scale (gemma3, gemma3n, gemma4) into a new [1, n_tokens] input. Give it a dynamic token dim and pad it per chunk on the static (NPU) path. * ggml-openvino: skip GPU MUL_MAT op tests with unbound Q4_1/Q4_K weights Op tests build Q4_1/Q4_K weights as u4 with an f16 zero point. The GPU plugin fails to compile that form for some row counts with "clFinish, error code: -5 CL_OUT_OF_RESOURCES", which aborts test-backend-ops on the MUL_MAT cases added in ggml-org#29869 (e.g. m=1000, n=2, k=1024). Model weights use a u4 zero point and are not affected. Report these cases as unsupported on GPU until the plugin is fixed. Op tests check support before allocating, so the check matches unbound weights only; model loading probes with a dummy buffer and keeps its weights on the GPU. * ggml-openvino: create FILL in the output type translate_fill always built an f32 constant, so an f16 FILL produced f32 data and the copy back overran the f16 output buffer. Use the output type for the constant. * ggml-openvino: reject CONCAT with a quantized type Quantized inputs are dequantized when translated, so the backend cannot write a quantized CONCAT output. Report it as unsupported, as for CPY to a quantized type. * ggml-openvino: handle the single recurrent state gather of build_rs ggml-org#29856 changed build_rs to gather all recurrent states with one GET_ROWS on the s_copy leaf and take the ubatch and extra states as views of it. The stateful path matched only the previous form, a GET_ROWS per view of s_copy, so Qwen3.5 failed with stateful execution on CPU and GPU ("is_axis_valid(axis, r)" in a Concat). For a single-slot cache, treat the GET_ROWS on the s_copy leaf as the active-state gather, keep the rank-4 layout of reshapes that read a view of it, and map the copy of the empty extra-state view to the single-slot remainder writeback. Do not warn about the dynamic dim of empty views. * openvino: align eltwise operand ranks to work around a GPU-plugin defect * openvino: match the MoE fusion on the rank-3 stateful graph * ggml-openvino: do not unsqueeze an RMS norm output in AlignEltwiseOperandRanks The pass unsqueezes the lower-rank operand of an Add/Multiply/Subtract whose operand ranks differ. In gemma-3 the lower-rank operand of the post-attention residual add is the norm output, and unsqueezing it makes the GPU plugin compute the layer wrongly: gemma-3 returns empty answers on GPU with stateful execution. Skip the rewrite when the lower-rank operand is an RMS norm output. * docs : update OpenVINO validated models --------- Co-authored-by: Mustafa Cavus <mustafa.cavus@intel.com> (cherry picked from commit b9a5a00)
(cherry picked from commit 0bb496d) Omitted the decision_order lines: that field comes from 99b9548 (clef decision model, ggml-org#29831) which is conflict-skipped in this fork. It will land when that commit is synced. Needed by gemma-embedding2 (4fbc76d) which calls build_inp_embd() with a token scale.
|
I have proposed this https://github.com/ggml-org/llama.cpp/pull/30151 as a better way to address the actualy issue without altering the intent of this PR. PTAL! |
|
Replaced by #30152 from my personal account based on most current project requirements. Same changes. |
…-10-10) Conflicts: - ggml-backend.cpp: upstream moved input copies to ggml_backend_sched_copy_input with host weights last (ggml-org#29943); the GGML_SCHED_BATCH_INPUTS fast path now runs in the first loop before falling back to it. - ggml-cuda/fattn-common.cuh: launch_fattn takes upstream async_kv_preload, then our parallel_blocks_fixed (GQA-6 tile launcher updated). - ggml-cuda/fattn-mma-f16.cuh: upstream swizzle refactor (ggml-org#29612); its helpers take ncols2 like the rest of the fork's config getters. - ggml-cuda/gated_delta_net.cu: keep the fork's kernel and launch geometry (its columns_per_block template argument is not upstream's cols_per_warp, ggml-org#30087). - ggml-cuda/mmq.cu: upstream per-tile src1 padding (ggml-org#29953); ggml_cuda_mul_mat_q_src1_nbytes and the MMQ input cache use the worst-case padding so producers and projections with other tiles can share the buffer. - ggml-cuda/mmq-vec-dot.cuh: fork vec_dot kernels pass GGML_PREC_Q8 to the config helpers (ggml-org#30168). - ggml-cuda/mmvf.cu/.cuh: keep both includes and the HIP declarations, upstream warp_size argument. - ggml-cuda/rope.cu: upstream grid-stride rms_norm+mul+rope (ggml-org#28175) with the fork's M-RoPE theta (nchannels instead of gridDim.y) and rope_store_cast. - ggml-cuda/top-k.cu: upstream shape-based selection (ggml-org#28713), RDNA4 tournament kept for few long rows. - src/llama-batch.*: mixed token/embd batches (ggml-org#29622) copy embd rows per entry; other batches keep pointing at the ext rows. - src/llama-context.cpp: expert copy callback also on the second MTP scheduler. - src/llama-graph.cpp: build_rs keeps the GET_ROWS of the ubatch states (GDN fusions, deferred state) and gathers the extra states from the tail like upstream's custom getter path (ggml-org#29856).
- llama-graph: build the mixed token/embd input branch (ggml-org#29622) only for mixed ubatches. The always-present inactive branch copied 3 inputs to the device per graph and its nodes took compute buffer space at the start of the graph, which shifted the first layer's layout so that the layer-0 residual_rms fusion no longer passed its overlap checks (1 kernel became 5). can_reuse rejects a graph when the mixedness of the ubatch changes. LLAMA_EMBD_MIXED_LAZY=0 restores the static topology. - norm.cu/rope.cu: the grid-stride loops of ggml-org#28175 end every iteration with __syncthreads; only sync when the block runs another iteration. - rope.cu: rms_norm_mul_rope uses a single-row kernel unless the grid is clamped (the loop version costs ~0.4 us per launch on RDNA4). MoE decode, 613 kernels per token again (620 after the sync); tg128 vs r55 -0.36% -> -0.12..-0.14%. Outputs identical to r55 (server compare ALL_IDENTICAL, production MTP bench same drafts).
6 participants
Overview
Extend
llama_batch_exthandling to allow adding both embd and text tokens into a batchCan be useful for models like paligemma where prompt is processed as non-causal
See #29622 (comment) for an explanation on how this works
Note that some models are still not supporting this,
llm_arch_supports_mixed_batchwill gate this internally and return an error "batch in valid" in this caseTODO: handle this in mtmd
Requirements