[TRTLLM-10947][perf] eagle3: use cudaMemcpy2DAsync custom op for hidden-state capture#14479
Merged
Merged
Conversation
…en-state capture Add trtllm::inplace_slice_copy custom op backed by cudaMemcpy2DAsync, and use it in Eagle3OneModelSpecMetadata.maybe_capture_hidden_states to replace the Python slice + Tensor.copy_(non_blocking=True) path. The op is registered with a Python fake and added to the piecewise CUDA graph inplace_map so it remains capturable. Signed-off-by: Pietro Cicotti <5833013+pcicotti@users.noreply.github.com>
Collaborator
Author
|
/bot run |
Contributor
📝 WalkthroughWalkthroughThis PR introduces ChangesInplace Slice Copy Custom CUDA Op
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/_torch/thop/parallel_hw_agnostic/test_inplace_slice_copy.py (1)
1-117:⚠️ Potential issue | 🟠 Major | ⚖️ Poor tradeoffAdd/extend CI perf regression coverage for the Eagle3 path
tests/integration/test_lists/qa/llm_perf_sanity.ymlalready has an Eagle3 entry (perf/test_perf.py::test_perf[qwen3_4b_eagle3-...]), so QA sanity runs have coverage.tests/integration/test_lists/test-db/l0_perf.ymlcontains noeagle/eagle3entries, meaning CI perf regression won’t track this optimization—add/update a matching Eagle3 scenario there (using the PR’s model/settings if different fromqwen3_4b_eagle3).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/thop/parallel_hw_agnostic/test_inplace_slice_copy.py` around lines 1 - 117, Add an Eagle3 perf scenario to the CI perf-regression list by updating the l0_perf.yml test list to include the same Eagle3 entry used by QA (the perf/test_perf.py::test_perf[qwen3_4b_eagle3-...] style scenario); if the PR introduces different model/settings, use those exact model ID and settings for the new entry so CI will track regressions for the Eagle3 path. Ensure the new entry matches the format of other entries in l0_perf.yml and passes existing job filters so it runs in the perf regression job.
🧹 Nitpick comments (2)
tests/unittest/_torch/thop/parallel_hw_agnostic/test_inplace_slice_copy.py (2)
104-116: ⚡ Quick winConsider adding error test coverage for dimension mismatch.
The current error tests validate dtype mismatch and out-of-bounds slice ranges. Consider adding a test for the case where
src.shape[1] != (dim1_end - dim1_start), which is another likely user error.Suggested additional test case
def test_dimension_mismatch_raises(): """src width must match slice width.""" dtype = torch.bfloat16 dest = torch.zeros(8, 32, dtype=dtype, device="cuda") src = torch.randn(8, 16, dtype=dtype, device="cuda") # 16 != (24 - 8) with pytest.raises(RuntimeError): torch.ops.trtllm.inplace_slice_copy(dest, src, 8, 24)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/thop/parallel_hw_agnostic/test_inplace_slice_copy.py` around lines 104 - 116, Add a unit test that asserts a RuntimeError is raised when the source width does not match the requested slice width: create a new test function (e.g., test_dimension_mismatch_raises) that constructs dest and src tensors with matching batch dim but src.shape[1] != (dim1_end - dim1_start) and calls torch.ops.trtllm.inplace_slice_copy(dest, src, dim1_start, dim1_end) inside pytest.raises(RuntimeError); this mirrors the existing test patterns (test_dtype_mismatch_raises, test_out_of_bounds_raises) and verifies the inplace_slice_copy validation for dimension mismatch.
1-117: QA test list updates are not required for this unit test.This is a focused unit test under
tests/unittest/validating a low-level custom CUDA operator. It does not require entries in the QA integration test lists (tests/integration/test_lists/qa/), as those are intended for end-to-end functional or multi-GPU integration scenarios.As per coding guidelines: unit-level tests do not need QA list updates unless they represent new end-to-end integration test scenarios.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/thop/parallel_hw_agnostic/test_inplace_slice_copy.py` around lines 1 - 117, This unit test file tests a low-level CUDA op (tests/unittest/_torch/thop/parallel_hw_agnostic/test_inplace_slice_copy.py) and does not require QA integration-list updates, so remove any added entries under tests/integration/test_lists/qa/ (if you added them) and do not modify those QA lists; ensure the PR description explicitly states these changes are unit-only (reference tests like test_full_dest_full_width, test_partial_rows, test_layered_capture_pattern) so reviewers know no QA-list changes are necessary.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/tensorrt_llm/thop/inplaceSliceCopyOp.cpp`:
- Around line 34-47: dim1Start and dim1End must be validated as non-negative and
device-safe before any pointer arithmetic and kernel launch: add a
TORCH_CHECK(dim1Start >= 0, "dim1Start must be non-negative") and ensure dim1End
> dim1Start (already checked) and dim1End <= dest.size(1) remains; also add a
device equality check (e.g., TORCH_CHECK(dest.device() == src.device(), "dest
and src must be on the same device")) before obtaining raw pointers or launching
CUDA work so destPtr/srcPtr aren't derived from tensors on different devices;
apply the same non-negative index and device checks in the other block
referenced (lines ~59-64) and perform pointer acquisition (data_ptr) only after
these checks.
In `@tensorrt_llm/_torch/speculative/eagle3.py`:
- Around line 443-446: The custom op call should only copy rows for active
tokens; before calling inplace_slice_copy in eagle3.py, truncate to_save to at
most self.num_tokens (e.g., use to_save = to_save[:self.num_tokens]) and pass
that smaller row count into inplace_slice_copy so you don't overwrite inactive
rows of self.hidden_states; update references around inplace_slice_copy,
hidden_states, to_save and ensure this honors any changes made in prepare().
---
Outside diff comments:
In `@tests/unittest/_torch/thop/parallel_hw_agnostic/test_inplace_slice_copy.py`:
- Around line 1-117: Add an Eagle3 perf scenario to the CI perf-regression list
by updating the l0_perf.yml test list to include the same Eagle3 entry used by
QA (the perf/test_perf.py::test_perf[qwen3_4b_eagle3-...] style scenario); if
the PR introduces different model/settings, use those exact model ID and
settings for the new entry so CI will track regressions for the Eagle3 path.
Ensure the new entry matches the format of other entries in l0_perf.yml and
passes existing job filters so it runs in the perf regression job.
---
Nitpick comments:
In `@tests/unittest/_torch/thop/parallel_hw_agnostic/test_inplace_slice_copy.py`:
- Around line 104-116: Add a unit test that asserts a RuntimeError is raised
when the source width does not match the requested slice width: create a new
test function (e.g., test_dimension_mismatch_raises) that constructs dest and
src tensors with matching batch dim but src.shape[1] != (dim1_end - dim1_start)
and calls torch.ops.trtllm.inplace_slice_copy(dest, src, dim1_start, dim1_end)
inside pytest.raises(RuntimeError); this mirrors the existing test patterns
(test_dtype_mismatch_raises, test_out_of_bounds_raises) and verifies the
inplace_slice_copy validation for dimension mismatch.
- Around line 1-117: This unit test file tests a low-level CUDA op
(tests/unittest/_torch/thop/parallel_hw_agnostic/test_inplace_slice_copy.py) and
does not require QA integration-list updates, so remove any added entries under
tests/integration/test_lists/qa/ (if you added them) and do not modify those QA
lists; ensure the PR description explicitly states these changes are unit-only
(reference tests like test_full_dest_full_width, test_partial_rows,
test_layered_capture_pattern) so reviewers know no QA-list changes are
necessary.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b14b389a-cb90-456c-8f8d-7138000a3e58
📒 Files selected for processing (7)
cpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/thop/inplaceSliceCopyOp.cpptensorrt_llm/_torch/compilation/utils.pytensorrt_llm/_torch/custom_ops/__init__.pytensorrt_llm/_torch/custom_ops/cpp_custom_ops.pytensorrt_llm/_torch/speculative/eagle3.pytests/unittest/_torch/thop/parallel_hw_agnostic/test_inplace_slice_copy.py
Collaborator
|
PR_Github #50046 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #50046 [ run ] completed with state |
Address review feedback on inplace_slice_copy: - TORCH_CHECK that dest and src are on the same CUDA device before obtaining raw pointers and launching the copy on dest's stream. - TORCH_CHECK that dim1Start is non-negative so the destPtr offset cannot underflow before the dest allocation. Add unit tests test_negative_dim1_start_raises and test_device_mismatch_raises covering the new error paths. Signed-off-by: Pietro Cicotti <5833013+pcicotti@users.noreply.github.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #50106 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #50106 [ run ] completed with state |
liji-nv
approved these changes
Jun 1, 2026
pcicotti
enabled auto-merge (squash)
June 1, 2026 13:47
hyukn
approved these changes
Jun 1, 2026
juney-nvidia
approved these changes
Jun 1, 2026
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
5 participants
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.
Summary
Replaces the Python
slice + Tensor.copy_(non_blocking=True)path inEagle3OneModelSpecMetadata.maybe_capture_hidden_stateswith a newtrtllm::inplace_slice_copycustom op backed bycudaMemcpy2DAsync.The op:
src[:numTokens, :]intodest[:numTokens, dim1_start:dim1_end]in one device-to-device 2-D memcpy on the current PyTorch CUDA stream.register_fake(no-op) so it composes withtorch.compile/ FX tracing.inplace_map(destis the mutated argument) so the capture continues to inline the write rather than treating it as an external mutation.Motivation
On the Eagle3 capture path the previous code dispatched a Python
aten::slice+aten::copy_per captured layer per step. Even withnon_blocking=Truethese show up as host-side overhead on the compute timeline and prevent the copy from being fully absorbed into the surrounding CUDA-graph region. Lowering the write to a singlecudaMemcpy2DAsyncissued on the current stream removes those launches from the timeline.Performance
Measured on H100 with the EAGLE3 hidden-state-capture repro used during development:
nsys timeline
Test plan
tests/unittest/_torch/thop/parallel_hw_agnostic/test_inplace_slice_copy.pycovers:/bot runonce this leaves draft.Files
cpp/tensorrt_llm/thop/inplaceSliceCopyOp.cpp— new op (cudaMemcpy2DAsync)cpp/tensorrt_llm/thop/CMakeLists.txt— wire op intoth_commontensorrt_llm/_torch/custom_ops/__init__.py— Python wrapper + exporttensorrt_llm/_torch/custom_ops/cpp_custom_ops.py—register_faketensorrt_llm/_torch/compilation/utils.py— add toinplace_info()maptensorrt_llm/_torch/speculative/eagle3.py— call the op inmaybe_capture_hidden_statestests/unittest/_torch/thop/parallel_hw_agnostic/test_inplace_slice_copy.py— new unit testsSummary by CodeRabbit
New Features
Tests