docker: add clean build and wheel-based install Dockerfiles

[sunway513]

Summary

• Add Dockerfile.wheels: multi-stage builder that produces all dependency wheels (torch, triton, triton_kernels, flydsl, mori, amd_aiter)
• Add Dockerfile.clean: wheel-only install with zero source compilation, using pre-built wheels from Dockerfile.wheels

No changes to the existing Dockerfile.

Build Time (MI300X node)

Dockerfile	Time	Notes
Dockerfile.wheels	~10 min	MORI compile 3.4min, AITER 1min, rest are downloads
Dockerfile.clean	~4 min	Pure pip install, zero compilation
Total	~14 min	vs ~2+ hours for full source build

Wheels Build Breakdown

Step	Time	Description
System packages	49s	apt-get install
PyTorch download	128s	torch + torchvision + torchaudio (5.8GB)
Triton download (PyPI 3.6.0)	20s	180MB
triton_kernels	37s	sparse-checkout from ROCm/triton, pure Python
FlyDSL download	11s	pre-built nightly (70MB)
torch+triton install	82s	needed for MORI/AITER builds
MORI compile	202s	slowest component
AITER compile (CK-free)	63s	ENABLE_CK=0

Usage

Build wheels

docker build -f docker/Dockerfile.wheels -t atom:wheels .

Install from wheels (zero compilation)

DOCKER_BUILDKIT=1 docker build \
  --build-context wheels=docker-image://atom:wheels \
  -f docker/Dockerfile.clean -t atom:clean .

Test plan

• Build Dockerfile.wheels — all 8 wheels produced
• Build Dockerfile.clean from wheels — all imports succeed (PyTorch, Triton, AITER, FlyDSL, MORI, ATOM)
• Sanity check clean image — 8x MI300X detected, GPU matmul OK
• E2E inference — blocked on AITER CK-free attention fallback (known limitation, being addressed by AITER team)

[sunway513]

Build Verification Results

Test	Status	Notes
Dockerfile.wheels build	PASS	8 wheels produced in ~10 min
Dockerfile.clean build	PASS	Zero compilation, ~3.7 min
Clean image sanity check	PASS	All imports OK, 8x MI300X detected, GPU matmul OK
Clean image E2E inference	BLOCKED	See below
Dockerfile full build (ENABLE_CK=1)	FAIL	Triton 3.5.x source compilation error (upstream issue, unrelated to this PR)

E2E Inference Limitation

E2E inference with the clean image (ENABLE_CK=0) fails because AITER's attention path (fmha_v3_varlen_fwd) still attempts to JIT-compile CK modules even when ATOM_CK_FREE=1 and AITER_FORCE_TRITON_ATTN=1 are set. This is a known limitation — the CK-free attention fallback in AITER is being addressed by the AITER team.

Once that fix lands, the clean Docker image will support full E2E inference out of the box.

[Copilot]

Pull request overview

Adds alternative Docker build paths that prebuild/download dependency wheels once and then create a “clean” runtime image by installing only from wheels (no source compilation), significantly reducing Docker build times for ROCm environments.

Changes:

• Add docker/Dockerfile.wheels to build/download wheels for major dependencies (PyTorch ROCm nightly, Triton, triton_kernels, FlyDSL, MORI, AITER).
• Add docker/Dockerfile.clean to install exclusively from those prebuilt wheels using BuildKit build-context mounting.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 5 comments.

File	Description
docker/Dockerfile.wheels	Wheel builder image that downloads/builds dependency wheels into /wheels for reuse.
docker/Dockerfile.clean	Runtime image that installs all dependencies from provided wheels via BuildKit mount, then installs ATOM from the build context.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

apt-get install includes git, but this Dockerfile does not use git (and the header says “No git clones”). Removing git reduces image size and keeps the “minimal” runtime promise accurate (or update the comment if git is intentionally required at runtime).

[Copilot]

The wheels mount copies from source=/ and then runs find /mnt/wheels -name '*.whl' .... When wheels is a docker-image build context, this will traverse the entire image filesystem and can accidentally pick up wheels from unexpected locations (pip cache, etc.) and slow the build. Consider limiting the mount/copy to the known wheels directory produced by Dockerfile.wheels (/wheels), or explicitly searching only that subdirectory when present.

[Copilot]

AITER is cloned but submodules aren’t initialized/updated before building. The existing docker/Dockerfile does a git submodule ... update --init --recursive for AITER, which suggests the build may rely on submodules; without this, the wheel build can fail or produce an incomplete package. Please add the submodule init/update step (or use git clone --recurse-submodules) before installing/building.

[Copilot]

PREBUILD_TRITON_ARCHS is hardcoded to gfx942,gfx950 even though the target arch list is configurable via ARG GPU_ARCH. If someone builds with a different GPU_ARCH, AITER will still prebuild for 942/950, which can lead to incorrect or missing kernels. Consider deriving PREBUILD_TRITON_ARCHS from the same build arg (or introducing a separate build arg) to keep them consistent.

[Copilot]

The header comment calls this a “Multi-stage wheel builder”, but the Dockerfile currently has a single stage (no FROM ... AS ... / multi-stage copy). Either adjust the wording or split into explicit stages if that was intended, to avoid confusing users.

[sunway513]

Related AITER PR: ROCm/aiter#2227 — Adds a Triton fallback for fused_rope_rms() in AITER. This is one of the pieces needed to unblock full E2E CK-free inference with the clean Docker image.