[None][feat] add detailed KV cache transfer time breakdown #8521

zhengd-nv · 2025-10-21T06:39:19Z

Summary by CodeRabbit

New Features
- Enhanced KV cache transfer performance metrics with detailed phase-level timing tracking (formatter, preprocess, transmissions, postprocess phases).
Refactor
- Improved API const-correctness for getter methods and timing utilities.
- Restructured timing data collection to use explicit time-point measurements instead of delay-based calculations.
Tests
- Updated performance metrics validation to reflect new timing data format with expanded phase-level information.

Description

Add more detailed KV cache transfer time and output to the folder specified by env TRTLLM_KVCACHE_TIME_OUTPUT_PATH.

Now the following timepoints are recorded in a transfer session:

Transfer Start (in perf metrics): For context, it's when sendAsync() is called and prepare for sending; for generation, it's just before RequestInfo is sent. Need to set return_perf_metrics in LLM args.
kTimeRequestInfo: For context, it's when RequestInfo is received; for generation, it's after RequestInfo is sent.
kTimeFormatter: At the beginning of CacheFormatter's format and unformat.
kTimePreprocess: Before the start of any network transmission.
kTimeTransmissions: After all transmissions finished.
kTimePostprocess: At the end of CacheFormatter's format and unformat.

And the time periods (in milliseconds) between adjacent timepoints above are recorded as

RequestInfo: Time waiting for/sending RequestInfo
Preparation: Other works before cache formatting/unformatting
Preprocess: Allocate buffers and (context) splitting cache
Transmissions: Network transmissions to different peers.
- Before this PR, only the delay(start of transmission), duration and bandwidth of each single transmission is recorded. Now they are appended to these intervals and may have multiple Delay,Duration,Bandwidth(Gbps) groups when TP>1 or PP>1.
Postprocess: (Generation) concatenating cache

The output of the file is like follows (after aligned):

rank_0_send.csv

RequestID ,RequestInfo ,Preparation ,Preprocess ,Transmissions ,Postprocess ,Delay   ,Duration ,Bandwidth(Gbps)
        1 ,  441.218   ,   8.60679  , 18.3809   ,     2.57771  ,   0.00119  ,0.00277 ,2.57384  ,        339.362
        2 ,   12.892   ,   0.046139 ,  2.56558  ,     2.66534  ,   0.00093  ,0.00088 ,2.66405  ,        327.871
        3 ,   10.7444  ,   0.086209 ,  2.54788  ,     2.6922   ,   0.00153  ,0.00128 ,2.68978  ,        324.735
        4 ,    9.98775 ,   0.043879 ,  0.447387 ,     2.67395  ,   0.00073  ,0.00068 ,2.67262  ,        313.088
        5 ,    9.77031 ,   0.046359 ,  0.448447 ,     2.66069  ,   0.00088  ,0.00066 ,2.65941  ,        314.643

rank_0_recv.csv

RequestID ,RequestInfo ,Preparation ,Preprocess ,Transmissions ,Postprocess ,Delay   ,Duration  ,Bandwidth(Gbps)
        1 , 304.73     , 126.048    ,  3.63886  ,    17.2395   ,   8.73037  ,0.00148 ,17.2366   ,        50.6748
        2 ,   0.364787 ,   0.296907 ,  0.03736  ,     5.18463  ,   0.565945 ,0.00086 , 5.18311  ,       168.521
        3 ,   0.283468 ,   0.384967 ,  0.039429 ,     5.18276  ,   0.504886 ,0.00082 , 5.18085  ,       168.595
        4 ,   0.226668 ,   0.243278 ,  0.02907  ,     3.08589  ,   0.515605 ,0.00063 , 3.08472  ,       271.261
        5 ,   0.227228 ,   0.272138 ,  0.02933  ,     3.07858  ,   0.512396 ,0.00068 , 3.07757  ,       271.891

Due to an extra copy when receiving, the effective bandwidth of recv side is lower than send side.

Test Coverage

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

coderabbitai · 2025-10-21T06:45:58Z

📝 Walkthrough

Walkthrough

Refactors timing instrumentation and const-correctness in TensorRT-LLM batch manager. Replaces simple delay/duration metrics with explicit TimePoint-based measurements, introduces timing phase markers (formatter, preprocess, transmissions, postprocess), renames environment variable from transfer to time output, updates static/mutable semantics in LlmRequest, and adjusts CSV export headers.

Changes

Cohort / File(s)	Summary
Header const-correctness and static semantics `cpp/include/tensorrt_llm/batch_manager/llmRequest.h`	Made `setKvCacheTransferStart/End` const methods with pass-by-value parameters; made `getKvCacheTransferStart/End` const; converted `getSteadyClockNow()` and `maybeToGlobalSteadyClock()` to static; renamed `mGlobalSteadyClockOffset` to `sGlobalSteadyClockOffset`; marked `mPerfMetrics` as mutable.
Timing metrics mutability `cpp/include/tensorrt_llm/executor/types.h`	Removed `mutable` qualifier from `RequestPerfMetrics::TimingMetrics::kvCacheSize`.
Environment utilities refactoring `cpp/tensorrt_llm/common/envUtils.h`, `cpp/tensorrt_llm/common/envUtils.cpp`	Renamed public function `getEnvKVCacheTransferOutputPath()` to `getEnvKVCacheTimeOutputPath()`.
Data transceiver timing infrastructure `cpp/tensorrt_llm/batch_manager/dataTransceiver.h`	Introduced `TimeNames` enum with phase markers (kTimeRequestInfo, kTimeFormatter, kTimePreprocess, kTimeTransmissions, kTimePostprocess); added `KVCacheTimes` struct; replaced `Measure` struct fields from delay/duration/bandwidth to start/end TimePoints; added `setTime(TimeNames)` method; added `appendMeasure(TimePoint, TimePoint, size_t)` overload; updated constructor parameter from `recordMeasure` to `recordTiming`.
Data transceiver implementation `cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp`	Implemented `setTime()` and `appendMeasure()` overloads; updated `exportMeasure()` to use mTimes with new CSV header format; integrated time-based delay calculations; updated environment variable usage to `getEnvKVCacheTimeOutputPath()`.
Cache formatting timing instrumentation `cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp`	Added `session.setTime()` calls at formatter, preprocess, transmissions, and postprocess phases; replaced ad-hoc delay calculations with `session.appendMeasure(startTime, endTime, size)`.
MLA cache formatter timing `cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp`	Added timing phase markers and per-block measurement calls in both format and unformat paths; removed legacy delay computation logic.
Cache transceiver environment check `cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp`	Updated KV cache transfer bandwidth update gating to check `getEnvKVCacheTimeOutputPath()` instead of transfer output path.
Python bindings `cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp`	Updated global static member binding from `mGlobalSteadyClockOffset` to `sGlobalSteadyClockOffset`.
Integration test `tests/integration/defs/disaggregated/test_disaggregated.py`	Updated CSV header expectation from "RequestID,Delay(ms),Duration(ms),Bandwidth(Gbps)" to "RequestID,RequestInfo,Preparation,Preprocess,Transmissions,Postprocess"; increased minimum field count from 4 to 9.

Sequence Diagram

