diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index db8c73912969..276c6795c554 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -51,7 +51,7 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models [^1]: FLUX models use embedded guidance and do not have a separate negative prompt path, so CFG parallelism is not applicable. -[^2]: Qwen-Image ships a native BF16 implementation with per-module numerical parity vs `diffusers.QwenImagePipeline` (cosine >= 0.999 on the full 20B transformer) and `trtllm-serve` / `/v1/images/generations` support. FP8 blockwise and NVFP4 use VisualGen dynamic quantization from BF16 checkpoints; no pre-quantized checkpoint is required. +[^2]: Qwen-Image ships a native BF16 implementation with per-module numerical parity vs `diffusers.QwenImagePipeline` (cosine >= 0.999 on the full 20B transformer) and `trtllm-serve` / `/v1/images/generations` support. FP8 blockwise and NVFP4 use VisualGen dynamic quantization from BF16 checkpoints; no pre-quantized checkpoint is required. ## Quick Start @@ -174,6 +174,51 @@ Configured under `VisualGenArgs.parallel_config`. Modes can be combined: - **Attention2D** (`attn2d_size: [N, M]`): Shards the sequence axis across an `N × M` device mesh (CP degree = `N · M`; total SP degree = `N · M · ulysses_size`). - **Ring Attention** (`ring_size: N`): Shards the sequence axis across a 1D ring of `N` ranks, streaming K/V blocks (CP degree = `N`; total SP degree = `N · ulysses_size`; mutually exclusive with Attention2D). - **Tensor Parallelism** (`tp_size: N`): Splits attention heads and transformer MLPs across GPUs for faster compute and reduced memory usage. + +### NCCL User-Buffer Registration + +Setting `nccl_buffer_reg: true` in `parallel_config` enables NCCL zero-copy collectives for VisualGen's inter-GPU communication (Ulysses all-to-all, ring all-gather, VAE all-reduce). Two mechanisms activate automatically: + +| Mechanism | When active | NCCL requirement | +|---|---|---| +| **CUMEM mode** (`NCCL_CUMEM_ENABLE=1`) | Always when `nccl_buffer_reg: true` | ≥ 2.21 | +| **Explicit `ncclCommRegister`** | When the raw `ncclComm_t` is accessible from the PyTorch process group | ≥ 2.19 | + +**CUMEM mode** instructs NCCL to allocate its internal scratch buffers via `cuMemCreate` (VMM). On NVLink-connected GPUs (e.g. H100, B200 NVL), this enables NCCL to fuse collectives with NVLS (NVLink SHARP), reducing SM usage and latency. On IB-connected multi-node setups, set `NCCL_PXN_DISABLE=1` alongside `NCCL_CUMEM_ENABLE=1` — PXN is incompatible with buffer registration. + +**Effect on VisualGen parallelism modes:** +- **Ulysses** (all-to-all around self-attention): most impactful — the all-to-all is on the critical path of every denoising step. +- **Ring Attention** (K/V streaming all-gather): moderate — one all-gather per layer per step. +- **CFG parallel** (all-gather of noise predictions): minor — one collective per step. + +**Example config** (`examples/visual_gen/serve/configs/flux1_nccl_ub.yml`): + +```yaml +attention_config: + backend: VANILLA +parallel_config: + ulysses_size: 2 # Ulysses across 2 GPUs + nccl_buffer_reg: true # enables NCCL_CUMEM_ENABLE + ncclCommRegister +``` + +**Benchmarking** — collect baseline vs. UB-registration numbers with: + +```bash +# Baseline (no UB registration) +torchrun --nproc-per-node=2 examples/visual_gen/bench_nccl_ub.py --tier micro + +# With NCCL_CUMEM_ENABLE=1 +torchrun --nproc-per-node=2 examples/visual_gen/bench_nccl_ub.py --tier micro --nccl-cumem + +# Full pipeline (requires model weights) +torchrun --nproc-per-node=2 examples/visual_gen/bench_nccl_ub.py \ + --tier pipeline \ + --flux-path /path/to/FLUX.1-dev \ + --wan-path /path/to/Wan2.1-T2V-1.3B-Diffusers +``` + +Results are saved to `bench_nccl_ub_results.json`. + ## Developer Guide ### Architecture Overview diff --git a/examples/visual_gen/bench_cold_ring.py b/examples/visual_gen/bench_cold_ring.py new file mode 100644 index 000000000000..3e4879a3061b --- /dev/null +++ b/examples/visual_gen/bench_cold_ring.py @@ -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 + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--nccl-cumem", action="store_true") + ap.add_argument("--warmup", type=int, default=0) + ap.add_argument("--iters", type=int, default=7) + ap.add_argument("--tag", default="") + a = ap.parse_args() + + if a.nccl_cumem: + os.environ["NCCL_CUMEM_ENABLE"] = "1" + else: + os.environ["NCCL_CUMEM_ENABLE"] = "0" + os.environ.setdefault("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "1") + + from tensorrt_llm import VisualGen, VisualGenArgs, VisualGenParams + + PROMPTS = [ + "A majestic lion resting on a sun-drenched savanna", + "Futuristic city skyline reflected in a calm river", + "A waterfall cascading into a turquoise pool", + "Northern lights above a snow-covered forest", + "Ancient temple ruins illuminated by morning mist", + "Desert dunes at sunset with long shadows", + "Underwater coral reef teeming with tropical fish", + ] + + tag = a.tag or ("cumem" if a.nccl_cumem else "base") + print(f"\n{'='*60}") + print(f"COLD RING-4 NCCL_CUMEM_ENABLE={os.environ['NCCL_CUMEM_ENABLE']} tag={tag}") + print(f"warmup={a.warmup} iters={a.iters}") + print(f"{'='*60}\n") + + args = VisualGenArgs( + attention_config={"backend": "CUTEDSL"}, + parallel_config={"ring_size": 4}, + ) + print("Loading model...") + vg = VisualGen("/workspace/models/Cosmos3-Nano", args=args) + + gen_kwargs = dict(height=720, width=1280, num_inference_steps=20, + guidance_scale=6.0, seed=42, num_frames=57) + + 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() + result = { + "tag": tag, + "nccl_cumem": os.environ["NCCL_CUMEM_ENABLE"], + "warmup": a.warmup, + "iters": a.iters, + "all_s": [round(t, 3) for t in times], + "min_s": round(min(times), 3), + "median_s": round(statistics.median(times), 3), + "mean_s": round(statistics.mean(times), 3), + } + print(f"\nResult: {result}") + out = f"/workspace/results/cold_ring4_{tag}.json" + os.makedirs("/workspace/results", exist_ok=True) + with open(out, "w") as f: + json.dump(result, f, indent=2) + print(f"Saved → {out}") + +if __name__ == "__main__": + main() diff --git a/examples/visual_gen/bench_nccl_ub.py b/examples/visual_gen/bench_nccl_ub.py new file mode 100644 index 000000000000..e9a3d3032910 --- /dev/null +++ b/examples/visual_gen/bench_nccl_ub.py @@ -0,0 +1,324 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Benchmark NCCL user-buffer registration for VisualGen collectives. + +Two benchmark tiers: + + Tier 1 — Collective micro-benchmark (no model weights needed). + Measures latency of all-to-all and all-gather collectives that dominate + Ulysses sequence-parallel step time, with and without NCCL_CUMEM_ENABLE. + Run this tier to get numbers quickly on any multi-GPU node. + + Tier 2 — Full pipeline benchmark (model weights required). + Runs end-to-end generation with trtllm VisualGen and records wall-clock + step latency over N iterations. + +Usage — collective micro-benchmark (recommended first): + torchrun --nproc-per-node=2 bench_nccl_ub.py --tier micro + +Usage — full pipeline (requires weights): + torchrun --nproc-per-node=2 bench_nccl_ub.py \\ + --tier pipeline \\ + --flux-path /path/to/FLUX.1-dev \\ + --wan-path /path/to/Wan2.1-T2V-1.3B-Diffusers + +Results are written to bench_nccl_ub_results.json in the current directory +(only from rank 0). +""" + +import argparse +import json +import os +import time +from typing import Dict, List + +import torch +import torch.distributed as dist + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _is_rank0() -> bool: + return not dist.is_initialized() or dist.get_rank() == 0 + + +def _timed_iters(fn, warmup: int, iters: int) -> List[float]: + """Run fn() warmup times, then time iters calls. Returns latencies in ms.""" + for _ in range(warmup): + fn() + torch.cuda.synchronize() + times = [] + for _ in range(iters): + t0 = time.perf_counter() + fn() + torch.cuda.synchronize() + times.append((time.perf_counter() - t0) * 1000.0) + return times + + +def _stats(times: List[float]) -> Dict: + import statistics + return { + "mean_ms": statistics.mean(times), + "median_ms": statistics.median(times), + "min_ms": min(times), + "max_ms": max(times), + "stdev_ms": statistics.stdev(times) if len(times) > 1 else 0.0, + "n": len(times), + } + + +# --------------------------------------------------------------------------- +# Tier 1: Collective micro-benchmark +# --------------------------------------------------------------------------- + +def _bench_collective_one( + group: dist.ProcessGroup, + world_size: int, + shape: tuple, + dtype: torch.dtype, + collective: str, + warmup: int, + iters: int, +) -> Dict: + """Benchmark a single collective op on tensors of given shape.""" + device = torch.device(f"cuda:{torch.cuda.current_device()}") + tensor = torch.randn(*shape, dtype=dtype, device=device) + + if collective == "all_to_all": + seq_dim_size = shape[1] # [B, S/U, H/U, D] + assert seq_dim_size % world_size == 0 + out = torch.empty_like(tensor) + + def fn(): + dist.all_to_all_single(out, tensor, group=group) + + elif collective == "all_gather": + out = torch.empty( + shape[0], shape[1] * world_size, *shape[2:], dtype=dtype, device=device + ) + + def fn(): + dist.all_gather_into_tensor(out, tensor, group=group) + + elif collective == "all_reduce": + t = tensor.clone() + + def fn(): + dist.all_reduce(t, group=group) + + else: + raise ValueError(f"Unknown collective: {collective}") + + times = _timed_iters(fn, warmup=warmup, iters=iters) + return _stats(times) + + +def run_micro_benchmark(args) -> Dict: + """Tier-1: no model weights, pure collective timing.""" + rank = dist.get_rank() + world_size = dist.get_world_size() + device = torch.device(f"cuda:{rank}") + torch.cuda.set_device(device) + + results = {} + + # Shapes representative of FLUX (seq=4096 after patch) and Wan-1.3B (seq=7680) + configs = [ + # (name, collective, shape=(B, S/world_size, H/world_size, D)) + ("flux_all_to_all", "all_to_all", (1, 4096 // world_size, 24 // world_size, 128)), + ("flux_all_gather", "all_gather", (1, 4096 // world_size, 24, 128)), + ("wan_all_to_all", "all_to_all", (1, 7680 // world_size, 40 // world_size, 128)), + ("wan_all_gather", "all_gather", (1, 7680 // world_size, 40, 128)), + ("all_reduce_small", "all_reduce", (1, 512, 512)), + ] + + for name, collective, shape in configs: + stat = _bench_collective_one( + group=dist.group.WORLD, + world_size=world_size, + shape=shape, + dtype=torch.bfloat16, + collective=collective, + warmup=args.warmup, + iters=args.iters, + ) + results[name] = stat + if rank == 0: + print(f" {name:30s} mean={stat['mean_ms']:.3f} ms median={stat['median_ms']:.3f} ms") + + return results + + +# --------------------------------------------------------------------------- +# Tier 2: Full pipeline benchmark +# --------------------------------------------------------------------------- + +def _bench_pipeline_one(model_path: str, config_path: str, prompt: str, n_steps: int, + warmup: int, iters: int, height: int, width: int) -> Dict: + """Run trtllm VisualGen for n_steps, measure wall-clock per-image latency.""" + from tensorrt_llm import VisualGen, VisualGenArgs, VisualGenParams + + args = VisualGenArgs.from_yaml(config_path) + args.model = model_path + + vg = VisualGen(args) + params = VisualGenParams(height=height, width=width, num_inference_steps=n_steps) + + # Warmup + for _ in range(warmup): + vg.generate(inputs=prompt, params=params) + + times = [] + for _ in range(iters): + t0 = time.perf_counter() + vg.generate(inputs=prompt, params=params) + times.append((time.perf_counter() - t0) * 1000.0) + + vg.shutdown() + return _stats(times) + + +def run_pipeline_benchmark(args) -> Dict: + """Tier-2: end-to-end generation latency.""" + results = {} + prompt = "A photo of a cat sitting on a red velvet throne" + + benchmarks = [] + if args.flux_path: + benchmarks.append( + ("flux_baseline", args.flux_path, + "examples/visual_gen/serve/configs/flux1.yml", 1024, 1024) + ) + benchmarks.append( + ("flux_nccl_ub", args.flux_path, + "examples/visual_gen/serve/configs/flux1_nccl_ub.yml", 1024, 1024) + ) + if args.wan_path: + benchmarks.append( + ("wan_baseline", args.wan_path, + "examples/visual_gen/serve/configs/wan21.yml", 480, 832) + ) + benchmarks.append( + ("wan_nccl_ub", args.wan_path, + "examples/visual_gen/serve/configs/wan21_nccl_ub.yml", 480, 832) + ) + + for name, model_path, config_path, h, w in benchmarks: + if _is_rank0(): + print(f"\nBenchmarking {name} ({model_path}) ...") + stat = _bench_pipeline_one( + model_path=model_path, + config_path=config_path, + prompt=prompt, + n_steps=args.steps, + warmup=args.warmup, + iters=args.iters, + height=h, + width=w, + ) + results[name] = stat + if _is_rank0(): + print(f" {name:20s} mean={stat['mean_ms']:.1f} ms ({stat['n']} iters)") + + return results + + +# --------------------------------------------------------------------------- +# Summary table +# --------------------------------------------------------------------------- + +def _print_summary(results: Dict, tier: str): + print(f"\n{'='*70}") + print(f"NCCL User-Buffer Registration Benchmark — {tier}") + print(f"{'='*70}") + print(f"{'Key':<30} {'Mean ms':>10} {'Median ms':>10} {'Min ms':>8}") + print(f"{'-'*30} {'-'*10} {'-'*10} {'-'*8}") + for k, v in results.items(): + print(f"{k:<30} {v['mean_ms']:>10.3f} {v['median_ms']:>10.3f} {v['min_ms']:>8.3f}") + print() + + # For pipeline tier: compute speedup + if tier == "pipeline": + for prefix in ("flux", "wan"): + base = results.get(f"{prefix}_baseline") + ub = results.get(f"{prefix}_nccl_ub") + if base and ub: + speedup = base["mean_ms"] / ub["mean_ms"] + delta_ms = base["mean_ms"] - ub["mean_ms"] + print(f"{prefix.upper()} speedup: {speedup:.3f}x ({delta_ms:+.1f} ms per image)") + print(f"{'='*70}\n") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--tier", choices=["micro", "pipeline"], default="micro", + help="micro: collective microbenchmark (no weights); " + "pipeline: full end-to-end (requires --flux-path or --wan-path)") + parser.add_argument("--warmup", type=int, default=3, help="Warmup iterations") + parser.add_argument("--iters", type=int, default=10, help="Timed iterations") + parser.add_argument("--steps", type=int, default=20, + help="Denoising steps for pipeline tier") + parser.add_argument("--flux-path", type=str, default="", + help="Path to FLUX.1-dev checkpoint dir") + parser.add_argument("--wan-path", type=str, default="", + help="Path to Wan2.1-T2V-1.3B-Diffusers checkpoint dir") + parser.add_argument("--out", type=str, default="bench_nccl_ub_results.json", + help="Output JSON path (rank-0 only)") + parser.add_argument("--nccl-cumem", action="store_true", default=False, + help="Force NCCL_CUMEM_ENABLE=1 for this process " + "(micro tier: test WITH registration; " + "use two separate torchrun invocations to compare)") + args = parser.parse_args() + + # CUMEM must be set before init_process_group + if args.nccl_cumem: + os.environ["NCCL_CUMEM_ENABLE"] = "1" + + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="nccl", init_method="env://") + + rank = dist.get_rank() + world_size = dist.get_world_size() + + if rank == 0: + nccl_cumem = os.environ.get("NCCL_CUMEM_ENABLE", "0") + print(f"\nNCCL UB Benchmark tier={args.tier} " + f"world_size={world_size} NCCL_CUMEM_ENABLE={nccl_cumem}") + print(f"warmup={args.warmup} iters={args.iters}\n") + + if args.tier == "micro": + results = run_micro_benchmark(args) + else: + if not args.flux_path and not args.wan_path: + parser.error("Pipeline tier requires at least one of --flux-path / --wan-path") + results = run_pipeline_benchmark(args) + + if rank == 0: + _print_summary(results, args.tier) + payload = { + "tier": args.tier, + "world_size": world_size, + "nccl_cumem_enable": os.environ.get("NCCL_CUMEM_ENABLE", "0"), + "warmup": args.warmup, + "iters": args.iters, + "results": results, + } + with open(args.out, "w") as f: + json.dump(payload, f, indent=2) + print(f"Results saved to {args.out}") + + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/examples/visual_gen/serve/configs/flux1_nccl_ub.yml b/examples/visual_gen/serve/configs/flux1_nccl_ub.yml new file mode 100644 index 000000000000..c50bfc52fe15 --- /dev/null +++ b/examples/visual_gen/serve/configs/flux1_nccl_ub.yml @@ -0,0 +1,9 @@ +cache_config: + cache_backend: teacache + teacache_thresh: 0.6 +attention_config: + backend: VANILLA +parallel_config: + cfg_size: 1 + ulysses_size: 2 # NCCL UB registration is most effective with Ulysses > 1 + nccl_buffer_reg: true diff --git a/examples/visual_gen/serve/configs/wan21_nccl_ub.yml b/examples/visual_gen/serve/configs/wan21_nccl_ub.yml new file mode 100644 index 000000000000..c5444dbeea28 --- /dev/null +++ b/examples/visual_gen/serve/configs/wan21_nccl_ub.yml @@ -0,0 +1,7 @@ +cache_config: + cache_backend: teacache + teacache_thresh: 0.2 +parallel_config: + cfg_size: 1 + ulysses_size: 2 # NCCL UB registration is most effective with Ulysses > 1 + nccl_buffer_reg: true diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 8b4f2dd876d2..7fa4e75966db 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -387,6 +387,14 @@ def run_diffusion_worker( if torch.cuda.is_available(): torch.cuda.set_device(device_id) + # NCCL user-buffer registration: set NCCL_CUMEM_ENABLE before + # init_process_group so NCCL uses VMM-backed scratch buffers. + # Must happen before the communicator is created. + _parallel_cfg = getattr(visual_gen_args, "parallel_config", None) + if _parallel_cfg is not None and getattr(_parallel_cfg, "nccl_buffer_reg", False): + from tensorrt_llm._torch.visual_gen.nccl_ub_reg import enable_nccl_cumem + enable_nccl_cumem() + dist.init_process_group( backend="cuda:nccl,cpu:gloo" if torch.cuda.is_available() else "gloo", init_method="env://", diff --git a/tensorrt_llm/_torch/visual_gen/nccl_ub_reg.py b/tensorrt_llm/_torch/visual_gen/nccl_ub_reg.py new file mode 100644 index 000000000000..f302a77375c9 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/nccl_ub_reg.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""NCCL user-buffer registration for VisualGen collectives. + +Two complementary mechanisms are provided: + +1. **CUMEM mode** (``enable_nccl_cumem()``): sets ``NCCL_CUMEM_ENABLE=1`` + before ``init_process_group``. NCCL then allocates its internal scratch + via cuMemCreate (VMM), enabling zero-copy for all collectives on VMM-backed + tensors. No per-tensor API call required. Requires NCCL ≥ 2.21. + +2. **Explicit registration** (``NCCLBufferRegistrar``): calls + ``ncclCommRegister`` / ``ncclCommDeregister`` on a specific tensor via + ctypes. Useful for pinning a persistent latent buffer that is used in the + same collective call across all denoising steps (same pointer, same + communicator, same offset alignment). Requires NCCL ≥ 2.19. + + **Constraint (from NCCL docs)**: if *any* rank passes a registered buffer + to a collective, *all* other ranks in the same communicator must also pass + their registered counterpart. The registrar is therefore all-or-nothing + per communicator. + +Reference: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/bufferreg.html +""" + +import ctypes +import logging +import os +from typing import Optional + +import torch +import torch.distributed as dist + +logger = logging.getLogger(__name__) + +# Minimum NCCL version codes for each feature. +_NCCL_MIN_REG = (2, 19, 0) # ncclCommRegister / ncclCommDeregister +_NCCL_MIN_CUMEM = (2, 21, 0) # NCCL_CUMEM_ENABLE / VMM allocation + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _load_nccl() -> Optional[ctypes.CDLL]: + """Try to load libnccl.so from standard locations.""" + for lib_name in ("libnccl.so.2", "libnccl.so"): + try: + return ctypes.CDLL(lib_name, mode=ctypes.RTLD_GLOBAL) + except OSError: + continue + return None + + +def _nccl_version(lib: ctypes.CDLL) -> tuple: + """Return (major, minor, patch) from ncclGetVersion.""" + ver = ctypes.c_int(0) + lib.ncclGetVersion(ctypes.byref(ver)) + v = ver.value + return (v // 10000, (v // 100) % 100, v % 100) + + +def _check_nccl_version(lib: ctypes.CDLL, required: tuple, feature: str) -> bool: + ver = _nccl_version(lib) + if ver < required: + logger.warning( + f"NCCL {ver[0]}.{ver[1]}.{ver[2]} < " + f"{required[0]}.{required[1]}.{required[2]}: " + f"{feature} not available" + ) + return False + return True + + +# --------------------------------------------------------------------------- +# Mechanism 1: CUMEM env-var (set before init_process_group) +# --------------------------------------------------------------------------- + +def enable_nccl_cumem() -> bool: + """Enable NCCL VMM-based buffer registration via ``NCCL_CUMEM_ENABLE=1``. + + Must be called **before** ``torch.distributed.init_process_group()``. + Returns True if the env var was set (or was already set), False if NCCL + is too old to support it. + + When enabled, NCCL allocates its scratch buffers with ``cuMemCreate`` + (VMM). Any CUDA tensor allocated by PyTorch's caching allocator that + happens to be VMM-backed (``expandable_segments:True`` on CUDA 11.8+) + will also be eligible for zero-copy collectives automatically. + """ + lib = _load_nccl() + if lib is None: + logger.warning("libnccl.so not found; cannot enable CUMEM mode") + return False + if not _check_nccl_version(lib, _NCCL_MIN_CUMEM, "NCCL_CUMEM_ENABLE"): + return False + + if os.environ.get("NCCL_CUMEM_ENABLE") == "1": + logger.debug("NCCL_CUMEM_ENABLE already set") + else: + os.environ["NCCL_CUMEM_ENABLE"] = "1" + logger.info("Set NCCL_CUMEM_ENABLE=1 (VMM-based buffer registration)") + return True + + +# --------------------------------------------------------------------------- +# Mechanism 2: Explicit per-tensor ncclCommRegister +# --------------------------------------------------------------------------- + +def _extract_comm_ptr(group: dist.ProcessGroup) -> Optional[int]: + """Best-effort: extract the raw ``ncclComm_t`` pointer from a PyTorch process group. + + PyTorch does not expose this pointer in a stable public API. We try a + sequence of progressively more fragile introspection methods and return + ``None`` if none succeed. + """ + try: + device = torch.device(f"cuda:{torch.cuda.current_device()}") + backend = group._get_backend(device) + + # PyTorch ≥ 2.6 ProcessGroupNCCL exposes _get_comm() returning the + # ncclComm_t as an integer. + if hasattr(backend, "_get_comm"): + return int(backend._get_comm(device)) + + # 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 + ctypes.pythonapi.PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p] + ptr = ctypes.pythonapi.PyCapsule_GetPointer(cap, None) + return ptr + + except Exception as exc: + logger.debug(f"Could not extract NCCL comm ptr: {exc}") + return None + + +class NCCLBufferRegistrar: + """Registers specific tensors with ``ncclCommRegister`` for zero-copy collectives. + + All ranks in the communicator **must** register the corresponding tensor + before using it in a collective; partial registration is undefined behaviour. + + Typical usage:: + + registrar = NCCLBufferRegistrar(ulysses_group) + if registrar.available: + registrar.register_all([latent, scratch_q, scratch_kv]) + # ... run denoising loop (same tensor pointers) ... + registrar.deregister_all() + + Registration is a no-op (and ``available`` is False) when: + - libnccl.so cannot be loaded + - NCCL version < 2.19 + - the raw ncclComm_t pointer cannot be extracted from the process group + """ + + def __init__(self, process_group: dist.ProcessGroup): + self._group = process_group + self._lib: Optional[ctypes.CDLL] = None + self._comm_ptr: Optional[int] = None + self._handles: dict[int, ctypes.c_void_p] = {} # data_ptr → opaque handle + self.available = False + + lib = _load_nccl() + if lib is None: + logger.warning("NCCLBufferRegistrar: libnccl.so not found") + return + if not _check_nccl_version(lib, _NCCL_MIN_REG, "ncclCommRegister"): + return + comm_ptr = _extract_comm_ptr(process_group) + if comm_ptr is None: + logger.info( + "NCCLBufferRegistrar: could not extract ncclComm_t; " + "CUMEM env-var mode is still active if enabled" + ) + return + + self._lib = lib + self._comm_ptr = comm_ptr + self.available = True + logger.info("NCCLBufferRegistrar ready (ncclCommRegister available)") + + 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() + + def __del__(self): + try: + self.deregister_all() + except Exception: + pass diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 8daefd5d7bb6..bd5b128ff9f1 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -232,6 +232,19 @@ class ParallelConfig(StrictBaseModel): status="prototype", description=("Tensor parallel group size. Heads are sharded across tp_size GPUs."), ) + nccl_buffer_reg: bool = Field( + False, + status="prototype", + description=( + "Enable NCCL user-buffer registration for zero-copy collectives. " + "Sets NCCL_CUMEM_ENABLE=1 before init_process_group and, when " + "possible, calls ncclCommRegister on persistent latent buffers. " + "Requires NCCL ≥ 2.19 (CUMEM mode requires ≥ 2.21). " + "Most effective with Ulysses or ring-attention sequence parallelism " + "where all-to-all / all-gather dominate step latency. " + "See tensorrt_llm/_torch/visual_gen/nccl_ub_reg.py." + ), + ) @property def seq_parallel_size(self) -> int: