Skip to content

[https://nvbugs/6262973][feat] Pool NCCL window tensor outputs - #1

Open
nv-lschneider wants to merge 9 commits into
graph-sage-symmetricfrom
perf/nccl-window-output-tensor-pool-only-review
Open

[https://nvbugs/6262973][feat] Pool NCCL window tensor outputs#1
nv-lschneider wants to merge 9 commits into
graph-sage-symmetricfrom
perf/nccl-window-output-tensor-pool-only-review

Conversation

@nv-lschneider

Copy link
Copy Markdown
Owner

@coderabbitai summary

Description

Reuse engine-owned NCCL-window storage for eligible row-parallel GEMM outputs,
allowing the following all-reduce to consume the GEMM result directly.

The pool keeps a stable allocation for warmup and CUDA-graph capture, normalizes
its capacity across TP ranks, and falls back to the allocating path when a later
shape exceeds the retained capacity.

Performance improvements were observed for NVFP4 Llama 3.3 70B and FP8
QDQ/cuBLAS Llama 3.3 70B. FP8 rowwise Llama 3.1 8B was also exercised.
The cublas_mm path remains explicit opt-in because it did not show a benefit
for GPT-OSS 20B.

Test Coverage

  • NCCL-window tensor-pool unit tests, including fixed-capacity fallback.
  • TP-wide capacity-selection unit test.
  • Llama 3.3 70B NVFP4: exact pool-off/pool-on outputs.
  • Llama 3.3 70B FP8 QDQ/cuBLAS: exact outputs with NCCL_SYMMETRIC.
  • Llama 3.1 8B FP8 rowwise: exact pool-off/pool-on outputs.
  • GPT-OSS 20B cublas_mm: exercised as an explicit opt-in path.

PR Checklist

  • PR description explains the change and motivation.
  • Relevant tests and model-path validation were run.
  • No API change.
  • No new dependencies.
  • Documentation is not required for this internal implementation change.

Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>
Support caller-provided NVFP4 outputs across CUTLASS, cuBLASLt, CUDA-core, CuteDSL, and Marlin backends while preserving unified autotuning.

Register the output variants with compilation metadata and fake implementations, enable the validated FP8-QDQ pooling policy, and expand backend correctness coverage.
@nv-lschneider
nv-lschneider force-pushed the perf/nccl-window-output-tensor-pool-only-review branch from 74e8b74 to 27e573a Compare August 25, 2026 14:48
@nv-lschneider
nv-lschneider force-pushed the graph-sage-symmetric branch 5 times, most recently from 12ea73e to 14a045a Compare September 3, 2026 17:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant