[https://nvbugs/6541356][fix] align cache transceiver harness setup - #17121
Open
chuangz0 wants to merge 3 commits into
Open
[https://nvbugs/6541356][fix] align cache transceiver harness setup#17121chuangz0 wants to merge 3 commits into
chuangz0 wants to merge 3 commits into
Conversation
chuangz0
marked this pull request as ready for review
July 31, 2026 09:48
chuangz0
force-pushed
the
fix/precheck-model-env-fabric-memory
branch
from
July 31, 2026 09:48
458ac02 to
ce1f8d9
Compare
chuangz0
requested review from
QiJune,
allisonlim-nv,
brnguyen2,
dhansen-nvidia,
hyukn and
laikhtewari
July 31, 2026 09:48
Collaborator
Author
|
/bot run --stage-list "GB300-56_GPUs-14_Nodes-PyTorch-Disagg-PerfSanity-CTX12-NODE1-GPU4-GEN1-NODE2-GPU8-Post-Merge*" |
Collaborator
|
PR_Github #63016 [ run ] triggered by Bot. Commit: |
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (3)
WalkthroughThe change centralizes fabric-memory setup for Python transceivers, applies it before KV-cache allocation, propagates shell-quoted ChangesFabric-memory setup and precheck integration
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant PerfSubmit
participant PrecheckConfig
participant run_precheck
participant FabricMemoryHelper
participant KVCacheManager
PerfSubmit->>PrecheckConfig: pass shell-quoted LLM_MODELS_ROOT
PrecheckConfig->>run_precheck: generate context and generation commands
run_precheck->>FabricMemoryHelper: configure resolved transceiver and cache manager
FabricMemoryHelper->>KVCacheManager: set fabric-memory default before construction
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py (1)
38-46: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the public helper parameters.
Add a Google-style
Args:section forcache_transceiver_configandkv_cache_manager_cls. This helper is imported outside its defining module.Proposed documentation update
"""Default the C++ V1 KV pool to fabric memory for the Python transceiver. This must run before any KV pool allocation because the C++ environment getter caches the value on first read. Explicit user settings are always respected. + + Args: + cache_transceiver_config: Configuration that selects the transceiver. + kv_cache_manager_cls: KV-cache manager class to evaluate. """🤖 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 `@tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py` around lines 38 - 46, Update the docstring for maybe_enable_fabric_memory_for_python_transceiver with a Google-style Args: section documenting cache_transceiver_config and kv_cache_manager_cls, including their roles and expected types. Preserve the existing behavior and explanatory text.Source: Coding guidelines
tests/unittest/others/test_kv_cache_transceiver.py (1)
47-67: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the new test to
tests/integration/test_lists/test-db/l0_sanity_check.yml.The test covers four parameterized cases, but only other tests from
test_kv_cache_transceiver.pyare listed. Coverage verdict: insufficient.🤖 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/others/test_kv_cache_transceiver.py` around lines 47 - 67, Add test_maybe_enable_fabric_memory_for_python_transceiver from test_kv_cache_transceiver.py to tests/integration/test_lists/test-db/l0_sanity_check.yml so all four parameterized cases are included in the integration test list.Source: Path instructions
🤖 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 `@tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py`:
- Line 167: The shlex.quote() call on line 167 produces single-quote escaping
that breaks when the quoted value is embedded in the double-quoted export
statements on lines 177-180. Replace the quoting strategy for model_root_env to
survive embedding in double-quoted strings, either by passing LLM_MODELS_ROOT as
a runtime environment variable without shell string interpolation or by quoting
the entire export value instead of the path alone. Add a regression test that
validates the fix with a path containing an apostrophe.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py`:
- Around line 38-46: Update the docstring for
maybe_enable_fabric_memory_for_python_transceiver with a Google-style Args:
section documenting cache_transceiver_config and kv_cache_manager_cls, including
their roles and expected types. Preserve the existing behavior and explanatory
text.
In `@tests/unittest/others/test_kv_cache_transceiver.py`:
- Around line 47-67: Add test_maybe_enable_fabric_memory_for_python_transceiver
from test_kv_cache_transceiver.py to
tests/integration/test_lists/test-db/l0_sanity_check.yml so all four
parameterized cases are included in the integration test list.
🪄 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: 49468c21-5a6c-412d-820a-cb887e0184ef
📒 Files selected for processing (7)
examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.pytests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.pytests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.pytests/unittest/others/test_cache_transceiver_precheck_config.pytests/unittest/others/test_kv_cache_transceiver.py
Collaborator
|
PR_Github #63016 [ run ] completed with state |
Mgluhovskoi
reviewed
Jul 31, 2026
Comment on lines
+557
to
+571
| # The C++ fabric-memory env getter is cached on its first KV pool | ||
| # allocation. Enable the default before any matrix case builds a pool, even | ||
| # when a C++ transceiver case appears before Python+V1 in the matrix. | ||
| python_v1_case = next( | ||
| (case for case in cases if case["runtime"] == "PYTHON" and case["cache_manager"] == "V1"), | ||
| None, | ||
| ) | ||
| if python_v1_case is not None: | ||
| maybe_enable_fabric_memory_for_python_transceiver( | ||
| CacheTransceiverConfig( | ||
| backend=python_v1_case["backend"], | ||
| transceiver_runtime="PYTHON", | ||
| ), | ||
| KVCacheManager, | ||
| ) |
Collaborator
There was a problem hiding this comment.
Is using fabric-backed pools for C++/V1 intentional in this harness? Could this affect the C++ baseline or bandwidth comparison relative to previous runs? If C++/V1 should keep ordinary pools while Python/V1 uses fabric memory, would the cases need to run in separate processes?
Collaborator
Author
There was a problem hiding this comment.
When the KV cache pool is allocated from fabric memory, KV cache transfers on GB200 can use MNNVL and achieve higher bandwidth.
This is a standalone KV cache transfer test, should use the aligned configuration.
We have not observed any measurable performance regression from using fabric memory compared with memory allocated through cudaMalloc. In addition, KV Cache Manager V2 also allocates its
memory using VMM with fabric-memory support.
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63342 [ run ] triggered by Bot. Commit: |
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@tests/unittest/others/test_cache_transceiver_precheck_config.py`:
- Around line 368-376: Update the test setup before the precheck_prefix_lines
call to also remove the TRTLLM_DISAGG_CT_PRECHECK environment variable with
monkeypatch, while preserving the existing LLM_MODELS_ROOT cleanup and test
arguments.
In `@tests/unittest/scripts/test_perf_submit.py`:
- Around line 103-123: Add a test near test_ci_missing_model_root_is_detectable
that passes an unterminated-quote pytestCommand prefix to
ci_submit_module._get_pytest_command_env_var and asserts it raises ValueError.
Keep the existing valid-command and missing-variable coverage unchanged.
🪄 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: 0e901ae8-3a8a-490a-a39c-a64512c2bbcc
📒 Files selected for processing (8)
jenkins/scripts/perf/cluster_env.pyjenkins/scripts/perf/local/submit.pyjenkins/scripts/perf/submit.pytests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.pytests/unittest/others/test_cache_transceiver_precheck_config.pytests/unittest/others/test_cache_transceiver_precheck_run.pytests/unittest/scripts/test_cluster_env.pytests/unittest/scripts/test_perf_submit.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py
chuangz0
force-pushed
the
fix/precheck-model-env-fabric-memory
branch
from
August 3, 2026 06:44
d26e663 to
f6e7500
Compare
Collaborator
|
PR_Github #63342 [ run ] completed with state
|
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/others/test_cache_transceiver_precheck_config.py (1)
365-365: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winApply the function-annotation rule consistently to both new tests.
tests/unittest/others/test_cache_transceiver_precheck_config.py#L365-L365: annotate the fixture and model-root parameters and returnNone, or verify an explicit permissive typing exception.tests/unittest/others/test_cache_transceiver_precheck_run.py#L400-L400: annotate the fixture parameters and returnNone, or verify an explicit permissive typing exception.As per coding guidelines, Python functions must be annotated. Based on learnings, permissive untyped test methods require an explicit submodule typing configuration.
🤖 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/others/test_cache_transceiver_precheck_config.py` at line 365, Annotate the new test function test_precheck_commands_propagate_model_root in tests/unittest/others/test_cache_transceiver_precheck_config.py at lines 365-365 with types for its fixture and model-root parameters and a None return type, or add the required explicit permissive typing configuration. Apply the same fixture-parameter and None-return annotations, or equivalent typing exception, to the new test at tests/unittest/others/test_cache_transceiver_precheck_run.py lines 400-400.Sources: Coding guidelines, Learnings
🤖 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.
Nitpick comments:
In `@tests/unittest/others/test_cache_transceiver_precheck_config.py`:
- Line 365: Annotate the new test function
test_precheck_commands_propagate_model_root in
tests/unittest/others/test_cache_transceiver_precheck_config.py at lines 365-365
with types for its fixture and model-root parameters and a None return type, or
add the required explicit permissive typing configuration. Apply the same
fixture-parameter and None-return annotations, or equivalent typing exception,
to the new test at tests/unittest/others/test_cache_transceiver_precheck_run.py
lines 400-400.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 082e16ca-d7db-4e84-b872-c679fca70dd2
📒 Files selected for processing (13)
examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.pyjenkins/scripts/perf/cluster_env.pyjenkins/scripts/perf/local/submit.pyjenkins/scripts/perf/submit.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.pytests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.pytests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.pytests/unittest/others/test_cache_transceiver_precheck_config.pytests/unittest/others/test_cache_transceiver_precheck_run.pytests/unittest/others/test_kv_cache_transceiver.pytests/unittest/scripts/test_cluster_env.pytests/unittest/scripts/test_perf_submit.py
🚧 Files skipped from review as they are similar to previous changes (11)
- tests/unittest/others/test_kv_cache_transceiver.py
- jenkins/scripts/perf/local/submit.py
- tests/unittest/scripts/test_cluster_env.py
- tensorrt_llm/_torch/pyexecutor/_util.py
- jenkins/scripts/perf/cluster_env.py
- tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py
- tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py
- tests/unittest/scripts/test_perf_submit.py
- jenkins/scripts/perf/submit.py
- examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py
- tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Collaborator
Author
|
/bot run --stage-list "GB300-56_GPUs-14_Nodes-PyTorch-Disagg-PerfSanity-CTX12-NODE1-GPU4-GEN1-NODE2-GPU8-Post-Merge*" |
Collaborator
|
PR_Github #63388 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #63388 [ run ] completed with state
|
This was referenced Aug 3, 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
3 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.
Description
Fixes NVBug 6541356.
LLM_MODELS_ROOTandPYTEST_COMMON_VARSto the cache-transceiver precheck command so automatic KV-cache-manager selection can load the model configuration instead of silently falling back to V1.examples/disaggregated/slurm/cache_transceiver_testbefore KV-pool allocation.TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORYvalues and leave C++ transceiver and KV cache manager V2 behavior unchanged.Root cause
The generated precheck command did not forward the model-root/common environment. Consequently, auto manager resolution could not read the model configuration and selected cache manager V1. In addition, the standalone precheck and SLURM transceiver harness instantiate KV-cache managers directly, bypassing the serving path where #16832 enables fabric memory for the Python transceiver with V1.
Validation
pre-commit run --files <changed files>pytest -q tests/unittest/others/test_cache_transceiver_precheck_config.py tests/unittest/others/test_cache_transceiver_precheck_run.py tests/unittest/scripts/test_perf_submit.py tests/unittest/disaggregated/test_cache_transceiver_harness_report.py -k 'not TestInternalApiContract'(104 passed, 11 deselected)python3 -m py_compilefor the changed Python modules and testsThe full internal-API/GPU E2E was not run in this clean worktree because the available local TensorRT-LLM bindings predate the current
mainAPI.Dev Engineer Review
KVCacheManagerV2behavior unchanged.LLM_MODELS_ROOTandPYTEST_COMMON_VARSwith shell-safe quoting.oci-agaUCX configuration consistently.QA Engineer Review
Test code changed in:
tests/unittest/others/test_cache_transceiver_precheck_config.pytests/unittest/others/test_kv_cache_transceiver.pytests/unittest/others/test_cache_transceiver_precheck_run.pyKVCacheManagerV2selection for DeepSeek V4.tests/unittest/scripts/test_cluster_env.pyoci-aga-cs-001UCX environment.tests/unittest/scripts/test_perf_submit.pyLLM_MODELS_ROOTfrom quoted and unquoted commands, including special characters and missing values.The new fabric-memory test is listed in
tests/integration/test_lists/test-db/l0_sanity_check.yml. The other test functions are not listed intests/integration/test_lists/test-db or QA files.Targeted validation passed.
Verdict: sufficient