Reduce decode CUDA-graph memory with shared pool and shared logits buffer

[gameofdimension]

Description

Reduces the memory cost of the SID-GR decode CUDA graph so decode at beam_width=1024 no longer OOMs. At beam=1024 the decode graph captured ~30 entries, each reserving its own CUDA-graph private pool and a [batch, beam, vocab] logits tensor, raising idle GPU memory from ~24 GiB to ~45 GiB and triggering OOM during decode. Both are fixed without altering any computation, so outputs are numerically unchanged.

1. Shared graph private pool — the prefill piecewise runner already shared one graph_pool_handle() across its captures; the decode runner did not, and got a fresh private pool per capture. The lazy-handle + pool= kwargs + env opt-out are hoisted onto CudaGraphCacheMixin (deduplicating the two runners' identical helpers), and each runner sets a _separate_pools_env class attribute. Decode captures now reuse one pool instead of one per entry.
   Opt-out: GR_INFERENCE_DECODE_CUDA_GRAPH_SEPARATE_POOLS / GR_INFERENCE_PREFILL_CUDA_GRAPH_SEPARATE_POOLS.
2. Shared logits output buffer — _linear_project gains an out= path (matmul out= into a caller buffer) and forward_decode_step a logits_out= arg; the decode runner allocates one [batch, beam, vocab] buffer per (active_beam_width, bucket) and writes every captured graph's lm_head output into it, so the logits lives once per shape rather than once per graph.
   Opt-out: GR_INFERENCE_DECODE_CUDA_GRAPH_SEPARATE_LOGITS.

New /status fields: decode_cuda_graph_shared_pool, decode_cuda_graph_shared_output_buffers.

Measured (NVIDIA L20 46 GiB, beam=1024, context_len=4000): idle after warmup 45.4 → 26.9 GiB (decode-graph overhead ~21 → ~3 GiB); beam=1024 decode that previously OOM'd now succeeds; decode_cuda_graph_shared_pool=1, shared_output_buffers=5, replays confirm the shared path is hit.

Accuracy: graph-ON output_ids match the eager (graph-OFF) baseline token-for-token across single / batch2 / batch4 (24/24, 0 mismatch).

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[greptile-apps]

Greptile Summary

This PR reduces decode CUDA graph memory use for large beam sizes. The main changes are:

• Shared CUDA graph private pools for decode and prefill piecewise captures.
• A shared decode logits output buffer per beam and bucket shape.
• LRU-based cleanup for shared logits buffers when graph shapes are evicted.
• Status fields and tests for pool sharing and shared logits reuse.

Confidence Score: 5/5

This looks safe to merge.

• No blocking issues found in the changed code.

Important Files Changed

Filename	Overview
examples/sid-gr-inference/src/gr_inference/gr_serving/decode_cuda_graph.py	Adds decode pool sharing, shared logits buffers, status counters, and buffer cleanup tied to graph cache lifetime.
examples/sid-gr-inference/src/gr_inference/gr_models/qwen3/model.py	Adds an optional projection output buffer path for decode logits.
examples/sid-gr-inference/src/gr_inference/gr_serving/cuda_graph_utils.py	Moves CUDA graph pool-sharing helper logic into the shared cache mixin.
examples/sid-gr-inference/src/gr_inference/gr_serving/prefill_cuda_graph.py	Switches prefill piecewise capture to the inherited pool-sharing helper.
examples/sid-gr-inference/tests/test_decode_shared_logits.py	Adds tests for shared logits projection, reuse, eviction cleanup, and caching-disabled behavior.
examples/sid-gr-inference/tests/test_cuda_graph_shared_pool.py	Adds tests for graph pool handle reuse and env opt-out behavior.

_{Reviews (4): Last reviewed commit: "Merge branch 'main' into enh-reduce-deco..." | Re-trigger Greptile}

[gameofdimension]

@cb521 @shijieliu would you please have a look

[cb521]

@gameofdimension yes，I have a fix for this feature, I will reply the customer, THX.

[gameofdimension]

@gameofdimension yes，I have a fix for this feature, I will reply the customer, THX.

Is the approach in this PR a correct solution? Does it have any chance of being merged? @cb521

[cb521]

@gameofdimension Sure, it's a common impl, I will merge it to here.

[gameofdimension]

@gameofdimension Sure, it's a common impl, I will merge it to here.

I'm a bit confused...

Specifically, I'm referring to the GPU memory issue and the proposed fix in this PR.