[https://nvbugs/6503299][fix] Default fabric memory KV pool for Python cache transceiver - #16832
Conversation
chuangz0
changed the title
|
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 (10)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (9)
WalkthroughThe change enables Fabric-memory pooling for compatible Python/V1 KV-cache setups and adds a 512 MB KV-cache bounce size to eight GB200 GPT-OSS NIXL benchmark configurations. ChangesKV-cache Fabric memory configuration
Estimated code review effort: 2 (Simple) | ~15 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/_util.py (1)
347-347: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the required return annotation.
As per coding guidelines, every Python function must be annotated; declare this helper as returning
None.Proposed fix
-def _maybe_enable_fabric_memory_for_python_transceiver(self): +def _maybe_enable_fabric_memory_for_python_transceiver(self) -> None:🤖 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/_util.py` at line 347, Update _maybe_enable_fabric_memory_for_python_transceiver to include an explicit None return annotation, preserving its existing behavior.Source: Coding guidelines
🤖 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 `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Line 347: Update _maybe_enable_fabric_memory_for_python_transceiver to include
an explicit None return annotation, preserving its existing behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e54489eb-50c5-42cb-9963-81b94b2899ed
📒 Files selected for processing (9)
tensorrt_llm/_torch/pyexecutor/_util.pytests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_1k1k_con64_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yamltests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yamltests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yamltests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb200_stress-gpt-oss-120b-fp4_8k1k_ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yaml
442491c to
ab93d7a
Compare
|
/bot run --disable-fast-fail |
|
PR_Github #61828 Bot args parsing error: usage: /bot [-h] |
|
/bot run --disable-fail-fast |
|
PR_Github #61833 [ run ] triggered by Bot. Commit: |
|
PR_Github #61833 [ run ] completed with state
|
Default TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY=1 when the Python cache transceiver (transceiver_runtime='PYTHON') is used with the C++ V1 KV cache manager, so the KV pool is allocated with fabric memory for MNNVL transfers. Explicit user settings are respected, and unsupported platforms fall back to standard allocation in C++. Also set kv_cache_bounce_size_mb=512 in the GPT-OSS disagg perf configs where the ctx/gen KV-head sharding differs (ctx TP1 -> gen TP4), which requires the bounce buffer to regroup KV heads during transfer. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
7014802 to
7821757
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #62056 [ run ] triggered by Bot. Commit: |
| perf/test_perf_sanity.py::test_e2e[aggr_upload-k25_thinking_fp4_2_nodes_grace_blackwell-k25_thinking_fp4_tep8_32k8k] SKIP (https://nvbugs/6422339) | ||
| perf/test_perf_sanity.py::test_e2e[aggr_upload-super_ad_blackwell-super_ad_ws1_1k1k] SKIP (https://nvbugs/6153575) | ||
| perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb200_deepseek-r1-fp4_128k8k_con128_ctx1_pp8_gen1_dep16_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6426890) | ||
| perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_glm-5-fp4_8k1k_con1024_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL] SKIP (https://nvbugs/6517846) |
There was a problem hiding this comment.
I see 7 cases being waived, do you missed to unwaive others? Thanks!
There was a problem hiding this comment.
Not missed — https://nvbugs/6517846 is a post-merge bug that covers multiple failures with different root causes. I think only the 3 cases are related by the fix.
| if not issubclass(self._kv_cache_manager_cls, KVCacheManager): | ||
| return | ||
| if os.environ.get("TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY") is None: | ||
| os.environ["TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY"] = "1" |
There was a problem hiding this comment.
Nit: Modifying an environment variable defined by ourselves is not a good idea in the long term. If you have a chance can you move the default heuristic to the consumer of the TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY env var instead of modifying it here?
There was a problem hiding this comment.
moving the heuristic to the consumer side is not straightforward here. The consumer of this env var is the C++ BlockManager at pool-allocation time, which has no knowledge of the transceiver runtime — the transceiver is created after the KV cache manager, and transceiver_runtime only exists in the Python-side config.
Also, we want users to be able to override the heuristic via the env var: the default is only applied when TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY is unset, so an explicit =0 (or =1) always wins.
ZhanruiSunCh
left a comment
ZhanruiSunCh
left a comment
There was a problem hiding this comment.
LGTM for infra part.
|
PR_Github #62056 [ run ] completed with state
|
11 participants
Description
When the Python cache transceiver (
cache_transceiver_config.transceiver_runtime='PYTHON') is used together with the C++ V1 KV cache manager, the KV cache pool should be allocated with fabric memory so KV blocks can be transferred over MNNVL.This PR:
Defaults
TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY=1inKvCacheCreator.__init__when:transceiver_runtime='PYTHON', andKVCacheManagerand its hybrid subclasses;KVCacheManagerV2is excluded).The default is applied before any pool allocation (the C++ env getter caches on first read). An explicit user setting is always respected, and platforms without fabric memory support fall back to standard GPU allocation in C++ with a warning.
Sets
kv_cache_bounce_size_mb: 512in the GPT-OSS disaggregated perf configs (from [None][perf] Validate GPT-OSS transceiver v2 performance #15765) where the ctx/gen KV-head sharding differs (ctx TP1 → gen TP4, no attention DP), which requires the bounce buffer to regroup KV heads during transfer. Configs with matching per-rank KV-head layout (equal TP, or gen with attention DP) are intentionally left unchanged.Test Coverage
transceiver_runtime: PYTHON.PR Checklist
🤖 Generated with Claude Code
Dev Engineer Review
KvCacheCreator._maybe_enable_fabric_memory_for_python_transceiver()intensorrt_llm/_torch/pyexecutor/_util.py, called fromKvCacheCreator.__init__.cache_transceiver_config.transceiver_runtime == "PYTHON") with the C++ V1 KV cache manager (KVCacheManager), the helper defaultsTRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY=1only if the env var isn’t already set; otherwise it leaves configuration unchanged so other/unsupported combinations behave as before.kv_cache_bounce_size_mb: 512under bothworker_config.gen.cache_transceiver_configandworker_config.ctx.cache_transceiver_configfor the listed NIXL variants:tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_1k1k_con64_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yamltests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yamltests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yamltests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yamltests/scripts/perf/disaggregated/gb200_stress-gpt-oss-120b-fp4_8k1k_ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yamlQA Engineer Review
tests/scripts/to addkv_cache_bounce_size_mb: 512in thegenandctxcache_transceiver_configblocks for the files listed in the Dev Engineer Review section.tests/integration/test_lists/waives.txtby removing threeperf/test_perf_sanity.py::test_e2ewaiver entries fordisagg_uploadtargetinggb300_glm-5-fp4_8k1k:disagg_upload-e2eforcon1024/ctx1/dep2/gen1/dep8/eplb256/mtp1disagg_upload-gen_onlyforcon1024/ctx1/dep2/gen1/dep32disagg_upload-gen_onlyforcon512/ctx1/dep2/gen1/dep32(the variant described in the waiver removal list, including theeplb0/mtp3portion)