graph : reorder get_rows for embeddings

[ggerganov]

Overview

cont #29622
fix #30033

With mixed batch support, we have 3 possible graph branches that produce the input embeddings:

• from tokens referencing the token embeddings weight tensor
• raw embeddings
• mixed: both tokens + raw

The way we currently implement it leads to interleaving graph ops (such as PAD, MUL, SCALE) with the token embeddings extraction (GET_ROWS). This is not desired because it creates unnecessary splits or leads to running some of the ops on the CPU.

To fix that, we have to rework the logic such that it orders all GET_ROWS to be the first nodes in the graph. We can achieve this with the ggml_build_forward_order.

The goal is SPLIT 0 to always contain just the GET_ROWS ops and nothing else:

Before:

After:

TODO:

• Refactor to make the implementation cleaner
• Fix Gemma4 PLE
• Fix a reshaped CPU tensor reaching the Meta backend (same as model : add K2 Horizon dense and MoVA support #29535 (comment))

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: NO

[deadprogram]

Hello @ggerganov I did local benchmarks with this PR and here are my results:

RTX 4070 Laptop GPU, CUDA, Q4_K_M models, llama-bench -ngl 99 -p 512 -n 128 -r 5, with -t as the thread count below. Values are tokens per second.

model	threads	test	b11399	master	PR 30160
Qwen3-VL-2B	8	tg128	183.9	166.2	182.5
Qwen3-VL-2B	32	tg128	183.8	153.9	182.4
Qwen3-VL-2B	8	pp512	11912	10210	11571
Qwen3-VL-2B	32	pp512	12474	10517	11372
Gemma 4 E2B	8	tg128	131.4	122.3	121.7
Gemma 4 E2B	32	tg128	130.8	114.9	115.1
Gemma 4 E2B	8	pp512	7618	8057	7842
Gemma 4 E2B	32	pp512	7699	7731	8078

master is c811cb8, the base of PR 30160 at 4e67ebb. b11399 is the release before PR 29622. pp512 varies about 3 percent from run to run.

So in short it does handle the performance slowdown thank you!

Of course, as already noted in the current TODO Gemma 4 still suffers from a performance hit, but perhaps that if for another PR?

[moriz1]

here are my results using the same prompt and text file as in my initial bug report:

PP: 582.23 t/s
TG: 32.75 t/s

vs #29622 :

PP: ~365 t/s
TG: ~24t /s

looks like it's good.

[ggerganov]

@deadprogram The Gemma4 regressions should be fixed with the latest version of this branch. Please run the tests again. Thanks.

[ggerganov]

@ngxson This change should improve the input embedding logic, though it's still not perfect and not very easy to follow. Will be thinking how to make it even better. Also, I expect that with LoRAs enabled, we'll have similar issues - I've added TODO to add LoRA tests and figure out how to prevent unnecessary splits.

[deadprogram]

Here are my benchmarks, however it appears that Gemma4 with embeddibng only batches now crashes.

First the numbers:

RTX 4070 Laptop GPU, CUDA, Q4_K_M models, llama-bench -ngl 99 -p 512 -n 128 -r 5, with -t as the thread count below. Values are tokens per second.

model	threads	test	b11399	master	PR 30160
Qwen3-VL-2B	8	tg128	183.8	166.7	182.5
Qwen3-VL-2B	32	tg128	183.8	154.4	182.4
Qwen3-VL-2B	8	pp512	12454	10402	12079
Qwen3-VL-2B	32	pp512	12463	10574	11958
Gemma 4 E2B	8	tg128	131.4	123.3	130.9
Gemma 4 E2B	32	tg128	131.4	115.9	130.9
Gemma 4 E2B	8	pp512	7714	8153	8081
Gemma 4 E2B	32	pp512	7545	8199	8099

My coding tools gave this reason for the crash:

• With 1 row: ggml-cpu/ops.cpp:4987: fatal error, because ggml_scale gets the Q5_K tensor without the cast to F32.
• With 8 rows: ggml.c:3758: GGML_ASSERT(ggml_nelements(a) == ne0ne1ne2), because one padding row is reshaped to n_tokens rows.

[ggerganov]

@deadprogram Thanks for the quick feedback - should be fixed in the new commit.

[deadprogram]

@ggerganov thanks for working on this! Crash fixed, and here are the new numbers:

RTX 4070 Laptop GPU, CUDA, Q4_K_M models, llama-bench -ngl 99 -p 512 -n 128 -r 5, with -t as the thread count below. Values are tokens per second.

model	threads	test	b11399	master	PR 30160
Qwen3-VL-2B	8	tg128	183.8	166.8	182.3
Qwen3-VL-2B	32	tg128	181.7	155.6	182.3
Qwen3-VL-2B	8	pp512	12460	10478	11620
Qwen3-VL-2B	32	pp512	12254	10293	12007
Gemma 4 E2B	8	tg128	131.4	122.6	130.9
Gemma 4 E2B	32	tg128	131.4	115.8	130.7
Gemma 4 E2B	8	pp512	7981	8080	8070
Gemma 4 E2B	32	pp512	7960	8193	8125

master is e60eff9, the base of PR 30160 at bd40fd1. b11399 is the release before PR 29622. pp512 varies about 3 percent from run to run.

[ngxson]

looks good, I'm running benchmark on my side, will report results a bit later

[ngxson]

Results:

CUDA box (RTX 5060 Ti, i7-12700KF)

model	threads	test	b11399	master	PR	PR vs master	PR vs b11399
Qwen3-VL-2B	8	pp512	13956	11390	13136	+15.3%	−5.9%
Qwen3-VL-2B	20	pp512	14043	11460	13255	+15.7%	−5.6%
Qwen3-VL-2B	8	tg128	287.3	266.7	281.5	+5.5%	−2.0%
Qwen3-VL-2B	20	tg128	287.2	267.0	281.4	+5.4%	−2.0%
Gemma 4 E4B	8	pp512	5285	5302	5276	−0.5%	−0.2%
Gemma 4 E4B	20	pp512	5310	5326	5303	−0.4%	−0.1%
Gemma 4 E4B	8	tg128	106.3	101.0	105.5	+4.5%	−0.8%
Gemma 4 E4B	20	tg128	106.3	100.1	105.5	+5.4%	−0.8%

Mac (M5 Max, Metal, 8 threads)

model	test	b11399	master	PR
Qwen3-VL-2B	pp512	11354	10930	11198
Qwen3-VL-2B	tg128	290.3	274.3	286.9
Gemma 4 E4B	pp512	4367	3848	4334
Gemma 4 E4B	tg128	107.4	104.7	106.7