Repository: ggml-org/llama.cpp
Pull request: models : fix backend assignment for Granite/Nemotron graphs

rel https://github.com/ggml-org/llama.cpp/discussions/18598

The `ggml_backend_forward_expand()` call in `build_rs()` causes the recurrent state nodes to be materialized at the beginning of the graph, before the typical token embeddings `get_rows` (which is typically assigned to the CPU backend):

https://github.com/ggml-org/llama.cpp/blob/f6ab317edb71a976031549f495e5606b6c81cb79/src/llama-graph.cpp#L1858-L1876

This causes an extra graph split on Mac which degrades the performance, especially at high thread counts:

```bash
# master
## SPLIT #0: Metal # 2 inputs: [ (view) (   0K)] [ (view) (   0K)] 
node #  2 (     SCALE): cache_r_l0 (reshaped (   0K) [Metal         ] use=0: cache_r_l0 (reshaped (   0K) [Metal         ]
node #  4 (  GET_ROWS):               node_4 (  39K) [Metal         ] use=1: cache_r_l0 (reshaped (  39K) [Metal         ]      Metal# (view)#0 (   0K) [ NULL         ]
node #  6 (  GET_ROWS):               node_6 (   0K) [Metal         ] use=1: cache_r_l0 (reshaped (  39K) [Metal         ]      Metal# (view)#0 (   0K) [ NULL         ]
node #  8 (       CPY): cache_r_l0 (view) (c (   0K) [Metal         ] use=0:               node_6 (   0K) [Metal         ]    cache_r_l0 (view) (   0K) [Metal         ]
## SPLIT #1: CPU # 0 inputs
node # 10 (  GET_ROWS):              node_10 (   6K) [  CPU         ] use=1:    token_embd.weight ( 120M) [ BLAS         ]               leaf_4 (   0K) [  CPU         ]
## SPLIT #2: Metal # 5 inputs: [node_10 (   6K)] [leaf_82 (   0K)] [leaf_84 (   0K)] [leaf_86 (   1K)] [leaf_589 (   0K)] 
node # 11 (     SCALE):             inp_embd (   6K) [Metal         ] use=2:      Metal#node_10#0 (   6K) [ NULL         ]
node # 12 (  RMS_NORM):               norm-0 (   6K) [Metal         ] use=1:             inp_embd (   6K) [Metal         ]
node # 13 (       MUL):          attn_norm-0 (   6K) [Metal         ] use=1:               norm-0 (   6K) [Metal         ] blk.0.attn_norm.weig (   6K) [Metal         ]
node # 15 (   MUL_MAT):              node_15 (  25K) [Metal         ] use=3:  blk.0.ssm_in.weight (   5M) [Metal         ] attn_norm-0 (reshape (   6K) [Metal         ]
...
```

Fix by forcing the token embeddings to be materialized early in the graph.

```bash
# PR
## SPLIT #0: CPU # 0 inputs
node #  0 (  GET_ROWS):               node_0 (   6K) [  CPU         ] use=1:    token_embd.weight ( 120M) [ BLAS         ]               leaf_1 (   0K) [  CPU         ]

## SPLIT #1: Metal # 7 inputs: [node_0 (   6K)] [ (view) (   0K)] [ (view) (   0K)] [leaf_82 (   0K)] [leaf_84 (   0K)] [leaf_86 (   1K)] [leaf_589 (   0K)] 
node #  1 (     SCALE):             inp_embd (   6K) [Metal         ] use=2:       Metal#node_0#0 (   6K) [ NULL         ]
node #  4 (     SCALE): cache_r_l0 (reshaped (   0K) [Metal         ] use=0: cache_r_l0 (reshaped (   0K) [Metal         ]
node #  6 (  GET_ROWS):               node_6 (  39K) [Metal         ] use=1: cache_r_l0 (reshaped (  39K) [Metal         ]      Metal# (view)#0 (   0K) [ NULL         ]
node #  8 (  GET_ROWS):               node_8 (   0K) [Metal         ] use=1: cache_r_l0 (reshaped (  39K) [Metal         ]      Metal# (view)#0 (   0K) [ NULL         ]
node # 10 (       CPY): cache_r_l0 (view) (c (   0K) [Metal         ] use=0:               node_8 (   0K) [Metal         ]    cache_r_l0 (view) (   0K) [Metal         ]
node # 12 (  RMS_NORM):               norm-0 (   6K) [Metal         ] use=1:             inp_embd (   6K) [Metal         ]
node # 13 (       MUL):          attn_norm-0 (   6K) [Metal         ] use=1:               norm-0 (   6K) [Metal         ] blk.0.attn_norm.weig (   6K) [Metal         ]
node # 15 (   MUL_MAT):              node_15 (  25K) [Metal         ] use=3:  blk.0.ssm_in.weight (   5M) [Metal         ] attn_norm-0 (reshape (   6K) [Metal         ]
```

Perf changes:

| Model                         |   Threads | Test        |   t/s master |   t/s gg/granite-graph-fix-backend-assignment |   Speedup |
|:------------------------------|----------:|:------------|-------------:|----------------------------------------------:|----------:|
| granitehybrid 1B Q4_0         |         1 | pp512@d8192 |      3006.02 |                                       3012.20 |      1.00 |
| granitehybrid 1B Q4_0         |         1 | tg32@d8192  |       119.35 |                                        121.88 |      1.02 |
| granitehybrid 1B Q4_0         |        16 | pp512@d8192 |      2922.94 |                                       2998.27 |      1.03 |
| granitehybrid 1B Q4_0         |        16 | tg32@d8192  |        75.51 |                                        120.90 |      1.60 |
| granitehybrid 1B Q8_0         |         1 | pp512@d8192 |      2548.82 |                                       2551.58 |      1.00 |
| granitehybrid 1B Q8_0         |         1 | tg32@d8192  |        77.80 |                                         80.65 |      1.04 |
| granitehybrid 1B Q8_0         |        16 | pp512@d8192 |      2488.23 |                                       2546.97 |      1.02 |
| granitehybrid 1B Q8_0         |        16 | tg32@d8192  |        55.93 |                                         79.87 |      1.43 |
| nemotron_h_moe 31B.A3.5B Q8_0 |         1 | pp512@d8192 |      1576.94 |                                       1572.13 |      1.00 |
| nemotron_h_moe 31B.A3.5B Q8_0 |         1 | tg32@d8192  |        87.81 |                                         87.50 |      1.00 |
| nemotron_h_moe 31B.A3.5B Q8_0 |        16 | pp512@d8192 |      1571.38 |                                       1573.27 |      1.00 |
| nemotron_h_moe 31B.A3.5B Q8_0 |        16 | tg32@d8192  |        86.98 |                                         86.75 |      1.00 |

Changed files:
filename: src/llama-graph.cpp
patch: @@ -1326,6 +1326,10 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const {
 
     res->add_input(std::move(inp));
 
+    // make sure the produced embeddings are immediately materialized in the ggml graph
+    // ref: https://github.com/ggml-org/llama.cpp/pull/18599
+    ggml_build_forward_expand(gf, cur);
+
     return cur;
 }