[joelagnel]

Implement graph shape keying to support multiple CUDA graphs with different first node shapes.

Rewrote original PR (PR #28666) by @Lisiyuan233 to be simpler and added bounded
caching with frequency-based eviction. With this, I retain most of the benefits
I mentioned in my comment on #28666 (comment)

Model: Qwen3.8-27B Q4_K_M target; DFlash2 rows also used the Qwen3.8-27B-DFlash2 Q4_K_M draft.

Device	Test	Upstream tok/s	Candidate tok/s	Change [95% CI]
GB300	llama-server DFlash2-7 decode, P=1, 512+256 tokens	122.244	126.801	+3.728% [+3.470%, +3.988%]
RTX PRO 6000 Blackwell	llama-server DFlash2-7 decode, P=1, 512+256 tokens	146.681	154.355	+5.230% [+4.781%, +5.681%]
GB300	llama-server ordinary decode, P=1, 512+256 tokens	91.414	91.564	+0.164% [+0.085%, +0.243%]
RTX PRO 6000 Blackwell	llama-server ordinary decode, P=1, 512+256 tokens	72.268	72.374	+0.147% [+0.106%, +0.188%]

llama-server test: eight fresh-server baseline/candidate pairs per device and
mode, alternating arm order. Each server received three sequential /completion
requests with distinct fixed 512-token prompts and n_predict=256; the first
conditioned the server. Decode tok/s was the average of 256000 / predicted_ms
for responses 2 and 3. DFlash2-7 used a maximum of seven draft tokens.

I also ran llama-batched-bench and do not see any change in performance with this PR.

Tested also Aman's PR #30804 and I see improvements with this PR both
with/without that PR, so this PR is additive in improvement on top of the
chunked graph PR.

• AI usage disclosure: YES - Used AI for tracing the issue, designing the algorithm, and as coding assistant.

CC @ORippler @gaugarg-nv @am17an

[am17an]

I think this problem needs to be solved the llama.cpp level where materially different graphs (i.e. differing n_tokens) should come via different memory arenas. Currently there are rebuilds at

1. scheduler level
2. llama.cpp level
3. CUDA backend level

This PR seems to address 3 but IMO we can tackle at least 2. Spec-dec workloads have become common and the rebuilds on the scheduler side probably are costing some cpu time. I think the least we can do here it to create separate arenas like done for MTP in #28549. For small batches (n_tokesn <=8) we can separate llama_graph_result_ptr

llama.cpp/src/llama-context.h

Lines 414 to 415 in 3982664

	// Separate arenas give batches with and without outputs distinct CUDA graph cache keys.
	std::array<llm_graph_result_ptr, 2> gf_res_prev;

Note this would miss the uid check as the scheduler will rebuild the graph and generate a new uid, but we should ideally fix this either having a scheduler per graph or modifying the scheduler to store multiple graphs.

[joelagnel]

Thank you so much, Aman, for the feedback! This is very useful, I agree that addressing at higher levels first would likely be a better approach. I will investigate the issue along the same lines in the coming days and let you know what I find. FWIW, I tried to keep this PR small so complexity wise it is a small change, but we can explore solving it at higher level first as you suggest.