cp: build: Fix ngc pytorch build with deep-ep (1234) into r0.4.0

[chtruong814]

beep boop [🤖]: Hi @chtruong814 👋,

we've cherry picked #1234 into  for you! 🚀

Please review and approve this cherry pick by your convenience!

Summary by CodeRabbit

• Chores
  • Updated Docker image build to include CUDA architectures for latest-generation NVIDIA GPUs (e.g., H100), improving compatibility and performance.
  • Enhanced build process with an additional dependency sync step to ensure required GPU communication libraries are present before compilation.
  • Expanded dependency groups resolved during image build for broader coverage and fewer setup issues, resulting in more reliable container builds.

[coderabbitai]

📝 Walkthrough

Walkthrough

Updates the NGC PyTorch Dockerfile: sets TORCH_CUDA_ARCH_LIST to "9.0 10.0" in the hermetic build stage and modifies the build RUN steps to ensure nvshmem is installed before building DeepEP, followed by an additional uv sync with expanded extra groups and flags.

Changes

Cohort / File(s)	Summary
Docker build updates docker/Dockerfile.ngc_pytorch	Added ENV TORCH_CUDA_ARCH_LIST="9.0 10.0". Updated build RUN commands to sync/install nvshmem before DeepEP, then perform an additional uv sync with extended --extra groups (vllm, mcore, automodel, all-groups) and flags, maintaining no-install-project behavior and improving dependency resolution order.

Sequence Diagram(s)

sequenceDiagram
    autonumber
    participant Dev as Developer
    participant Docker as Docker Build
    participant Herm as Hermetic Stage
    participant UV as uv
    participant NVSH as nvshmem
    participant DeepEP as DeepEP Build

    Dev->>Docker: build
    Docker->>Herm: start build stage
    Note over Herm: ENV TORCH_CUDA_ARCH_LIST="9.0 10.0"
    Herm->>UV: uv sync (ensure nvshmem)
    UV->>NVSH: install/verify nvshmem
    NVSH-->>UV: ready
    Herm->>DeepEP: build DeepEP
    DeepEP-->>Herm: built
    Herm->>UV: uv sync --extra vllm,mcore,automodel,all-groups ...<br/>--no-install-project
    UV-->>Herm: dependencies synced
    Herm-->>Docker: stage complete

Loading

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• build: Fix ngc pytorch build with deep-ep #1234 — Also adjusts docker/Dockerfile.ngc_pytorch to set TORCH_CUDA_ARCH_LIST for DeepEP; closely related to CUDA arch targeting and build-stage configuration.

Suggested labels

r0.4.0

Suggested reviewers

• guyueh1
• terrykong

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Title Check	⚠️ Warning	The title primarily documents the cherry-pick operation and branch target instead of succinctly summarizing the actual change, and it includes backticks, PR numbers, and branch names that add noise. It does not clearly convey the key update to the Dockerfile, namely adding CUDA architecture settings and sync steps for DeepEP compatibility.	Please revise the title to focus on the functional change, for example “build: add TORCH_CUDA_ARCH_LIST and sync steps in ngc_pytorch Dockerfile for DeepEP support,” and remove cherry-pick metadata and branch references.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage	✅ Passed	No functions found in the changes. Docstring coverage check skipped.
Test Results For Major Changes	✅ Passed	Based on the summary, this PR only updates a Dockerfile to set TORCH_CUDA_ARCH_LIST and adjusts build steps/dependency sync for DeepEP; there are no model, algorithm, or runtime logic changes. While the PR description lacks test results, the changes are build/infrastructure-only and unlikely to affect numerics or runtime performance. Therefore, this qualifies as a minor change under the check criteria, and explicit test results are not required for passing.

✨ Finishing touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch cherry-pick-1234-r0.4.0

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

docker/Dockerfile.ngc_pytorch (1)

100-103: Avoid installing all extras; add an explicit NVSHMEM presence check before DeepEP build

Installing with --all-groups can pull dev/test extras, inflating image size and build time. Also, the comment implies nvshmem must pre-exist; make that explicit with a fast failure check.

• Keep extras minimal (only what’s required to build/use DeepEP).
• Verify where nvshmem actually comes from (system vs pip meta). If it’s a system lib, add a preflight check.

Apply:

-# Ensure nvshmem is installed before building DeepEP
-uv sync --link-mode symlink --locked --inexact --no-install-project $UV_NO_INSTALL_PACKAGES
-uv sync --link-mode symlink --locked --inexact --extra vllm --extra mcore --extra automodel --all-groups --no-install-project $UV_NO_INSTALL_PACKAGES
+# Ensure NVSHMEM runtime is present before building DeepEP (fail fast if missing)
+ldconfig -p | grep -q nvshmem || { echo "ERROR: NVSHMEM runtime not found on system path"; exit 1; }
+# Prime base deps without installing the project to keep layer size lower
+uv sync --link-mode symlink --locked --inexact --no-install-project $UV_NO_INSTALL_PACKAGES
+# Install only the extras needed (avoid --all-groups to prevent dev/test bloat)
+uv sync --link-mode symlink --locked --inexact \
+  --extra vllm --extra mcore --extra automodel \
+  --no-install-project $UV_NO_INSTALL_PACKAGES

If DeepEP requires an additional extra (e.g., --extra deep_ep), prefer adding that explicitly over --all-groups. Please confirm which extra actually triggers DeepEP’s build in pyproject.

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 82ab565 and b7e687a.

📒 Files selected for processing (1)

• docker/Dockerfile.ngc_pytorch (2 hunks)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (5)

• GitHub Check: Lint check
• GitHub Check: Lint check
• GitHub Check: Lint check
• GitHub Check: Post submodule check comment / Comment on PR
• GitHub Check: Post automodel integration comment / Comment on PR

🔇 Additional comments (1)

docker/Dockerfile.ngc_pytorch (1)

86-87: Use semicolon-separated TORCH_CUDA_ARCH_LIST with +PTX; verify sm_100 support

Replace:

- ENV TORCH_CUDA_ARCH_LIST="9.0 10.0"
+ ENV TORCH_CUDA_ARCH_LIST="9.0;10.0+PTX"

Confirm that nvcr.io/nvidia/pytorch:25.06-py3’s toolchain supports sm_100; gate or remove it if unsupported.