sequenceDiagram
    participant Client
    participant Session
    participant Formatter
    participant Transfer
    participant Export

    Client->>Session: create TransferSession(recordTiming=true)
    Session->>Session: mTimes = make_unique<KVCacheTimes>()
    
    Formatter->>Session: setTime(kTimeFormatter)
    Session->>Session: times[kTimeFormatter] = now()
    
    Formatter->>Session: setTime(kTimePreprocess)
    Session->>Session: times[kTimePreprocess] = now()
    
    Transfer->>Session: setTime(kTimeTransmissions)
    Session->>Session: times[kTimeTransmissions] = now()
    
    Transfer->>Session: appendMeasure(startTime, endTime, size)
    Session->>Session: measures.push_back({start, end, size})
    
    Transfer->>Session: setTime(kTimePostprocess)
    Session->>Session: times[kTimePostprocess] = now()
    
    Client->>Export: exportMeasure()
    Export->>Export: calc delays from times[] and measures
    Export->>Export: write "RequestID,RequestInfo,Preparation,..."
    Export->>Client: CSV with expanded timing info

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

This refactor spans multiple interconnected systems with both API-breaking changes and semantic shifts. Review requires understanding: (1) const-correctness implications and static method migrations in LlmRequest; (2) the new TimePoint-based measurement model and how it replaces delay/duration calculations; (3) the cascading updates across cache formatters, transceiver, and export logic; (4) the environment variable rename propagation; and (5) test assertion adjustments reflecting the new CSV structure. The heterogeneity of changes across timing infrastructure, header files, and implementations increases cognitive load despite individual file changes being relatively localized.

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 15.22% which is insufficient. The required threshold is 80.00%.	You can run `@coderabbitai generate docstrings` to improve docstring coverage.
Description Check	⚠️ Warning	The PR description provides a well-structured explanation with a properly formatted title `[None][feat] add detailed KV cache transfer time breakdown` that follows the required template. The Description section is comprehensive and clearly explains the feature being added, including detailed timepoints, measured periods, and CSV output examples that effectively illustrate the changes. The PR Checklist section is present with the final confirmation checkbox marked. However, the Test Coverage section is completely empty with no information about relevant tests, which is a required section explicitly outlined in the template that should "list clearly what are the relevant test(s) that can safeguard the changes in the PR." While test changes are visible in the code modifications (e.g., updates to test_disaggregated.py), they are not documented in the designated section.	Fill out the Test Coverage section with a description of the relevant tests that validate these changes. At minimum, document that the disaggregated test (tests/integration/defs/disaggregated/test_disaggregated.py) has been updated to validate the new CSV header format and structure. Additionally, explicitly document which other existing or new tests provide coverage for the new timepoint recording functionality and the changes to the TransferSession API.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title Check	✅ Passed	The PR title "[None][feat] add detailed KV cache transfer time breakdown" directly and clearly describes the primary change in the pull request. The changeset focuses on adding granular timing tracking with multiple timepoints (kTimeRequestInfo, kTimeFormatter, kTimePreprocess, kTimeTransmissions, kTimePostprocess) and computing time periods between them to provide a detailed breakdown of KV cache transfer durations. The title accurately reflects this main objective without being vague or misleading.

✨ Finishing touches

📝 Generate docstrings

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

coderabbitai

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp (1)

416-436: Possible buffer index mismatch can send the wrong slice.

Else-path uses outputSplitCaches[processIdx] while if-path uses outputSplitCaches[bufferIdx]. These indices can differ; sizes/tensors may not correspond to the intended peer, risking corruption or mis-accounting.

Apply consistent indexing:

-                size_t remainSendSize = outputSplitCaches[processIdx]->getSize();
-                size_t needSendSize = outputSplitCaches[processIdx]->getSize();
+                size_t remainSendSize = outputSplitCaches[bufferIdx]->getSize();
+                size_t needSendSize = outputSplitCaches[bufferIdx]->getSize();
@@
-                    auto copySlice = runtime::ITensor::slice(
-                        outputSplitCaches[bufferIdx], needSendSize - remainSendSize, sendSize);
+                    auto copySlice = runtime::ITensor::slice(
+                        outputSplitCaches[bufferIdx], needSendSize - remainSendSize, sendSize);
@@
-                    session.send(processIdx, copyTargetSlice->data(), copyTargetSlice->getSizeInBytes());
+                    session.send(processIdx, copyTargetSlice->data(), copyTargetSlice->getSizeInBytes());

Also set size from the same index:

-            size_t size = outputSplitCaches[bufferIdx]->getSizeInBytes();
+            size_t size = outputSplitCaches[bufferIdx]->getSizeInBytes();

cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp (1)

743-792: Data race: kvCacheSize updates from parallel receive tasks.

recvBufferFun calls llmRequest.updateKvCacheSize(...) on multiple threads when parallel receive is enabled. If the setter isn’t atomic/thread‑safe, this is undefined behavior.

Proposed fix: accumulate per-task sizes and update once on the main thread.

@@
-            auto recvBufferFun = [&](int deviceId, size_t processIdx)
+            // Per-connection byte counters
+            std::vector<size_t> perConnBytes(pickUpConnections.size(), 0);
+            auto recvBufferFun = [&](int deviceId, size_t processIdx)
             {
@@
-                size_t size = 0;
+                size_t size = 0;
@@
-                    llmRequest.updateKvCacheSize((*recvSplitCaches.at(processIdx)).getSizeInBytes());
                     auto& buffer = recvSplitCaches[processIdx];
                     size = buffer->getSizeInBytes();
@@
-                        size += recvSlice->getSizeInBytes();
-                        llmRequest.updateKvCacheSize((*recvSlice).getSizeInBytes());
+                        size += recvSlice->getSizeInBytes();
@@
-                auto endTime = llmRequest.getSteadyClockNow();
-                session.appendMeasure(startTime, endTime, size);
+                auto endTime = llmRequest.getSteadyClockNow();
+                session.appendMeasure(startTime, endTime, size);
+                perConnBytes[processIdx] = size;
             };
@@
-            else
-            {
-                recvBufferFun(deviceId, 0);
-            }
+            else { recvBufferFun(deviceId, 0); }
+            // Single-thread update to avoid data races
+            {
+                size_t total = 0;
+                for (auto b : perConnBytes) { total += b; }
+                if (total) { llmRequest.updateKvCacheSize(total); }
+            }

If updateKvCacheSize is an additive API, the above remains correct; otherwise replace with a set/assign variant.

🧹 Nitpick comments (2)

cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp (1)

392-441: Minor: measure size accounting on send path can be wrong in fallback branch.

size is taken from outputSplitCaches[bufferIdx], but the fallback copies from a potentially different slice (processIdx). Use the same index for both branches to keep Duration/Bandwidth accurate.

See the related fix suggested in cacheTransceiver.cpp for consistent indexing.
cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp (1)
246-246: Consider calling static method directly for clarity.

The method getSteadyClockNow() is now static (as changed in llmRequest.h), so calling it via the instance (llmRequest.getSteadyClockNow()) works but is less clear than calling it directly as LlmRequest::getSteadyClockNow().

Apply this diff to improve clarity:
-        auto startTime = llmRequest.getSteadyClockNow();
+        auto startTime = LlmRequest::getSteadyClockNow();
(Repeat for line 448 and line 279/481 if capturing endTime similarly)

Also applies to: 448-448

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between a4227cf and e76a32d.

📒 Files selected for processing (11)

cpp/include/tensorrt_llm/batch_manager/llmRequest.h (5 hunks)
cpp/include/tensorrt_llm/executor/types.h (1 hunks)
cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp (9 hunks)
cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp (1 hunks)
cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp (10 hunks)
cpp/tensorrt_llm/batch_manager/dataTransceiver.h (3 hunks)
cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp (9 hunks)
cpp/tensorrt_llm/common/envUtils.cpp (1 hunks)
cpp/tensorrt_llm/common/envUtils.h (1 hunks)
cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp (1 hunks)
tests/integration/defs/disaggregated/test_disaggregated.py (1 hunks)

🧰 Additional context used

📓 Path-based instructions (8)

**/*.{h,hpp,hh,hxx,cpp,cxx,cc,cu,cuh}