Skip to content

[fix](kt-kernel): keep CPU expert buffers alive during CUDA graph capture - #2174

Open
05yuki wants to merge 1 commit into
kvcache-ai:mainfrom
05yuki:fix/kt-cpu-buffer-cuda-graph-capture
Open

[fix](kt-kernel): keep CPU expert buffers alive during CUDA graph capture#2174
05yuki wants to merge 1 commit into
kvcache-ai:mainfrom
05yuki:fix/kt-cpu-buffer-cuda-graph-capture

Conversation

@05yuki

@05yuki 05yuki commented Aug 27, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

KExpertsCPUBuffer.get_buffer() only promotes a buffer into capture_buffers
when its batch size appears in capture_bs, and capture_bs is populated
exclusively by the decode CUDA graph runner through
KTMoEWrapper.set_capture_batch_sizes(). The prefill graph runner never
registers its own num-token sizes, in this tree or upstream, so every
buffer a prefill capture touches is uncovered.

Such a capture cannot allocate the buffer itself, since pinned-host
allocation is not permitted during capture. It gets the buffer that a
warmup iteration left in the single temp slot, and the pinned pointers of
that buffer are baked into the graph. The next request at a different
batch size overwrites the temp slot, the tensors lose their last
reference, and replaying the graph reads freed pinned memory.

Retain the buffer whenever a capture is actually underway, covering both
the warmup buffer a capture adopts from the temp slot and a buffer
allocated during capture. torch.cuda.is_current_stream_capturing() is
documented to return False without initializing a CUDA context, so the
check is safe on CPU-only paths.

Reproduced and fixed end to end on 2x RTX 5070 Ti (SM120), serving
Ornith-1.5-35B-A3B (Qwen3.5-MoE, NVFP4 experts) with KT CPU experts and
the breakable prefill CUDA graph, which is the CUDA default:

  • without this patch the scheduler dies with SIGSEGV once prefill starts
    replaying, in torch/cuda/graphs.py replay() under
    prefill_cuda_graph_runner.replay_layer_forward
  • with it the same configuration starts clean and serves: an 801-token
    prefill and 128-token decode at 75.8 tok/s, coherent output

Unit-level checks against the class and a real torch.cuda.CUDAGraph
capture show the mechanism directly:

  • unpatched: the capture adopts the warmup buffer, the buffer is not
    retained, a later differently sized request replaces it, and asking for
    the captured size again returns a different object
  • patched: the same sequence keeps returning the captured buffer
  • bounding: 38 distinct eager-mode batch sizes with no capture in flight
    retain 0 buffers, where caching every size unconditionally retains all
    38. Prefill-sized entries are hundreds of MB of pinned host memory plus
    device memory each, so this distinction matters.

Before submitting

No new test files: this is a lifetime fix in an existing hot path, and the
reproduction needs a multi-GPU host serving a MoE checkpoint with the prefill
CUDA graph enabled. The repro steps and both outcomes are described above so
they can be re-run on such a host.

…ture

KExpertsCPUBuffer.get_buffer() only promotes a buffer into capture_buffers
when its batch size appears in capture_bs, and capture_bs is populated
exclusively by the decode CUDA graph runner through
KTMoEWrapper.set_capture_batch_sizes(). The prefill graph runner never
registers its own num-token sizes, in this tree or upstream, so every
buffer a prefill capture touches is uncovered.

Such a capture cannot allocate the buffer itself, since pinned-host
allocation is not permitted during capture. It gets the buffer that a
warmup iteration left in the single temp slot, and the pinned pointers of
that buffer are baked into the graph. The next request at a different
batch size overwrites the temp slot, the tensors lose their last
reference, and replaying the graph reads freed pinned memory.

Retain the buffer whenever a capture is actually underway, covering both
the warmup buffer a capture adopts from the temp slot and a buffer
allocated during capture. torch.cuda.is_current_stream_capturing() is
documented to return False without initializing a CUDA context, so the
check is safe on CPU-only paths.

Reproduced and fixed end to end on 2x RTX 5070 Ti (SM120), serving
Ornith-1.5-35B-A3B (Qwen3.5-MoE, NVFP4 experts) with KT CPU experts and
the breakable prefill CUDA graph, which is the CUDA default:

- without this patch the scheduler dies with SIGSEGV once prefill starts
  replaying, in torch/cuda/graphs.py replay() under
  prefill_cuda_graph_runner.replay_layer_forward
- with it the same configuration starts clean and serves: an 801-token
  prefill and 128-token decode at 75.8 tok/s, coherent output

Unit-level checks against the class and a real torch.cuda.CUDAGraph
capture show the mechanism directly:

- unpatched: the capture adopts the warmup buffer, the buffer is not
  retained, a later differently sized request replaces it, and asking for
  the captured size again returns a different object
- patched: the same sequence keeps returning the captured buffer
- bounding: 38 distinct eager-mode batch sizes with no capture in flight
  retain 0 buffers, where caching every size unconditionally retains all
  38. Prefill-sized entries are hundreds of MB of pinned host memory plus
  device memory each, so this distinction matters.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant