Add ragged FP8 MoE path and pinned runtime image for GPT-OSS 20B MLPerf

.gitignore

@@ -26,6 +26,8 @@ pp_simulation_result
 
 *.log
 *.nohup
+slurm-*.out
+spur-*.out
 *.zip
 .triton_cache_shared/
 .cursor/

examples/mlperf/gpt_oss_20b/.dockerignore

@@ -0,0 +1,10 @@
+**
+!Dockerfile.runtime
+!runtime_versions.sh
+!prepare_triton_runtime.sh
+!prepare_primus_turbo_runtime.sh
+!prepare_mlperf_runtime.sh
+!install_asm_attention.sh
+!verify_runtime.py
+!aiter_hd64_asm_override.py
+!aiter_hd64_asm_override.pth

examples/mlperf/gpt_oss_20b/Dockerfile.runtime

@@ -0,0 +1,55 @@
+# syntax=docker/dockerfile:1.7
+
+ARG BASE_IMAGE=rocm/primus:v26.3@sha256:1a02b74a94d82131f3f119a94ccbed45bb5b4cd77a0616cb4e6a29e64a048482
+FROM ${BASE_IMAGE}
+
+ARG TRITON_REF=09500db9f0fe66fd176d1f080e2017b37e7e995d
+ARG PRIMUS_TURBO_REF=7593fc4443b11f7b4714f3f7a284b4873d398177
+ARG FLYDSL_VERSION=0.2.0
+ARG TE_REF=5235bae2cc683a0ad4bf15221c746ab3c1e229e7
+ARG AITER_REF=c4b33df03faae1c4e470420d950f8a9589e9634d
+ARG FWD_ATTN_ASM_REF=53d3dadc3f3b0ac35ae536f2d1d7864a3e07ba22
+ARG BWD_ATTN_ASM_REF=9b9fb6444f3fee388617f62432c3faea74079377
+ARG BWD_ATTN_ASM_SYMBOL=_ZN5aiter43fmha_bwd_hd64_bf16_causal_a16_rtz_recompileE
+ARG BUILD_JOBS=96
+
+ENV MLPERF_TRITON_REF=${TRITON_REF} \
+    MLPERF_PRIMUS_TURBO_REF=${PRIMUS_TURBO_REF} \
+    MLPERF_FLYDSL_VERSION=${FLYDSL_VERSION} \
+    MLPERF_TE_REF=${TE_REF} \
+    MLPERF_AITER_REF=${AITER_REF} \
+    MLPERF_FWD_ATTN_ASM_REF=${FWD_ATTN_ASM_REF} \
+    MLPERF_BWD_ATTN_ASM_REF=${BWD_ATTN_ASM_REF} \
+    MLPERF_BWD_ATTN_ASM_SYMBOL=${BWD_ATTN_ASM_SYMBOL} \
+    MLPERF_ENABLE_FWD_ATTN_ASM=0
+
+LABEL org.opencontainers.image.title="Primus GPT-OSS-20B MLPerf runtime" \
+      org.opencontainers.image.description="Pinned Triton, Primus-Turbo, TransformerEngine, and gfx950 ASM attention" \
+      com.amd.mlperf.triton.revision="${TRITON_REF}" \
+      com.amd.mlperf.primus-turbo.revision="${PRIMUS_TURBO_REF}" \
+      com.amd.mlperf.transformer-engine.revision="${TE_REF}" \
+      com.amd.mlperf.aiter.revision="${AITER_REF}" \
+      com.amd.mlperf.fwd-attn-asm.revision="${FWD_ATTN_ASM_REF}" \
+      com.amd.mlperf.bwd-attn-asm.revision="${BWD_ATTN_ASM_REF}"
+
+WORKDIR /opt/mlperf-gpt-oss-20b
+
+COPY runtime_versions.sh prepare_triton_runtime.sh \
+     prepare_primus_turbo_runtime.sh prepare_mlperf_runtime.sh \
+     install_asm_attention.sh verify_runtime.py \
+     aiter_hd64_asm_override.py aiter_hd64_asm_override.pth ./
+
+RUN --mount=type=cache,id=gptoss20b-mlperf-build,target=/opt/mlperf-build,sharing=locked \
+    chmod 0755 ./*.sh ./verify_runtime.py \
+    && PRIMUS_PATH=/opt/mlperf-build \
+       MLPERF_DEPS_DIR=/opt/mlperf-build/deps \
+       MLPERF_TRITON_WHEEL_CACHE=/opt/mlperf-build/wheels/triton \
+       MLPERF_PRIMUS_TURBO_WHEEL_CACHE=/opt/mlperf-build/wheels/primus-turbo \
+       MLPERF_TE_WHEEL_CACHE=/opt/mlperf-build/wheels/transformer-engine \
+       MAX_JOBS="${BUILD_JOBS}" \
+       bash -lc 'source /opt/mlperf-gpt-oss-20b/prepare_primus_turbo_runtime.sh \
+           && source /opt/mlperf-gpt-oss-20b/prepare_mlperf_runtime.sh \
+           && python3 /opt/mlperf-gpt-oss-20b/verify_runtime.py' \
+    && rm -rf /root/.cache/pip
+
+WORKDIR /workspace

examples/mlperf/gpt_oss_20b/README.md

@@ -1,39 +1,94 @@
 # GPT-OSS-20B Pretraining Benchmark
 
-GPT-OSS 20B (Mixture of Experts)
+This benchmark trains the GPT-OSS 20B MoE model with FP8 grouped GEMM on
+8×MI355X.
 
+## Runtime image
 
-## Setup
+The measured software stack can be built once instead of rebuilding Triton,
+Primus-Turbo, TransformerEngine, and the attention kernels for every job:
 
-### Start Docker Image
+```bash
+cd examples/mlperf/gpt_oss_20b
+bash build_runtime_image.sh primus:gpt-oss-20b-mlperf-v1
+```
+
+The image is based on `rocm/primus:v26.3` (pinned by digest) and contains:
+
+- Triton 3.7 at the revision pinned in `runtime_versions.sh`;
+- Primus-Turbo built against that Triton revision;
+- TransformerEngine/AITER with the HD64 causal backward ASM code object
+  embedded in `te_libmha_bwd.so`;
+- the HD64 causal forward ASM code object and its gated Python dispatch shim.
+
+Docker images are local to the build node. Push the resulting tag to a shared
+registry before using it on other nodes:
 
 ```bash
-docker run -it     --device /dev/dri     --device /dev/kfd     --device /dev/infiniband     --network host --ipc host     --group-add video     --cap-add SYS_PTRACE     --security-opt seccomp=unconfined     --privileged     -v $HOME:$HOME   --shm-size 128G     --name primus_training_env rocm/primus:v26.5
+PUSH_IMAGE=1 bash build_runtime_image.sh \
+  registry.example.com/team/primus:gpt-oss-20b-mlperf-v1
+```
 
-cd /workspace/Primus
+Verify the image on an 8-GPU node:
+
+```bash
+docker run --rm \
+  --device /dev/dri --device /dev/kfd \
+  --network host --ipc host --privileged \
+  primus:gpt-oss-20b-mlperf-v1 \
+  python3 /opt/mlperf-gpt-oss-20b/verify_runtime.py --require-gpus 8
 ```
 
+Run the full 1.2M-iteration configuration with the prebuilt image:
 
-### Configuration
+```bash
+DATA_HOST_PATH=/path/to/data \
+MODEL_HOST_PATH=/path/to/model \
+bash run_runtime_longrun.sh
+```
 
-This benchmark trains a 20B parameter GPT model with Mixture of Experts (MoE) architecture using the Primus framework on AMD GPUs.
+`TRAIN_ITERS` may be overridden for a bounded stability run. With the configured
+128-step warmup, use more than 256 iterations so the decay schedule remains
+valid. The script verifies all eight GPUs and the pinned runtime before training
+starts.
 
-**Key Features:**
-- 20B parameter MoE model
-- Expert Parallelism (EP=8)
-- FP8 hybrid precision training
-- Primus Turbo optimizations (DeepEP, sync-free MoE)
+`install_asm_attention.sh` owns both ASM integrations. Forward dispatch remains
+gated by `MLPERF_ENABLE_FWD_ATTN_ASM=1`; backward ASM is fixed at image build
+time.
 
-## Key Files
+## Start the training container
 
-- `configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml` - Model and training config
-    - Update `train_data_path` and `train_data_path` to your local downloaded location
-- `config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh` - System config and env vars
-   - Update `PRIMUS_PATH` to clone Primus Repo
-   - Update `EXP`to `<PRIMUS_PATH>/examples/mlperf/configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml`
-- `run_and_time.sh` - Run script
+```bash
+docker run -it --rm \
+  --device /dev/dri \
+  --device /dev/kfd \
+  --device /dev/infiniband \
+  --network host \
+  --ipc host \
+  --cap-add SYS_PTRACE \
+  --security-opt seccomp=unconfined \
+  --privileged \
+  --shm-size 128G \
+  -v /path/to/Primus:/workspace/Primus \
+  -v /path/to/data:/data \
+  -v /path/to/results:/results \
+  primus:gpt-oss-20b-mlperf-v1
+```
 
-### Data
+## Key files
+
+- `runtime_versions.sh`: validated dependency revisions.
+- `Dockerfile.runtime`, `build_runtime_image.sh`: reproducible runtime image.
+- `install_asm_attention.sh`: forward/backward ASM build and installation.
+- `prepare_triton_runtime.sh`: pinned Triton wheel builder.
+- `prepare_primus_turbo_runtime.sh`: pinned Primus-Turbo wheel builder.
+- `prepare_mlperf_runtime.sh`: pinned TransformerEngine/AITER wheel builder.
+- `verify_runtime.py`: image/runtime consistency checks.
+- `configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml`: model config.
+- `config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh`: MI355X environment.
+- `run_and_time.sh`: benchmark entry point.
+
+## Data
 
 Download preprocessed C4 dataset:
 
@@ -52,13 +107,10 @@ After download, you should see files with the following naming conventions:
 
 The data directory is approximately **80 GB** and model directory is approximately **30 GB**.
 
-### How to run
+## Run
 
 ```bash
 export HF_TOKEN=<your_huggingface_token>
 source config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh
 bash run_and_time.sh
 ```
-## Notes
-
-- `log_interval: 99999999` suppresses regular Primus logs

examples/mlperf/gpt_oss_20b/aiter_hd64_asm_override.pth

@@ -0,0 +1 @@
+import aiter_hd64_asm_override