[gpt-oss] add triton-kernels build for gpt-oss in ROCm base Dockerfile (#637)

* add triton-kernel build for gpt-oss in base dockerfile

Signed-off-by: Hongxia Yang <[email protected]>

* add conditional check

Signed-off-by: Hongxia Yang <[email protected]>

* dealing with different setup.py locations in different triton branches

Signed-off-by: Hongxia Yang <[email protected]>

* fixed the dist location and reduce one copy operation

Signed-off-by: Hongxia Yang <[email protected]>

---------

Signed-off-by: Hongxia Yang <[email protected]>

Author: hongxiayang

docker/Dockerfile.rocm_base

@@ -74,9 +74,11 @@ ARG TRITON_REPO
 RUN git clone ${TRITON_REPO}
 RUN cd triton \
     && git checkout ${TRITON_BRANCH} \
-    && cd python \
-    && python3 setup.py bdist_wheel --dist-dir=dist
-RUN mkdir -p /app/install && cp /app/triton/python/dist/*.whl /app/install
+    && if [ ! -f setup.py ]; then cd python; fi \
+    && python3 setup.py bdist_wheel --dist-dir=dist \
+    && mkdir -p /app/install && cp dist/*.whl /app/install
+RUN if [ -d triton/python/triton_kernels ]; then pip install build && cd triton/python/triton_kernels \
+    && python3 -m build --wheel && cp dist/*.whl /app/install; fi
 
 FROM base AS build_amdsmi
 RUN cd /opt/rocm/share/amd_smi \