MindSpeed-Bridge Docker Image Overview

This directory provides a configurable Docker build entry for MindSpeed-Bridge on Ascend NPU.

The image uses the Python environment provided by the CANN base image directly and does not install Conda.

Default stack:

• AscendHub CANN base image: 9.1.0-beta.3
• TorchNPU: 2.9.0
• PyTorch: 2.9.0
• torch_npu: 2.9.0, installed from the default torch_npu-2.9.0.post5-cp312-cp312-manylinux_2_28_aarch64.whl URL
• triton-ascend: 3.2.1
• MindSpeed: core_r0.16.0
• MindSpeed-Ops: master
• flash-linear-attention-npu: v26.1.0
• Megatron-LM: core_v0.16.1
• Megatron-Bridge: v0.3.1
• Python: 3.12
• Supported NPU types: 910b, a3, 950

flash-linear-attention-npu Ops Build

During image build, the Dockerfile sources the CANN environment after cloning flash-linear-attention-npu, then builds and installs the GDN custom operator run package and the torch_custom/fla_npu wheel.

The FLA NPU --soc value is mapped from the selected NPU type by default:

NPU type	FLA NPU --soc
910b	ascend910b
a3	ascend910_93
950	ascend950

Override the mapping if needed:

bash image_build.sh --fla-npu-soc ascend910_93

The FLA NPU operator list is maintained in the FLA_NPU_OPS array in docker/image_build.sh. Add new operator names to that array, and the script will convert it to the comma-separated value required by build.sh --ops.

Build

cd docker
bash image_build.sh

Use a full AscendHub base image name when the published image tag differs from the default naming rule:

bash image_build.sh \
  --base-image swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:9.1.0-beta.3-910b-openeuler24.03-py3.12

The build uses the configured TorchNPU wheel URL by default. Override it only when another wheel is required:

bash image_build.sh \
  --torch-npu-whl-url https://gitcode.com/Ascend/pytorch/releases/download/v26.1.0-beta.2-pytorch2.9.0/torch_npu-2.9.0.post5-cp312-cp312-manylinux_2_28_aarch64.whl

To install torch_npu from the configured pip source instead, pass an empty wheel URL. The fallback pip package version follows TORCH_NPU_VERSION.

bash image_build.sh --torch-npu-whl-url ""

The build script uses --network=host by default. If proxy environment variables such as http_proxy, https_proxy, HTTP_PROXY, HTTPS_PROXY, no_proxy, or NO_PROXY are set, they are automatically forwarded to Docker build as build arguments.

The image installs triton-ascend==3.2.1 from https://triton-ascend.osinfra.cn/pypi/simple by default. Override the version with:

bash image_build.sh --triton-ascend-version 3.2.1

See OVERVIEW.zh.md for detailed Chinese instructions.