Repository: ggml-org/llama.cpp
Pull request: llama : fix pooling assertion crash in chunked GDN detection path

## Summary

Fix `ggml_mul_mat` assertion failure (`ggml_can_mul_mat(a, b)`) when using `--pooling mean` with embedding models (e.g. bge-m3).

Regression introduced by #20340 (commit d28961d, first affected release: b8278). Same class of bug as #12517, fixed by #12545.

## Root cause

The chunked fused GDN detection in `sched_reserve()` calls:

```cpp
graph_reserve(16*n_seqs, n_seqs, n_outputs, mctx.get(), true);
```

where `n_outputs = n_seqs`. This builds the full computation graph including `build_pooling()`. For embedding models with mean pooling:

- `build_inp_mean()` creates a tensor with shape `[n_tokens = 16*n_seqs, n_seqs_unq]`
- `res->t_embd` is reduced to `[n_outputs = n_seqs, ...]` via `ggml_get_rows(cur, inp_out_ids)`
- `ggml_mul_mat(transpose(t_embd), inp_mean)` asserts because `n_seqs != 16*n_seqs`

The other detection calls (Flash Attention, GDN AR) pass `n_tokens=1`, which triggers the rounding guard (`n_tokens % n_seqs != 0` → `n_outputs = max(n_outputs, n_tokens)`), so they don't crash. The chunked path with `n_tokens = 16*n_seqs` passes the guard evenly and hits the mismatch.

## Fix

Pass `n_tokens` as `n_outputs` in the chunked GDN graph reservation, matching the pattern used by the pp/tg worst-case reservations:

```cpp
const uint32_t n_tokens_ch = 16*n_seqs;
auto * gf = graph_reserve(n_tokens_ch, n_seqs, n_tokens_ch, mctx.get(), true);
```

Also adds `test_embedding_pooling_mean` and `test_embedding_pooling_mean_multiple` to the server test suite — this pooling mode was previously untested.

## Reproduction

```bash
llama-server -m bge-m3-Q8_0.gguf --pooling mean --embedding -c 8192 --parallel 2
# GGML_ASSERT(ggml_can_mul_mat(a, b)) in build_pooling
```

- **Crashes on**: b8278+ (master as of 2026-03-12)
- **Works on**: b8155 and earlier
- `--pooling cls` and `--pooling last` are not affected (use `ggml_get_rows`, not `ggml_mul_mat`)

## Testing

Tested on AMD gfx1151 (ROCm 7.11) and CPU-only builds:

| Model | Architecture | Use case | Unpatched | Patched |
|-------|-------------|----------|-----------|---------|
| bge-m3 FP16 | BERT | `--pooling mean --embedding` | CRASH | OK |
| bge-m3 FP16 | BERT | `--pooling cls --embedding` | OK | OK |
| bge-m3 FP16 | BERT | `--pooling last --embedding` | OK | OK |
| Qwen3.5-4B Q4_K_M | GDN | Chat completion | OK (not affected) | OK |

Existing server test suite: 21/21 passed. 2 new mean pooling tests added and passing locally.

AI was used in an assistive capacity (analysis and drafting). All code was manually reviewed and tested.

Changed files:
filename: src/llama-context.cpp
patch: @@ -512,7 +512,12 @@ void llama_context::sched_reserve() {
 
         if (cparams.fused_gdn_ch) {
             // more than one token in the batch per sequence in order to take the chunked path
-            auto * gf = graph_reserve(16*n_seqs, n_seqs, n_outputs, mctx.get(), true);
+            // note: n_outputs must match n_tokens for embedding models with mean/rank pooling,
+            // because build_pooling creates inp_mean with shape [n_tokens, n_seqs] and multiplies
+            // it with t_embd which is reduced to [n_outputs, ...] via out_ids. if n_outputs != n_tokens,
+            // the ggml_mul_mat assertion fails. this matches the pp reservation below (line ~553).
+            const uint32_t n_tokens_ch = 16*n_seqs;
+            auto * gf = graph_reserve(n_tokens_ch, n_seqs, n_tokens_ch, mctx.get(), true);
             if (!gf) {
                 throw std::runtime_error("failed to reserve graph for fused Gated Delta Net check (chunked)");
             }

filename: tools/server/tests/unit/test_embedding.py
patch: @@ -101,6 +101,40 @@ def test_embedding_mixed_input(input, is_multi_prompt: bool):
         assert len(data[0]['embedding']) > 1
 
 
+def test_embedding_pooling_mean():
+    global server
+    server.pooling = 'mean'
+    server.start()
+    res = server.make_request("POST", "/v1/embeddings", data={
+        "input": "I believe the meaning of life is",
+    })
+    assert res.status_code == 200
+    assert len(res.body['data']) == 1
+    assert 'embedding' in res.body['data'][0]
+    assert len(res.body['data'][0]['embedding']) > 1
+
+    # make sure embedding vector is normalized
+    assert abs(sum([x ** 2 for x in res.body['data'][0]['embedding']]) - 1) < EPSILON
+
+
+def test_embedding_pooling_mean_multiple():
+    global server
+    server.pooling = 'mean'
+    server.start()
+    res = server.make_request("POST", "/v1/embeddings", data={
+        "input": [
+            "I believe the meaning of life is",
+            "Write a joke about AI",
+            "This is a test",
+        ],
+    })
+    assert res.status_code == 200
+    assert len(res.body['data']) == 3
+    for d in res.body['data']:
+        assert 'embedding' in d
+        assert len(d['embedding']) > 1
+
+
 def test_embedding_pooling_none():
     global server
     server.pooling = 'none'