RPC Tensor-parallel: Caching graphs to reduce network overhead

[Alexander2995]

Hi llama community,

I was doing research on USB4 transport for RPC in split_mode tensor and found that before almost every compute phase server needs to get new graph again.

There is existing feature to reuse a graph but it works just for a single graph instance, for Gemma 31B there had been 241 graphs.

With help of AI I have made a patch to handle caching every reused graph by UID (used at least twice).

On my setup of 2 GMKtec K12 with AMD Ryzen 7 H255 PCs I have tested it and got roughly 2 times boost on token generation phase compared to master branch.

My setup:

K12-1 Ryzen H255 64gb DDR5 5600
K12-2 Ryzen H255 32gb DDR5 5600
transport - Ethernet 2.5G
API - Vulkan

Tests performed:

Model and mode	pp64 tok/s	tg32 tok/s
Qwen3 4B Q8_0, K12-1 solo	393.674985 +/- 12.040010	18.040894 +/- 0.429658
Qwen3 4B Q8_0, upstream TCP RPC	63.902555 +/- 1.447661	5.411568 +/- 0.002232
Qwen3 4B Q8_0, graph-cache TCP RPC	66.181528 +/- 0.286268	13.211557 +/- 0.324566
Gemma 4 31B Q4_0, K12-1 solo	73.446822 +/- 0.900541	4.486488 +/- 0.019191
Gemma 4 31B Q4_0, upstream TCP RPC	18.288838 +/- 0.026488	1.637130 +/- 0.001639
Gemma 4 31B Q4_0, graph-cache TCP RPC	18.775400 +/- 0.012552	4.111661 +/- 0.139460

Could it be considered for PR/Enhancement? (code is tested on my setup)

I have made runtime option to set size for the cache.
Server reports feature in HELLO message and size is requested with separate new command.
Client manages caches for each server separately.
When overfilled cache resets on both sides with short log message. (tested)
Modification is backward-compatible. (older v7 nodes/clients will work without caching)

my fork with mod

Feedback/reports of tests on another setups are appreciated.

Regards,
Alexander2995

Topics related:

PR #26610 - RPC: add '-sm tensor'
PR #22701 — rpc: use graph uid instead of graph cache
Issue #22850 — Essential RPC performance degradation

[am17an]

Have you seen #26610?

[Alexander2995]

Hi, yes I've seen it couple weeks ago, when was starting experiments with usb4stream as transport.

After adding simple transport wrapper when I started testing I found that graphs are creating most of the overhead and switched my direction.

Thanks for reminding, I have checked again and found it has much in common.
Could the patch be considered as improvement for graph cache feature in that branch then?

It has flexible implementation, and caching from second appearance was introduced, since caching from the first time UID met was reducing prefill speeds