pytorch / pytorch/executorch

Support CUDA graph replay with multi-session LLM serving

Open
#20,310 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
5k
Forks
1.2k
Avg merge
2d 10h
Merged PRs (30d)
581

Description

Description:

Today the Qwen3.5-MoE engine treats CUDA graph decode and multi-session mutable-state rebinding as mutually exclusive. This is intentional: CUDA graph replay captures device pointers, while multi-session serving swaps per-session mutable
buffers before execution. Replaying a graph captured for one session after rebinding to another session can silently use the wrong KV/recurrent buffers.

We should investigate a design that supports both:

  • one physical model / one weight load
  • multiple isolated logical sessions
  • CUDA graph optimized decode
  • no state bleed across sessions

A likely direction is to make CUDA graph state session-aware. For example, graph capture could be associated with the active session’s mutable-buffer bindings, and when execution switches to a different session the backend either recaptures
the graph for that session or selects a previously captured graph for that session.

Key questions:

  • Should CUDA graph configuration/state live at the LLM session layer, CUDA backend layer, or both?
  • Can we cache one captured decode graph per active session?
  • What is the memory overhead per captured graph?
  • Is recapture-on-session-switch cheap enough, or should capture happen lazily on first decode per session?
  • How do we fail closed if a graph is replayed with stale mutable-buffer bindings?
  • How does this interact with MutableStateContextOwner and weight_sharing_across_methods?

Acceptance criteria:

  • Multi-session no-bleed test still passes with CUDA graph enabled.
  • Decode graph replay uses the correct session’s mutable buffers after A → B → A session switches.
  • If graph capture cannot be made safe, the engine keeps rejecting/disable graph mode for multi-session with a clear diagnostic.
  • Benchmark single-session and multi-session decode throughput with and without graph capture.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by tracing MutableStateContextOwner and weight_sharing_across_methods, then run or inspect the existing multi-session no-bleed test. Determine where CUDA graph capture and mutable-buffer rebinding are owned, and evaluate per-session caching or safe rejection. Done means correct A → B → A replay behavior, a clear diagnostic when unsafe, and the requested throughput benchmarks.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.