MoE: correct + faster a16w4 (bf16 A x MXFP4 W) SiTUv2 kernel (replaces old a16w4) #14183
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: Atom Test | |
| on: | |
| push: | |
| branches: [main] | |
| pull_request: | |
| types: [opened, synchronize, reopened, ready_for_review, labeled] | |
| branches: [main] | |
| paths-ignore: | |
| - '**/*.md' | |
| - 'docs/**' | |
| - 'LICENSE' | |
| - '.gitignore' | |
| workflow_dispatch: | |
| concurrency: | |
| group: ${{ github.workflow }}-${{ github.ref }} | |
| cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} | |
| env: | |
| ATOM_BRANCH: "main" | |
| ATOM_REPOSITORY_URL: "ROCm/ATOM" | |
| BASE_IMAGE: "rocm/atom-dev:latest" | |
| DOCKER_PULL_MAX_ATTEMPTS: "3" | |
| DOCKER_PULL_RETRY_DELAY_SECONDS: "10" | |
| GITHUB_REPO_URL: ${{ github.event.pull_request.head.repo.clone_url || 'https://github.com/ROCm/Aiter.git' }} | |
| GITHUB_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.event.head_commit.id || github.sha }} | |
| jobs: | |
| check-signal: | |
| if: >- | |
| github.event_name != 'pull_request' || | |
| (github.event.pull_request.draft == false && | |
| (contains(github.event.pull_request.labels.*.name, 'ci:atom') || | |
| contains(github.event.pull_request.labels.*.name, 'ci:atom_full') || | |
| contains(github.event.pull_request.labels.*.name, 'ci:all'))) | |
| runs-on: ubuntu-latest | |
| steps: | |
| - name: Checkout code | |
| uses: actions/checkout@v4 | |
| - name: Wait for Checks workflow | |
| run: ./.github/scripts/check_signal.sh | |
| env: | |
| GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} | |
| GITHUB_SHA: ${{ github.sha }} | |
| load-atom-models: | |
| if: >- | |
| github.event_name != 'pull_request' || | |
| (github.event.pull_request.draft == false && | |
| (contains(github.event.pull_request.labels.*.name, 'ci:atom') || | |
| contains(github.event.pull_request.labels.*.name, 'ci:atom_full') || | |
| contains(github.event.pull_request.labels.*.name, 'ci:all'))) | |
| runs-on: ubuntu-latest | |
| outputs: | |
| models_json: ${{ steps.models.outputs.models_json }} | |
| steps: | |
| - name: Checkout ATOM model configs | |
| uses: actions/checkout@v4 | |
| with: | |
| repository: ${{ env.ATOM_REPOSITORY_URL }} | |
| ref: ${{ env.ATOM_BRANCH }} | |
| - name: Load ATOM accuracy model matrix | |
| id: models | |
| env: | |
| EVENT_NAME: ${{ github.event_name }} | |
| LABELS_JSON: ${{ toJSON(github.event.pull_request.labels.*.name) }} | |
| run: | | |
| python3 <<'PY' | |
| import json | |
| import os | |
| # Always-on PR gate set (every ci:atom / ci:all). Kept small so the | |
| # do-mi350x pool isn't overloaded (overload was flaking the Kimi + | |
| # perf gates). The InferenceX frontier models (DeepSeek-V4-Pro, | |
| # Qwen3.5-397B-A17B-FP8, GLM-5.1-FP8 — all validated PASS on do-mi350x) | |
| # are covered on-demand via the ci:atom_full label, which runs every | |
| # pr/main model in ATOM models_accuracy.json with the runner pins below. | |
| # MiniMax-M2.7 needs HSA_NO_SCRATCH_RECLAIM=1 (scratch reclaim corrupts | |
| # its MoE scratch buffers on gfx950); injected via model_env_overrides | |
| # below. It runs under ci:atom_full, not on every PR. | |
| default_models = [ | |
| {"model_name": "DeepSeek-R1-0528"}, | |
| {"model_name": "gpt-oss-120b"}, | |
| # Kimi-K2.7 ATOM (in-tree) downstream gate. Config (model_path, | |
| # extraArgs -tp4, threshold) is pulled from ATOM models_accuracy.json. | |
| {"model_name": "Kimi-K2.7-Code-MXFP4"}, | |
| ] | |
| labels_raw = os.environ.get("LABELS_JSON") or "[]" | |
| labels = json.loads(labels_raw) | |
| if labels is None: | |
| labels = [] | |
| with open(".github/benchmark/models_accuracy.json", encoding="utf-8") as f: | |
| atom_models = json.load(f) | |
| runner_overrides = { | |
| "atom-mi355-8gpu.predownload": "linux-aiter-do-mi350x-8", | |
| "linux-atom-do-mi350x-8": "linux-aiter-do-mi350x-8", | |
| "linux-atom-mi35x-4": "linux-aiter-do-mi350x-4", | |
| "linux-atom-mi35x-1": "linux-aiter-mi35x-1", | |
| } | |
| # Per-model runner pin (overrides the generic mapping above). Kimi-K2.7 | |
| # runs the downstream gate on the AITER MI350X pool (TP4 -> 4 GPUs). | |
| model_runner_overrides = { | |
| "Kimi-K2.7-Code-MXFP4": "linux-aiter-do-mi350x-4", | |
| "DeepSeek-V4-Pro": "linux-aiter-do-mi350x-8", # -tp 8 | |
| "Qwen3.5-397B-A17B-FP8": "linux-aiter-do-mi350x-4", # -tp 4 | |
| "MiniMax-M2.7-MXFP4": "linux-aiter-do-mi350x-4", # -tp 2 (MXFP4) | |
| "GLM-5.1-FP8": "linux-aiter-do-mi350x-8", # -tp 8 | |
| } | |
| # Per-model extra env injected into the test container (one KEY=VALUE per | |
| # line, merged with any env_vars already declared in the ATOM model | |
| # config). MiniMax MoE needs scratch reclaim disabled or its scratch | |
| # buffers get corrupted on gfx950, failing the accuracy gate. | |
| model_env_overrides = { | |
| "MiniMax-M2.7-MXFP4": "HSA_NO_SCRATCH_RECLAIM=1", | |
| } | |
| # Extra CLI args appended (after the ATOM config's extraArgs) for | |
| # specific models. MiniMax's trust_remote_code modeling breaks ATOM's | |
| # torch.compile/inductor path during warmup (BackendCompilerFailed -> | |
| # FileNotFoundError '<frozen posixpath>'); --enforce-eager only | |
| # disables cudagraph capture in ATOM, so --level 0 (NO_COMPILATION) is | |
| # what bypasses the failing backend. The MXFP4 ATOM entry has no -tp, | |
| # so pin TP2 to match the do-mi350x-2 runner. | |
| model_extra_args_append = { | |
| "MiniMax-M2.7-MXFP4": "-tp 2 --enforce-eager --level 0", | |
| } | |
| def prepare_model(model): | |
| model = dict(model) | |
| model["runner"] = runner_overrides.get(model["runner"], model["runner"]) | |
| model["runner"] = model_runner_overrides.get(model["model_name"], model["runner"]) | |
| extra_env = model_env_overrides.get(model["model_name"]) | |
| if extra_env: | |
| existing_env = str(model.get("env_vars", "")).strip() | |
| model["env_vars"] = ( | |
| f"{existing_env}\n{extra_env}" if existing_env else extra_env | |
| ) | |
| extra_args = model_extra_args_append.get(model["model_name"]) | |
| if extra_args: | |
| base_args = str(model.get("extraArgs", "")).strip() | |
| model["extraArgs"] = f"{base_args} {extra_args}".strip() | |
| model["label"] = model.get("runner", "") | |
| model["display_name"] = model.get("display_name", model["model_name"]) | |
| model["accuracy_test_threshold"] = str(model.get("accuracy_threshold", 0)) | |
| model["run_on_pr"] = True | |
| return model | |
| # Models force-included in ci:atom_full regardless of their ATOM | |
| # test_level. The MiniMax MXFP4 variant is tagged "nightly" in the ATOM | |
| # config, so the pr/main filter below would otherwise skip it. | |
| full_atom_extra_models = {"MiniMax-M2.7-MXFP4"} | |
| full_atom = os.environ["EVENT_NAME"] == "pull_request" and "ci:atom_full" in labels | |
| if not full_atom: | |
| default_model_names = [model["model_name"] for model in default_models] | |
| atom_models_by_name = {model["model_name"]: model for model in atom_models} | |
| missing_models = [ | |
| name for name in default_model_names if name not in atom_models_by_name | |
| ] | |
| if missing_models: | |
| raise SystemExit( | |
| "Missing default model(s) in ATOM accuracy config: " | |
| + ", ".join(missing_models) | |
| ) | |
| selected = [ | |
| prepare_model(atom_models_by_name[name]) for name in default_model_names | |
| ] | |
| else: | |
| selected = [ | |
| prepare_model(model) | |
| for model in atom_models | |
| if str(model.get("test_level", "")).lower() in {"pr", "main"} | |
| or model["model_name"] in full_atom_extra_models | |
| ] | |
| with open(os.environ["GITHUB_OUTPUT"], "a", encoding="utf-8") as out: | |
| out.write(f"models_json={json.dumps(selected)}\n") | |
| mode = "full" if full_atom else "default" | |
| print(f"Selected {len(selected)} ATOM accuracy model(s) for {mode} mode") | |
| PY | |
| atom_benchmark: | |
| if: >- | |
| github.event_name != 'pull_request' || | |
| (github.event.pull_request.draft == false && | |
| (contains(github.event.pull_request.labels.*.name, 'ci:atom') || | |
| contains(github.event.pull_request.labels.*.name, 'ci:atom_full') || | |
| contains(github.event.pull_request.labels.*.name, 'ci:all'))) | |
| needs: [check-signal, load-atom-models] | |
| name: Accuracy (${{ matrix.display_name }}, ${{ matrix.runner }}) | |
| strategy: | |
| fail-fast: false | |
| max-parallel: 2 | |
| matrix: | |
| include: ${{ fromJson(needs.load-atom-models.outputs.models_json) }} | |
| runs-on: ${{ matrix.runner }} | |
| steps: | |
| - name: Checkout code | |
| if: matrix.run_on_pr == true || github.event_name != 'pull_request' | |
| uses: actions/checkout@v4 | |
| with: | |
| repository: ${{ env.ATOM_REPOSITORY_URL }} | |
| ref: ${{ env.ATOM_BRANCH }} | |
| - name: Docker login | |
| run: docker login -u rocmshared -p ${{ secrets.DOCKER_PASSWORD }} || true | |
| - name: Download the ATOM base image | |
| if: matrix.run_on_pr == true || github.event_name != 'pull_request' | |
| run: | | |
| for attempt in $(seq 1 "${DOCKER_PULL_MAX_ATTEMPTS}"); do | |
| if docker pull "${{ env.BASE_IMAGE }}"; then | |
| echo "Docker pull succeeded on attempt ${attempt}" | |
| exit 0 | |
| fi | |
| echo "Docker pull attempt ${attempt} failed" | |
| if [ "${attempt}" -lt "${DOCKER_PULL_MAX_ATTEMPTS}" ]; then | |
| sleep "${DOCKER_PULL_RETRY_DELAY_SECONDS}" | |
| fi | |
| done | |
| echo "Docker pull failed after ${DOCKER_PULL_MAX_ATTEMPTS} attempts" | |
| exit 1 | |
| - name: Generate Dockerfile | |
| if: matrix.run_on_pr == true || github.event_name != 'pull_request' | |
| run: | | |
| cat <<EOF > Dockerfile.mod | |
| FROM ${{ env.BASE_IMAGE }} | |
| RUN echo "=== Aiter version BEFORE uninstall ===" && pip show amd-aiter || true | |
| RUN pip uninstall -y amd-aiter | |
| RUN pip install --upgrade "pybind11>=3.0.1" | |
| RUN pip show pybind11 | |
| RUN rm -rf /app/aiter-test | |
| RUN git clone --no-checkout ${{ env.GITHUB_REPO_URL }} /app/aiter-test && \\ | |
| cd /app/aiter-test && \\ | |
| git checkout --force ${{ env.GITHUB_COMMIT_SHA }} && \\ | |
| git submodule sync && git submodule update --init --recursive && \\ | |
| MAX_JOBS=64 PREBUILD_KERNELS=0 GPU_ARCHS=gfx950 pip install -e . && \\ | |
| ./.github/scripts/install_triton.sh | |
| RUN echo "=== triton version ===" && pip show triton || true | |
| RUN echo "=== Aiter version AFTER installation ===" && pip show amd-aiter || true | |
| EOF | |
| - name: Build the ATOM test image | |
| if: matrix.run_on_pr == true || github.event_name != 'pull_request' | |
| timeout-minutes: 30 | |
| run: | | |
| docker build --network=host \ | |
| --no-cache \ | |
| -t rocm/aiter-ci:atom-test \ | |
| -f Dockerfile.mod . | |
| - name: Start CI container | |
| if: matrix.run_on_pr == true || github.event_name != 'pull_request' | |
| run: | | |
| echo "Clean up containers..." | |
| docker ps -aq -f name=atom_test | xargs -r docker stop | xargs -r docker rm | |
| if [ -f "/etc/podinfo/gha-render-devices" ]; then | |
| DEVICE_FLAG=$(cat /etc/podinfo/gha-render-devices) | |
| else | |
| DEVICE_FLAG="--device /dev/dri" | |
| fi | |
| MODEL_MOUNT="" | |
| if [ -d "/models" ]; then | |
| MODEL_MOUNT="-v /models:/models" | |
| echo "Using host model cache mount: /models" | |
| else | |
| echo "Warning: /models directory not found on runner; model cache will stay container-local." | |
| fi | |
| cat > /tmp/env_file.txt << 'EOF' | |
| ${{ matrix.env_vars }} | |
| EOF | |
| echo "Starting container: rocm/aiter-ci:atom-test" | |
| echo "Model-specific environment variables for ${{ matrix.model_name }}:" | |
| cat /tmp/env_file.txt | |
| docker run -dt --device=/dev/kfd $DEVICE_FLAG \ | |
| $MODEL_MOUNT \ | |
| --ipc=host --group-add video \ | |
| --shm-size=16G \ | |
| --privileged \ | |
| --cap-add=SYS_PTRACE \ | |
| --env-file /tmp/env_file.txt \ | |
| --security-opt seccomp=unconfined \ | |
| --ulimit memlock=-1 \ | |
| --ulimit stack=67108864 \ | |
| -e ATOM_DISABLE_MMAP=true \ | |
| -e HF_TOKEN="${HF_TOKEN:-${{ secrets.HF_TOKEN_TEST }}}" \ | |
| -v "${{ github.workspace }}:/workspace" \ | |
| -w /workspace \ | |
| --name atom_aiter_test \ | |
| rocm/aiter-ci:atom-test | |
| env: | |
| GITHUB_WORKSPACE: ${{ github.workspace }} | |
| - name: Check shm size | |
| if: matrix.run_on_pr == true || github.event_name != 'pull_request' | |
| run: | | |
| df -h # for testing | |
| df -h /dev/shm | |
| docker exec atom_aiter_test df -h /dev/shm | |
| - name: Check version | |
| if: matrix.run_on_pr == true || github.event_name != 'pull_request' | |
| run: | | |
| docker exec atom_aiter_test bash -lc "pip show amd-aiter atom && pip list" | |
| - name: Download Models | |
| if: matrix.run_on_pr == true || github.event_name != 'pull_request' | |
| timeout-minutes: 150 | |
| run: | | |
| set -euo pipefail | |
| if docker exec atom_aiter_test bash -lc "[ -d /models ]"; then | |
| model_cache_dir="/models" | |
| else | |
| model_cache_dir="/run" | |
| fi | |
| target_dir="${model_cache_dir}/${{ matrix.model_path }}" | |
| echo "Downloading model to ${target_dir}" | |
| if ! docker exec \ | |
| -e HF_TOKEN="${HF_TOKEN:-${{ secrets.HF_TOKEN_TEST }}}" \ | |
| -e MODEL_ID="${{ matrix.model_path }}" \ | |
| -e TARGET_DIR="${target_dir}" \ | |
| -e MODEL_DOWNLOAD_TIMEOUT="${MODEL_DOWNLOAD_TIMEOUT}" \ | |
| -e MODEL_LOCK_WAIT_SECONDS="${MODEL_LOCK_WAIT_SECONDS}" \ | |
| -e MODEL_LOCK_POLL_INTERVAL="${MODEL_LOCK_POLL_INTERVAL}" \ | |
| -e MODEL_PROGRESS_INTERVAL="${MODEL_PROGRESS_INTERVAL}" \ | |
| atom_aiter_test bash -lc 'bash /workspace/.github/scripts/download_model_with_lock.sh "$MODEL_ID" "$TARGET_DIR"'; then | |
| echo "Model download failed for '${{ matrix.model_path }}'. Aborting." | |
| exit 1 | |
| fi | |
| env: | |
| MODEL_DOWNLOAD_TIMEOUT: "3h" | |
| MODEL_LOCK_WAIT_SECONDS: "7200" | |
| MODEL_LOCK_POLL_INTERVAL: "30" | |
| MODEL_PROGRESS_INTERVAL: "60" | |
| - name: Run ATOM accuracy test | |
| if: matrix.run_on_pr == true || github.event_name != 'pull_request' | |
| timeout-minutes: 90 | |
| env: | |
| MODEL_NAME: ${{ matrix.model_name }} | |
| DISPLAY_NAME: ${{ matrix.display_name }} | |
| MODEL_PATH: ${{ matrix.model_path }} | |
| MODEL_EXTRA_ARGS: ${{ matrix.extraArgs }} | |
| RUNNER_NAME: ${{ matrix.runner }} | |
| ACCURACY_TEST_THRESHOLD: ${{ matrix.accuracy_test_threshold }} | |
| CLIENT_COMMAND: ${{ matrix.client_command || '' }} | |
| run: | | |
| set -euo pipefail | |
| echo "ATOM accuracy config:" | |
| printf ' display_name: %s\n' "$DISPLAY_NAME" | |
| printf ' model_name: %s\n' "$MODEL_NAME" | |
| printf ' model_path: %s\n' "$MODEL_PATH" | |
| printf ' runner: %s\n' "$RUNNER_NAME" | |
| printf ' accuracy_threshold: %s\n' "$ACCURACY_TEST_THRESHOLD" | |
| echo " extraArgs:" | |
| if [ -n "$MODEL_EXTRA_ARGS" ]; then | |
| printf '%s\n' "$MODEL_EXTRA_ARGS" | |
| else | |
| echo "<empty>" | |
| fi | |
| echo "" | |
| echo "========== Launching ATOM server ==========" | |
| if docker exec atom_aiter_test bash -lc "[ -d /models ]"; then | |
| model_cache_dir="/models" | |
| else | |
| model_cache_dir="/run" | |
| fi | |
| if docker exec atom_aiter_test bash -lc "[ -d '${model_cache_dir}/${{ matrix.model_path }}' ]"; then | |
| model_path="${model_cache_dir}/${{ matrix.model_path }}" | |
| else | |
| model_path="${{ matrix.model_path }}" | |
| fi | |
| printf '.github/scripts/atom_test.sh launch %q %s\n' "$model_path" "$MODEL_EXTRA_ARGS" | \ | |
| docker exec -i atom_aiter_test bash -l | |
| echo "" | |
| echo "========== Running accuracy test ==========" | |
| docker exec \ | |
| -e CLIENT_COMMAND="${CLIENT_COMMAND}" \ | |
| atom_aiter_test bash -lc " | |
| .github/scripts/atom_test.sh accuracy $model_path | |
| " 2>&1 | tee atom_accuracy_output.txt | |
| - name: Check accuracy test results | |
| if: (matrix.run_on_pr == true || github.event_name != 'pull_request') && success() | |
| run: | | |
| result_file=$(ls -1t accuracy_test_results/*.json 2>/dev/null | head -n 1) | |
| if [ -z "$result_file" ] || [ ! -f "$result_file" ]; then | |
| echo "ERROR: No results JSON file found in accuracy_test_results/" | |
| exit 2 | |
| else | |
| echo "RESULT_FILE: $result_file" | |
| fi | |
| flexible_extract_value=$(jq '.results.gsm8k["exact_match,flexible-extract"]' "$result_file") | |
| echo "Flexible extract value: $flexible_extract_value" | |
| echo "Accuracy test threshold: ${{ matrix.accuracy_test_threshold }}" | |
| result=$(awk -v val="$flexible_extract_value" -v threshold="${{ matrix.accuracy_test_threshold }}" 'BEGIN {print (val < threshold) ? 1 : 0}') | |
| if [ "$result" -eq 1 ]; then | |
| echo "Accuracy test failed: Flexible extract value $flexible_extract_value is less than the threshold ${{ matrix.accuracy_test_threshold }}." | |
| exit 1 | |
| else | |
| echo "Accuracy test passed: Flexible extract value $flexible_extract_value is greater than or equal to the threshold ${{ matrix.accuracy_test_threshold }}." | |
| exit 0 | |
| fi | |
| - name: Collect Test Summary | |
| if: (matrix.run_on_pr == true || github.event_name != 'pull_request') && success() | |
| run: | | |
| echo "Accuracy Test Summary for ${{ matrix.model_name }}:" >> $GITHUB_STEP_SUMMARY | |
| awk '/\|Tasks\|Version\|/,/^$/ { if (NF > 0) print }' atom_accuracy_output.txt >> $GITHUB_STEP_SUMMARY | |
| - name: Upload output | |
| if: (matrix.run_on_pr == true || github.event_name != 'pull_request') && always() | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: ${{ matrix.model_name }}_${{ matrix.label }}_${{ matrix.runner }}_atom_accuracy_output.txt | |
| path: atom_accuracy_output.txt | |
| retention-days: 7 | |
| - name: Clean Up | |
| if: (matrix.run_on_pr == true || github.event_name != 'pull_request') && always() | |
| run: | | |
| docker stop atom_aiter_test || true | |
| docker rm atom_aiter_test || true |