[fix](kt-kernel): keep CPU expert buffers alive during CUDA graph capture - #2174
Open
05yuki wants to merge 1 commit into
Open
[fix](kt-kernel): keep CPU expert buffers alive during CUDA graph capture#217405yuki wants to merge 1 commit into
05yuki wants to merge 1 commit into
Conversation
…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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
replaying, in torch/cuda/graphs.py replay() under
prefill_cuda_graph_runner.replay_layer_forward
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:
retained, a later differently sized request replaces it, and asking for
the captured size again returns a different object
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.