[None][perf] enable NCCL user-buffer registration for VisualGen collectives - #15013
Closed
pkisfaludi-nv wants to merge 10 commits into
Closed
[None][perf] enable NCCL user-buffer registration for VisualGen collectives#15013pkisfaludi-nv wants to merge 10 commits into
pkisfaludi-nv wants to merge 10 commits into
Conversation
Contributor
📝 WalkthroughWalkthroughThis PR introduces HunyuanDiT, a bilingual text-to-image diffusion model with Ulysses sequence parallelism, alongside NCCL user-buffer registration (CUMEM) for optimized distributed collectives. It includes benchmark tooling, example configurations, and comprehensive documentation of both features. ChangesHunyuanDiT and NCCL User-Buffer Registration
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ 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: 3
🧹 Nitpick comments (4)
tensorrt_llm/_torch/visual_gen/models/hunyuandit/pipeline_hunyuandit.py (1)
507-507: 💤 Low valueRename unused loop variable
ito_i.The loop counter is not used within the denoise loop body.
Proposed fix
- for i, t in enumerate(timesteps): + for _i, t in enumerate(timesteps):🤖 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/models/hunyuandit/pipeline_hunyuandit.py` at line 507, Rename the unused loop counter in the denoising loop from i to _i to signal it's intentionally unused; change the loop header "for i, t in enumerate(timesteps):" to "for _i, t in enumerate(timesteps):" (the loop over timesteps in pipeline_hunyuandit.py) so linters won't flag the unused variable.examples/visual_gen/nccl_ub_benchmark_report.md (1)
69-78: 💤 Low valueAdd language specifier to satisfy markdownlint.
The fenced code block at line 69 lacks a language specifier. Adding
textwould satisfy the linter while maintaining readability.✨ Proposed fix
-``` +```text 2 GPU 4 GPU 8 GPU Collective Base CUMEM Base CUMEM Base CUMEM Best speedup🤖 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/nccl_ub_benchmark_report.md` around lines 69 - 78, The fenced code block that begins with ``` in nccl_ub_benchmark_report.md lacks a language specifier; update that opening fence to ```text (i.e., add the word "text" immediately after the three backticks) so the markdownlint rule is satisfied while preserving the table formatting.tensorrt_llm/_torch/visual_gen/nccl_ub_reg.py (2)
129-129: ⚡ Quick winRemove redundant
ctypesimport.
ctypesis already imported at the module level (line 26), so this local import inside_extract_comm_ptris unnecessary.♻️ Proposed fix
# Older releases: the comm is stored in backend._comm (a capsule). # We can cast the capsule's void* with ctypes. if hasattr(backend, "_comm"): - import ctypes cap = backend._comm # PyCapsule_GetPointer is exposed through ctypes.pythonapi ctypes.pythonapi.PyCapsule_GetPointer.restype = ctypes.c_void_p🤖 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 129, Remove the redundant local import of ctypes inside the _extract_comm_ptr function: delete the "import ctypes" statement (ctypes is already imported at module level), leaving the function to reference the existing module-level ctypes import; ensure no other references to a local ctypes name remain in _extract_comm_ptr.
213-220: ⚡ Quick winAdd availability check in
deregister()for consistency.While in practice
deregister()is only called for keys that exist in_handles(which can only be populated whenavailable=True), adding an explicit availability check mirrors the pattern inderegister_all()(line 223) and prevents potentialAttributeErrorif the method is called incorrectly.🛡️ Proposed fix
def deregister(self, tensor: torch.Tensor): + if not self.available or self._lib is None: + return key = tensor.data_ptr() if key in self._handles: self._lib.ncclCommDeregister(🤖 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 213 - 220, The deregister() method currently assumes the registry is available and accesses self._handles directly; add the same availability guard used in deregister_all() to avoid AttributeError when called incorrectly. In deregister(), first check if self.available (or return early if not), then proceed to compute key and pop from self._handles and call self._lib.ncclCommDeregister with self._comm_ptr; this mirrors deregister_all() and ensures safe access to _handles, _lib and _comm_ptr.
🤖 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_nccl_ub.py`:
- Around line 129-138: The comment above the configs list is misleading: update
the comment to reflect that each entry in configs is a 3-tuple (name,
collective, shape) rather than a 6-element tuple; locate the configs variable
and the unpacking that uses (name, collective, shape) and change the explanatory
comment to accurately describe the tuple structure and the meaning of the shape
elements (e.g., shape = (B, S/world_size, H/world_size, D) or shape = (B, S, H)
for the small case) so it matches the actual data and unpacking.
- Around line 160-181: In _bench_pipeline_one fix the VisualGen.generate call to
match its signature: pass the input text as the first positional argument
(inputs) and construct a VisualGenParams-like object for inference settings
(height, width, num_inference_steps) as the optional second parameter instead of
using keyword args; update both the warmup loop and timed call (both vg.generate
invocations) to pass (inputs, params) so VisualGen.generate is called with the
expected parameters.
In `@tensorrt_llm/_torch/visual_gen/models/hunyuandit/transformer_hunyuandit.py`:
- Around line 272-284: Add a runtime validation after computing S =
hidden_states.shape[1] to ensure S is divisible by self._ulysses_size: check if
S % ulysses_size != 0 and raise a ValueError with a clear message if so; this
prevents incorrect shard_size computation and subsequent incorrect all_gather
results in the Ulysses sharding code that slices hidden_states by rank *
shard_size:(rank + 1) * shard_size (and similarly for image_rotary_emb
freqs_cos/freqs_sin). Ensure the check references S and ulysses_size
(self._ulysses_size) and runs before computing shard_size or slicing
hidden_states.
---
Nitpick comments:
In `@examples/visual_gen/nccl_ub_benchmark_report.md`:
- Around line 69-78: The fenced code block that begins with ``` in
nccl_ub_benchmark_report.md lacks a language specifier; update that opening
fence to ```text (i.e., add the word "text" immediately after the three
backticks) so the markdownlint rule is satisfied while preserving the table
formatting.
In `@tensorrt_llm/_torch/visual_gen/models/hunyuandit/pipeline_hunyuandit.py`:
- Line 507: Rename the unused loop counter in the denoising loop from i to _i to
signal it's intentionally unused; change the loop header "for i, t in
enumerate(timesteps):" to "for _i, t in enumerate(timesteps):" (the loop over
timesteps in pipeline_hunyuandit.py) so linters won't flag the unused variable.
In `@tensorrt_llm/_torch/visual_gen/nccl_ub_reg.py`:
- Line 129: Remove the redundant local import of ctypes inside the
_extract_comm_ptr function: delete the "import ctypes" statement (ctypes is
already imported at module level), leaving the function to reference the
existing module-level ctypes import; ensure no other references to a local
ctypes name remain in _extract_comm_ptr.
- Around line 213-220: The deregister() method currently assumes the registry is
available and accesses self._handles directly; add the same availability guard
used in deregister_all() to avoid AttributeError when called incorrectly. In
deregister(), first check if self.available (or return early if not), then
proceed to compute key and pop from self._handles and call
self._lib.ncclCommDeregister with self._comm_ptr; this mirrors deregister_all()
and ensures safe access to _handles, _lib and _comm_ptr.
🪄 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: 0dc180d6-13a0-4db1-a8e2-d2559487e37f
📒 Files selected for processing (21)
docs/source/models/visual-generation.mdexamples/visual_gen/bench_nccl_ub.pyexamples/visual_gen/nccl_ub_benchmark_report.mdexamples/visual_gen/nccl_ub_results/baseline_2gpu.jsonexamples/visual_gen/nccl_ub_results/baseline_4gpu.jsonexamples/visual_gen/nccl_ub_results/baseline_8gpu.jsonexamples/visual_gen/nccl_ub_results/ub_2gpu.jsonexamples/visual_gen/nccl_ub_results/ub_4gpu.jsonexamples/visual_gen/nccl_ub_results/ub_8gpu.jsonexamples/visual_gen/serve/configs/flux1_nccl_ub.ymlexamples/visual_gen/serve/configs/hunyuandit.ymlexamples/visual_gen/serve/configs/wan21_nccl_ub.ymltensorrt_llm/_torch/visual_gen/executor.pytensorrt_llm/_torch/visual_gen/models/__init__.pytensorrt_llm/_torch/visual_gen/models/hunyuandit/__init__.pytensorrt_llm/_torch/visual_gen/models/hunyuandit/defaults.pytensorrt_llm/_torch/visual_gen/models/hunyuandit/pipeline_hunyuandit.pytensorrt_llm/_torch/visual_gen/models/hunyuandit/transformer_hunyuandit.pytensorrt_llm/_torch/visual_gen/nccl_ub_reg.pytensorrt_llm/_torch/visual_gen/pipeline_registry.pytensorrt_llm/visual_gen/args.py
pkisfaludi-nv
changed the title
venkywonka
reviewed
Jun 9, 2026
| | `Qwen/Qwen-Image-2512` | Text-to-Image | | ||
| | `nvidia/Cosmos3-Nano` | Text-to-Image, Text-to-Video, Image-to-Video | | ||
| | `nvidia/Cosmos3-Super` | Text-to-Image, Text-to-Video, Image-to-Video | | ||
| | `Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers` | Text-to-Image | |
Collaborator
There was a problem hiding this comment.
are these supposed to be part of the change?
Author
There was a problem hiding this comment.
you are right they belon to the other changeset in #14991
venkywonka
reviewed
Jun 9, 2026
Collaborator
There was a problem hiding this comment.
seems like this is a entire new model addition PR disguised as a [None][perf] enable NCCL user-buffer registration for VisualGen collectives?
Author
pkisfaludi-nv
added a commit
to pkisfaludi-nv/TensorRT-LLM
that referenced
this pull request
Jun 30, 2026
…lts from PR PR NVIDIA#15013 should only cover NCCL user-buffer registration. Move all HunyuanDiT model code to PR NVIDIA#14991 (feat/hunyuandit-visual-gen) and drop benchmark result files that are not part of the feature itself. Removed: - tensorrt_llm/_torch/visual_gen/models/hunyuandit/ (4 files) - examples/visual_gen/serve/configs/hunyuandit.yml - examples/visual_gen/nccl_ub_results/ (16 JSON result files) - examples/visual_gen/nccl_ub_benchmark_report.md Updated: - models/__init__.py: drop HunyuanDiTPipeline import + __all__ entry - pipeline_registry.py: drop HunyuanDiT detection branch - docs/source/models/visual-generation.md: remove HunyuanDiT rows and [^3] footnote from supported-models table and feature matrix Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Integrates Tencent HunyuanDiT into the TensorRT-LLM VisualGen framework, following the same BasePipeline pattern used by Qwen-Image and FLUX. New files: - tensorrt_llm/_torch/visual_gen/models/hunyuandit/pipeline_hunyuandit.py HunyuanDiTPipeline registered via @register_pipeline; supports v1.0-v1.2 checkpoints. Implements bilingual (BertModel + MT5EncoderModel) text encoding, DDPM denoising loop, resolution binning to training buckets, 2D RoPE embeddings, and VAE decode. - tensorrt_llm/_torch/visual_gen/models/hunyuandit/transformer_hunyuandit.py HunyuanDiT2DModelWrapper: thin nn.Module around diffusers HunyuanDiT2DModel with load_weights() compatible with the WeightLoader contract. - tensorrt_llm/_torch/visual_gen/models/hunyuandit/defaults.py Default generation params (1024×1024, 50 steps, cfg=7.5) and extra-param schema (negative_prompt, use_resolution_binning). - examples/visual_gen/serve/configs/hunyuandit.yml Serve config with VANILLA attention backend. Modified: - pipeline_registry.py: add HunyuanDiT detection in _detect_from_checkpoint - models/__init__.py: export HunyuanDiTPipeline - docs/source/models/visual-generation.md: add HunyuanDiT to model table and feature matrix Qwen-Image (Qwen/Qwen-Image, Qwen/Qwen-Image-2512) was already present in main; no code changes needed for it — confirmed in docs table. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Peter Kisfaludi <pkisfaludi@nvidia.com>
Implements DeepSpeed Ulysses sequence parallelism for HunyuanDiT
following the same pattern used by FLUX (SequenceSharder + all-to-all
around attention).
HunyuanDiTUlyssesAttnProcessor
Drop-in replacement for HunyuanAttnProcessor2_0. For self-attention
(encoder_hidden_states is None) it wraps F.scaled_dot_product_attention
with two all_to_all_4d calls:
[B, S/U, H, D] → all-to-all → [B, S, H/U, D] → SDPA → all-to-all → [B, S/U, H, D]
Cross-attention falls back to standard SDPA (text K/V is replicated on
every rank, so no all-to-all is needed).
HunyuanDiT2DModelUlysses
Monkeypatches HunyuanDiT2DModel.forward() via types.MethodType to:
1. Shard the patch-embedded latent sequence across Ulysses ranks.
2. Slice image_rotary_emb to the local sequence shard.
3. Run the transformer blocks (with custom processors for self-attn).
4. all_gather to reassemble the full sequence before norm_out/proj_out.
U-Net-style skip tensors are sharded the same way as hidden_states so
no special handling is needed for the skip connections.
Constraints
- num_attention_heads (16) must be divisible by ulysses_size.
- Validated at wrapper construction; raises ValueError otherwise.
- Requires vgm.ulysses_group to be initialised (VisualGenMapping.init_device_mesh).
Docs: HunyuanDiT Ulysses column updated from No → Yes in the feature matrix.
Qwen-Image already has Ulysses support through the TRT-LLM Attention module
(inherits UlyssesAttention wrapping); no code changes needed there.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Peter Kisfaludi <pkisfaludi@nvidia.com>
…lectives
Adds two complementary NCCL buffer-registration mechanisms to the VisualGen
parallel inference path:
1. CUMEM mode (NCCL_CUMEM_ENABLE=1): set before init_process_group so NCCL
allocates its scratch via cuMemCreate (VMM). Enables NVLS zero-copy on
NVLink systems; benefits all collectives automatically. Requires NCCL ≥ 2.21.
2. Explicit ncclCommRegister (NCCLBufferRegistrar): best-effort per-tensor
registration via ctypes, extracting ncclComm_t from the PyTorch process
group. Falls back gracefully if NCCL < 2.19 or the comm ptr cannot be
extracted. Requires NCCL ≥ 2.19.
Both are activated by setting nccl_buffer_reg: true in parallel_config.
The flag is threaded through executor.py, which calls enable_nccl_cumem()
before torch.distributed.init_process_group() when the flag is set.
New files:
- tensorrt_llm/_torch/visual_gen/nccl_ub_reg.py — registration utilities
- examples/visual_gen/serve/configs/flux1_nccl_ub.yml — FLUX config (ulysses=2 + UB)
- examples/visual_gen/serve/configs/wan21_nccl_ub.yml — Wan config (ulysses=2 + UB)
- examples/visual_gen/bench_nccl_ub.py — two-tier benchmark script:
tier=micro : collective latency without model weights (torchrun --nproc-per-node=N)
tier=pipeline: end-to-end generation timing (requires checkpoint paths)
Ref: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/bufferreg.html
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Peter Kisfaludi <pkisfaludi@nvidia.com>
Collective micro-benchmark (no model weights) run on umb-b200-138 (8× NVIDIA B200, NCCL 2.29.2) at world_size=2/4/8. Key findings: - 4-GPU Ulysses: CUMEM mode reduces flux all-to-all −13%, Wan all-to-all −11%, small all-reduce −19% (median latency) - 2-GPU: near-zero delta; 8-GPU: mixed (Wan all-to-all −12%, FLUX flat) - Recommendation: enable nccl_buffer_reg for 4-GPU deployments Raw JSON in nccl_ub_results/. Pipeline tier (end-to-end) pending model access. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Peter Kisfaludi <pkisfaludi@nvidia.com>
…s and parallelism modes Full 1/2/4/8-GPU E2E sweep for FLUX.1-dev and Cosmos3-Nano with Ulysses and Ring attention, measured with and without NCCL_CUMEM_ENABLE on 8×B200 SXM node. Key findings: - Ulysses 4-GPU achieves 1.95× (FLUX) and 2.06× (Cosmos) speedup over single GPU - CUMEM is negligible for Ulysses (±2%, noise floor) on both models - Ring attention is counter-productive for FLUX at 1024×1024 at all GPU counts - CUMEM is critical for video ring attention: Cosmos ring-4 drops from 32.45s to 14.52s (+55%) — without it ring-4 is 2.3× slower than single GPU Updated report combines micro-benchmark collective latency tables (from prior commit) with E2E pipeline timing tables and parallelism strategy recommendations. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Peter Kisfaludi <pkisfaludi@nvidia.com>
- Fix misleading comment: tuple structure is (name, collective, shape), not the old (name, B, S/world_size, H/world_size, D, collective) - Fix vg.generate() call: use correct signature generate(inputs, params) with VisualGenParams instead of raw keyword args (prompt/height/width/ num_inference_steps), which would raise TypeError at runtime Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Peter Kisfaludi <pkisfaludi@nvidia.com>
…art vs steady state Re-ran ring attention with warmup=2, iters=5 to capture NCCL buffer warming behavior. Finding: the original ring-4 Cosmos "55% improvement" from CUMEM was a cold-start artifact, not a sustained throughput difference. With insufficient warmup (1 iter), timed iterations measured NCCL's one-time staging buffer allocation overhead. Corrected ring-4 Cosmos results (per-iteration, after 2 warmup): no CUMEM: [32.4s, 32.4s, 32.8s, 16.2s, 14.3s] ← settles after 3-4 calls CUMEM: [15.1s, 14.6s, 14.1s, 14.3s, 14.3s] ← stable from first call Steady-state is the same (~14.3s) either way. CUMEM's actual benefit for ring attention is eliminating cold-start spikes on the first few inferences after service startup, not improving throughput. Updated report reflects this corrected interpretation. Ulysses 4-GPU remains the clear recommendation (0.93s FLUX, 6.77s Cosmos vs ring-4 ~14.3s). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Peter Kisfaludi <pkisfaludi@nvidia.com>
…A kernel JIT, not NCCL Ran each ring-4 config in a fresh isolated process (no prior calls to warm kernel caches) to correctly characterize cold-start behavior. Result (Cosmos ring-4, no warmup, 7 iters): no CUMEM: [34.0, 14.6, 14.2, 14.2, 14.2, 14.2, 14.2] CUMEM: [29.5, 15.9, 14.2, 14.2, 14.1, 14.1, 14.2] Both modes have the same cold-start profile and identical steady-state (~14.15s). CUMEM does not eliminate cold-start. The cold spike is CUDA kernel JIT compilation (CUTEDSL compiles ring attention kernels on first use for each shape config). Prior report incorrectly showed "55% improvement" from CUMEM — this was a test artifact where the CUMEM run ran in the same session after the baseline, benefiting from already-compiled kernels. Corrected by running each config as a fresh process. Correct conclusions: - CUMEM: negligible E2E impact for both Ulysses and ring attention - Micro-benchmark: 10-19% improvement on individual collectives at 4 GPU (real, but diluted in E2E because communication is not the bottleneck on B200) - Ulysses 4-GPU remains optimal: 1.95x (FLUX), 2.06x (Cosmos) - Ring attention slower than Ulysses at all tested GPU counts Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Peter Kisfaludi <pkisfaludi@nvidia.com>
…lts from PR PR NVIDIA#15013 should only cover NCCL user-buffer registration. Move all HunyuanDiT model code to PR NVIDIA#14991 (feat/hunyuandit-visual-gen) and drop benchmark result files that are not part of the feature itself. Removed: - tensorrt_llm/_torch/visual_gen/models/hunyuandit/ (4 files) - examples/visual_gen/serve/configs/hunyuandit.yml - examples/visual_gen/nccl_ub_results/ (16 JSON result files) - examples/visual_gen/nccl_ub_benchmark_report.md Updated: - models/__init__.py: drop HunyuanDiTPipeline import + __all__ entry - pipeline_registry.py: drop HunyuanDiT detection branch - docs/source/models/visual-generation.md: remove HunyuanDiT rows and [^3] footnote from supported-models table and feature matrix Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Peter Kisfaludi <pkisfaludi@nvidia.com>
pkisfaludi-nv
force-pushed
the
feat/nccl-buffer-reg
branch
from
July 1, 2026 02:01
f42f4a6 to
48c2b90
Compare
Signed-off-by: Peter Kisfaludi <pkisfaludi@nvidia.com>
Author
|
Hi @NVIDIA/tensorrt-llm-maintainers — could a maintainer please approve the GitHub Actions workflows for this PR? The |
Collaborator
|
Hi, I am interested in this feature. I tested locally with |
Collaborator
|
Also I‘ve tested on all-gather and it shows no perf gain on it. Do you know why would it not work on all-gather? |
Collaborator
|
/bot run |
4 tasks
|
Superseded by #17006. The replacement PR copies this change because I do not have write permission to this PR's head branch. |
Collaborator
|
close this PR in favor of #17006. |
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
6 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.
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.```yaml
parallel_config:
ulysses_size: 4
nccl_buffer_reg: true # sets NCCL_CUMEM_ENABLE=1 before init_process_group
```
Files changed
```
tensorrt_llm/_torch/visual_gen/nccl_ub_reg.py (new) — registration utilities
tensorrt_llm/_torch/visual_gen/executor.py (modified) — set CUMEM before init_process_group
tensorrt_llm/visual_gen/args.py (modified) — nccl_buffer_reg on ParallelConfig
examples/visual_gen/bench_nccl_ub.py (new) — micro + E2E benchmark script
examples/visual_gen/bench_cold_ring.py (new) — cold-start characterization script
```
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