[None][perf] enable NCCL user-buffer registration for VisualGen collectives - #17006
[None][perf] enable NCCL user-buffer registration for VisualGen collectives#17006daichu-nv wants to merge 1 commit into
Conversation
|
@luyiyun1021 let me also copy your comments from the original PR to this PR as I'll be addressing your comments here later.
|
📝 WalkthroughWalkthroughVisualGen adds an ChangesVisualGen NCCL User-Buffer Registration
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Configuration
participant DiffusionWorker
participant NCCL
participant VisualGen
Configuration->>DiffusionWorker: set nccl_buffer_reg
DiffusionWorker->>NCCL: enable CUMEM before init_process_group
DiffusionWorker->>NCCL: initialize distributed communication
VisualGen->>NCCL: execute configured collectives
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsGit: Failed to clone repository. Please run the Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (7)
tensorrt_llm/_torch/visual_gen/nccl_ub_reg.py (4)
129-135: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRedundant
import ctypesshadows the module-level import.
ctypesis already imported at Line 26; the localimport ctypeshere is unnecessary and shadows the outer name.As per coding guidelines, "Avoid shadowing outer variables."🧹 Proposed fix
if hasattr(backend, "_comm"): - import ctypes cap = backend._comm🤖 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/visual_gen/nccl_ub_reg.py` around lines 129 - 135, Remove the redundant local import ctypes from the pointer retrieval block, and reuse the existing module-level ctypes import when configuring PyCapsule_GetPointer in the surrounding function.Source: Coding guidelines
209-236: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing return/parameter annotations on several methods.
register_all'stensorsparameter (Line 209) is unannotated, andderegister(Line 213),deregister_all(Line 222), and__del__(Line 232) lack-> Nonereturn annotations.As per coding guidelines, "Annotate every function, use `None` for non-returning functions."🧹 Proposed fix
- def register_all(self, tensors) -> int: + def register_all(self, tensors: list[torch.Tensor]) -> int: """Register a list of tensors. Returns the number successfully registered.""" return sum(self.register(t) for t in tensors) - def deregister(self, tensor: torch.Tensor): + def deregister(self, tensor: torch.Tensor) -> None: key = tensor.data_ptr() ... - def deregister_all(self): + def deregister_all(self) -> None: if not self.available or self._lib is None: return ... - def __del__(self): + def __del__(self) -> None: try: self.deregister_all() except Exception: pass🤖 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/visual_gen/nccl_ub_reg.py` around lines 209 - 236, Add the missing type annotations in the registration class: annotate register_all’s tensors parameter with its appropriate tensor collection type, and add -> None to deregister, deregister_all, and __del__. Preserve the existing behavior and register_all’s integer return annotation.Source: Coding guidelines
29-29: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer
X | Noneovertyping.Optional[X].
Optional[...]is used throughout (e.g., Lines 45, 110, 164-166) instead of the built-inX | Noneunion syntax.As per coding guidelines, "prefer built-in generic types and
|."Also applies to: 45-45, 110-110, 164-166
🤖 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/visual_gen/nccl_ub_reg.py` at line 29, Replace the typing.Optional import and all Optional[...] annotations in nccl_ub_reg.py with the equivalent X | None syntax, including the annotations near the referenced lines. Preserve the existing types and annotation behavior while removing the now-unused import.Source: Coding guidelines
137-137: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueBroad
except Exceptionviolates the narrow-exception guideline.Both
_extract_comm_ptr(Line 137) and__del__(Line 235) catch bareException. Given the documented "best-effort, multiple fragile introspection paths" nature of_extract_comm_ptr, and that__del__swallowing errors during interpreter teardown is a common defensive pattern, this is understandable, but it still departs from the coding guideline to catch the narrowest exceptions possible.As per coding guidelines, "Catch the narrowest possible exceptions instead of using broad or bare exception handling."
Also applies to: 232-236
🤖 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/visual_gen/nccl_ub_reg.py` at line 137, Replace broad Exception handlers in _extract_comm_ptr and __del__ with the narrowest exception types raised by their introspection and cleanup operations. Preserve _extract_comm_ptr’s best-effort fallback behavior and __del__’s teardown-safe swallowing of expected cleanup failures, while allowing unexpected exceptions to propagate.Source: Coding guidelines
examples/visual_gen/bench_cold_ring.py (1)
42-42: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHardcoded absolute paths reduce portability.
/workspace/models/Cosmos3-Nanoand/workspace/resultsare dev-machine-specific paths with no CLI override, unlike--tag,--warmup,--iters. Anyone else running this script must edit source to point at their own paths.♻️ Proposed fix
ap.add_argument("--tag", default="") + ap.add_argument("--model-path", default="/workspace/models/Cosmos3-Nano") + ap.add_argument("--out-dir", default="/workspace/results") a = ap.parse_args() ... - vg = VisualGen("/workspace/models/Cosmos3-Nano", args=args) + vg = VisualGen(a.model_path, args=args) ... - out = f"/workspace/results/cold_ring4_{tag}.json" - os.makedirs("/workspace/results", exist_ok=True) + out = f"{a.out_dir}/cold_ring4_{tag}.json" + os.makedirs(a.out_dir, exist_ok=True)Also applies to: 72-73
🤖 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 `@examples/visual_gen/bench_cold_ring.py` at line 42, Update the VisualGen setup and result-output handling in the benchmark entrypoint to replace the hardcoded model and results paths with configurable CLI arguments, following the existing args pattern used by --tag, --warmup, and --iters. Provide sensible defaults while allowing callers to override both paths without editing source.examples/visual_gen/bench_nccl_ub.py (2)
61-71: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove
import statisticsto module scope.Importing
statisticsinside_statson every call is unconventional; hoist it to the top-level imports alongsidejson/os/time.🤖 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 `@examples/visual_gen/bench_nccl_ub.py` around lines 61 - 71, Move the statistics import from inside _stats to the module-level import section alongside json, os, and time, while leaving the function’s statistical calculations unchanged.
30-71: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer built-in generics (
dict,list) overtyping.Dict/List.
from typing import Dict, List(line 34) and the resultingDict/List[float]annotations throughout (_timed_iters,_stats,_bench_collective_one, etc.) could use PEP 585 built-in generics, consistent with the repo's Python 3.10+ target.As per coding guidelines, "prefer built-in generic types and
|."🤖 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 `@examples/visual_gen/bench_nccl_ub.py` around lines 30 - 71, Replace the typing.Dict and typing.List imports and all corresponding annotations, including _timed_iters, _stats, and _bench_collective_one, with Python 3.10 built-in dict and list generics while preserving the existing type parameters.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.
Inline comments:
In `@examples/visual_gen/bench_cold_ring.py`:
- Around line 47-60: Ensure benchmark cleanup runs when generation fails by
placing the warmup and iteration loops in bench_cold_ring.py within a
try/finally and keeping vg.shutdown() in the finally block. Apply the same
structure in _bench_pipeline_one in examples/visual_gen/bench_nccl_ub.py, with
vg.shutdown() guaranteed in finally; preserve the existing benchmark behavior
otherwise.
- Around line 1-3: Add the standard NVIDIA SPDX copyright header at the
beginning of the new file, matching the header format and current copyright year
used by the sibling bench_nccl_ub.py file, while preserving the existing shebang
and module docstring.
- Line 3: Update the imports at the top of the module to use separate import
statements for os, statistics, time, json, and argparse, ordered according to
the repository’s configured linter conventions.
In `@tensorrt_llm/_torch/visual_gen/nccl_ub_reg.py`:
- Around line 188-231: Update the registrar’s self._handles entries used by
register() and deregister() to retain a strong reference to each registered
tensor alongside its NCCL handle, preventing the address from being reused while
registration remains active. Adjust existing-handle checks, NCCL deregistration,
and deregister_all() to unpack and use the stored handle while preserving
successful registration counting and cleanup behavior.
- Line 1: Update the CUMEM NCCL version requirement consistently in
nccl_ub_reg.py and the related configuration description in args.py: lower the
documented and enforced minimum to NCCL 2.20 to match the required deadlock fix,
or explicitly document 2.21 as an intentional safety margin if retaining the
current gate.
---
Nitpick comments:
In `@examples/visual_gen/bench_cold_ring.py`:
- Line 42: Update the VisualGen setup and result-output handling in the
benchmark entrypoint to replace the hardcoded model and results paths with
configurable CLI arguments, following the existing args pattern used by --tag,
--warmup, and --iters. Provide sensible defaults while allowing callers to
override both paths without editing source.
In `@examples/visual_gen/bench_nccl_ub.py`:
- Around line 61-71: Move the statistics import from inside _stats to the
module-level import section alongside json, os, and time, while leaving the
function’s statistical calculations unchanged.
- Around line 30-71: Replace the typing.Dict and typing.List imports and all
corresponding annotations, including _timed_iters, _stats, and
_bench_collective_one, with Python 3.10 built-in dict and list generics while
preserving the existing type parameters.
In `@tensorrt_llm/_torch/visual_gen/nccl_ub_reg.py`:
- Around line 129-135: Remove the redundant local import ctypes from the pointer
retrieval block, and reuse the existing module-level ctypes import when
configuring PyCapsule_GetPointer in the surrounding function.
- Around line 209-236: Add the missing type annotations in the registration
class: annotate register_all’s tensors parameter with its appropriate tensor
collection type, and add -> None to deregister, deregister_all, and __del__.
Preserve the existing behavior and register_all’s integer return annotation.
- Line 29: Replace the typing.Optional import and all Optional[...] annotations
in nccl_ub_reg.py with the equivalent X | None syntax, including the annotations
near the referenced lines. Preserve the existing types and annotation behavior
while removing the now-unused import.
- Line 137: Replace broad Exception handlers in _extract_comm_ptr and __del__
with the narrowest exception types raised by their introspection and cleanup
operations. Preserve _extract_comm_ptr’s best-effort fallback behavior and
__del__’s teardown-safe swallowing of expected cleanup failures, while allowing
unexpected exceptions to propagate.
🪄 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: bee095fd-a4a3-4873-a53e-d51fb6c801a6
📒 Files selected for processing (8)
docs/source/models/visual-generation.mdexamples/visual_gen/bench_cold_ring.pyexamples/visual_gen/bench_nccl_ub.pyexamples/visual_gen/serve/configs/flux1_nccl_ub.ymlexamples/visual_gen/serve/configs/wan21_nccl_ub.ymltensorrt_llm/_torch/visual_gen/executor.pytensorrt_llm/_torch/visual_gen/nccl_ub_reg.pytensorrt_llm/visual_gen/args.py
| #!/usr/bin/env python3 | ||
| """Cold-start ring-4 test — each run is a fresh process with no prior ring-4 calls.""" | ||
| import os, statistics, time, json, argparse |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Missing NVIDIA copyright header.
This new file has no SPDX/copyright header, unlike its sibling bench_nccl_ub.py (lines 2-3), which includes SPDX-FileCopyrightText and SPDX-License-Identifier.
As per coding guidelines, "New files must include the NVIDIA copyright header; modified files must have the copyright year updated."
📄 Proposed fix
#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
"""Cold-start ring-4 test — each run is a fresh process with no prior ring-4 calls."""📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| #!/usr/bin/env python3 | |
| """Cold-start ring-4 test — each run is a fresh process with no prior ring-4 calls.""" | |
| import os, statistics, time, json, argparse | |
| #!/usr/bin/env python3 | |
| # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| """Cold-start ring-4 test — each run is a fresh process with no prior ring-4 calls.""" | |
| import os, statistics, time, json, argparse |
🤖 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 `@examples/visual_gen/bench_cold_ring.py` around lines 1 - 3, Add the standard
NVIDIA SPDX copyright header at the beginning of the new file, matching the
header format and current copyright year used by the sibling bench_nccl_ub.py
file, while preserving the existing shebang and module docstring.
Source: Coding guidelines
| @@ -0,0 +1,79 @@ | |||
| #!/usr/bin/env python3 | |||
| """Cold-start ring-4 test — each run is a fresh process with no prior ring-4 calls.""" | |||
| import os, statistics, time, json, argparse | |||
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Combine imports onto one line — flagged by pycodestyle E401.
import os, statistics, time, json, argparse places multiple modules on a single import statement, which pycodestyle's E401 rule flags. Ruff in this repo enables the E rule set.
As per coding guidelines, "Order imports according to the configured linter."
🔧 Proposed fix
-import os, statistics, time, json, argparse
+import argparse
+import json
+import os
+import statistics
+import time📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| import os, statistics, time, json, argparse | |
| import argparse | |
| import json | |
| import os | |
| import statistics | |
| import time |
🤖 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 `@examples/visual_gen/bench_cold_ring.py` at line 3, Update the imports at the
top of the module to use separate import statements for os, statistics, time,
json, and argparse, ordered according to the repository’s configured linter
conventions.
Source: Coding guidelines
| for i in range(a.warmup): | ||
| print(f" warmup {i+1}/{a.warmup}...") | ||
| vg.generate(inputs=PROMPTS[i % len(PROMPTS)], params=VisualGenParams(**gen_kwargs)) | ||
| print(f" warmup {i+1} done") | ||
|
|
||
| times = [] | ||
| for i in range(a.iters): | ||
| t0 = time.perf_counter() | ||
| vg.generate(inputs=PROMPTS[i % len(PROMPTS)], params=VisualGenParams(**gen_kwargs)) | ||
| elapsed = time.perf_counter() - t0 | ||
| times.append(elapsed) | ||
| print(f" iter {i+1}/{a.iters}: {elapsed:.2f}s") | ||
|
|
||
| vg.shutdown() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Worker process/GPU memory can leak if generate() raises before shutdown(). Both benchmark scripts call vg.generate() in a plain loop and only reach vg.shutdown() afterward; if any iteration raises, the executor's worker process(es) and GPU memory are never released, requiring manual cleanup before the next run on a shared node.
examples/visual_gen/bench_cold_ring.py#L47-L60: wrap the warmup/iteration loops andvg.shutdown()(line 60) in atry/finallysoshutdown()always runs.examples/visual_gen/bench_nccl_ub.py#L172-L182: wrap the warmup/iteration loops in_bench_pipeline_oneandvg.shutdown()(line 181) in atry/finallysoshutdown()always runs.
📍 Affects 2 files
examples/visual_gen/bench_cold_ring.py#L47-L60(this comment)examples/visual_gen/bench_nccl_ub.py#L172-L182
🤖 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 `@examples/visual_gen/bench_cold_ring.py` around lines 47 - 60, Ensure
benchmark cleanup runs when generation fails by placing the warmup and iteration
loops in bench_cold_ring.py within a try/finally and keeping vg.shutdown() in
the finally block. Apply the same structure in _bench_pipeline_one in
examples/visual_gen/bench_nccl_ub.py, with vg.shutdown() guaranteed in finally;
preserve the existing benchmark behavior otherwise.
| @@ -0,0 +1,236 @@ | |||
| # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
When was NCCL_CUMEM_ENABLE introduced in NCCL, and which NCCL version fixed CUMEM-related hangs?
💡 Result:
The environment variable NCCL_CUMEM_ENABLE was introduced in NCCL 2.18 [1]. While it was disabled by default in that version, it became auto-enabled by default starting in NCCL 2.19 if the system supports it [1]. Regarding CUMEM-related hangs, users encountered deadlock scenarios during communicator connection establishment (specifically reported in Alltoall operations) when NCCL_CUMEM_ENABLE was enabled (the default) in NCCL 2.19 [2][3]. Setting NCCL_CUMEM_ENABLE=0 was the recommended workaround for this issue [2][3]. The fix for these deadlocks was included in NCCL 2.20 [3].
Citations:
- 1: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html
- 2: https://docs.nvidia.com/deeplearning/nccl/release-notes/rel_2-19-3.html
- 3: NCCL alltoall_perf hangs via PXN in our non-superpod network nccl#1116
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n-- nccl_ub_reg.py --\n'
sed -n '1,140p' tensorrt_llm/_torch/visual_gen/nccl_ub_reg.py
printf '\n-- args.py excerpt --\n'
sed -n '220,260p' tensorrt_llm/visual_gen/args.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1924
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
for path in [
Path("tensorrt_llm/_torch/visual_gen/nccl_ub_reg.py"),
Path("tensorrt_llm/visual_gen/args.py"),
]:
print(f"\n== {path} ==")
for i, line in enumerate(path.read_text().splitlines(), start=1):
if path.name == "nccl_ub_reg.py" and i <= 130:
print(f"{i:4d}: {line}")
elif path.name == "args.py" and 225 <= i <= 255:
print(f"{i:4d}: {line}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 1924
Align the CUMEM floor with the actual NCCL support level
tensorrt_llm/_torch/visual_gen/nccl_ub_reg.py and tensorrt_llm/visual_gen/args.py both say CUMEM needs NCCL ≥ 2.21, but NCCL_CUMEM_ENABLE exists since 2.18, is auto-enabled from 2.19, and the deadlock fix landed in 2.20. If 2.21 is just a safety margin, say so explicitly; otherwise lower the gate/description to 2.20.
🤖 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/visual_gen/nccl_ub_reg.py` at line 1, Update the CUMEM
NCCL version requirement consistently in nccl_ub_reg.py and the related
configuration description in args.py: lower the documented and enforced minimum
to NCCL 2.20 to match the required deadlock fix, or explicitly document 2.21 as
an intentional safety margin if retaining the current gate.
| def register(self, tensor: torch.Tensor) -> bool: | ||
| """Register a single tensor. Returns True on success.""" | ||
| if not self.available: | ||
| return False | ||
| key = tensor.data_ptr() | ||
| if key in self._handles: | ||
| return True # already registered | ||
| handle = ctypes.c_void_p(None) | ||
| ret = self._lib.ncclCommRegister( | ||
| ctypes.c_void_p(self._comm_ptr), | ||
| ctypes.c_void_p(key), | ||
| ctypes.c_size_t(tensor.nbytes), | ||
| ctypes.byref(handle), | ||
| ) | ||
| if ret == 0: # ncclSuccess | ||
| self._handles[key] = handle | ||
| logger.debug(f"Registered tensor @{key:#x} ({tensor.nbytes} B)") | ||
| return True | ||
| logger.warning(f"ncclCommRegister returned error {ret} for tensor @{key:#x}") | ||
| return False | ||
|
|
||
| def register_all(self, tensors) -> int: | ||
| """Register a list of tensors. Returns the number successfully registered.""" | ||
| return sum(self.register(t) for t in tensors) | ||
|
|
||
| def deregister(self, tensor: torch.Tensor): | ||
| key = tensor.data_ptr() | ||
| if key in self._handles: | ||
| self._lib.ncclCommDeregister( | ||
| ctypes.c_void_p(self._comm_ptr), | ||
| self._handles.pop(key), | ||
| ) | ||
| logger.debug(f"Deregistered tensor @{key:#x}") | ||
|
|
||
| def deregister_all(self): | ||
| if not self.available or self._lib is None: | ||
| return | ||
| for key, handle in list(self._handles.items()): | ||
| self._lib.ncclCommDeregister( | ||
| ctypes.c_void_p(self._comm_ptr), | ||
| handle, | ||
| ) | ||
| self._handles.clear() | ||
|
|
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Registrar dict keyed by raw data_ptr() can go stale if a tensor is freed and its address reused.
register()/deregister() key self._handles purely by tensor.data_ptr(). If a registered tensor is freed (without deregister() being called) and PyTorch's allocator later reuses the same address for a different, differently-sized tensor, register() on the new tensor will see key in self._handles and silently skip re-registration — leaving NCCL's registration for that address bound to the old size/tensor. The registrar holds no reference to the tensor to keep it alive, and there's no size check to catch a mismatch. The documented intended usage (persistent latent buffer reused every step) largely avoids this, but nothing enforces it.
🛡️ Proposed mitigation
- key = tensor.data_ptr()
- if key in self._handles:
- return True # already registered
+ key = tensor.data_ptr()
+ if key in self._handles:
+ cached_tensor, _ = self._handles[key]
+ if cached_tensor.nbytes != tensor.nbytes:
+ logger.warning(f"Address {key:`#x`} reused with a different size; re-registering")
+ else:
+ return True # already registeredAlso consider storing (handle, tensor) (keeping a strong reference) in self._handles to prevent premature reuse of the address while registered.
🤖 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/visual_gen/nccl_ub_reg.py` around lines 188 - 231, Update
the registrar’s self._handles entries used by register() and deregister() to
retain a strong reference to each registered tensor alongside its NCCL handle,
preventing the address from being reused while registration remains active.
Adjust existing-handle checks, NCCL deregistration, and deregister_all() to
unpack and use the stored handle while preserving successful registration
counting and cleanup behavior.
d2f6cdd to
3e85c19
Compare
|
Pushed a rebase/conflict-resolution update on top of current Follow-up on the copied CUMEM vs. explicit-registration comments: this PR currently wires |
3e85c19 to
f3147b7
Compare
|
Update after addressing the reviewer comments and CI:
Local changed-file pre-commit passed for the PR files. |
…ctives Signed-off-by: Daisy Chu <daichu@nvidia.com>
f3147b7 to
e642d69
Compare
2 participants
This PR copies #15013 because I do not have write permission to the original PR branch.
What this changes
Adds opt-in NCCL user-buffer registration to the VisualGen parallel inference path via
ParallelConfig.nccl_buffer_reg. Pure TRT-LLM change — activates a feature NCCL already exposes.Files changed
Benchmark Results
System: 8× NVIDIA B200 SXM 183 GB · NCCL 2.29.2 · TRT-LLM 1.3.0rc17
Models: FLUX.1-dev (1024×1024, 20 steps) and Cosmos3-Nano (720×1280, 57 frames, 20 steps)
Micro-benchmark (collective latency, 4-GPU, 50 iters)
CUMEM improves individual collective latency at 4 GPU:
E2E Ulysses parallelism — CUMEM effect
CUMEM has negligible E2E impact (±2%, within noise). Communication is not the bottleneck on B200 — the collective-level improvements are real but diluted by compute time.
Ring attention cold-start (each config run as fresh isolated process)
Both CUMEM modes show the same cold-start profile and identical steady-state:
The cold spike is CUDA kernel JIT compilation (CUTEDSL compiles ring attention kernels on first use for each shape config), not an NCCL artifact. CUMEM does not eliminate cold-start.
Ring attention is not competitive with Ulysses for either model: ring-4 Cosmos steady-state is ~14.15 s vs Ulysses 4-GPU 6.77 s.
Conclusion
nccl_buffer_reg: trueis safe to enable as a default. It may benefit configurations with higher communication-to-compute ratios (longer sequences, more layers, lower-end hardware) not tested here.Test plan
python -m tensorrt_llm.visual_gen ...(default) — no regressionnccl_buffer_reg: true— CUMEM flag visible in NCCL logsnccl_buffer_reg: false(default) path unchanged — no CUMEM env var set🤖 Generated with Claude Code
Summary by CodeRabbit